path: root/abstract/abs_interp_edd.ml

open Ast
open Ast_util
open Formula
open Typing
open Cmdline

open Util
open Num_domain



exception Top

exception Found_int of int

module I (ND : NUMERICAL_ENVIRONMENT_DOMAIN) : sig

    val do_prog : cmdline_opt -> rooted_prog -> unit

    val test : unit -> unit

end = struct


    (*  **********************
              EDD Domain  
        **********************  *)

    type item = string

    type evar = id * item list
    type nvar = id * bool

    type varenv = {
        evars         : evar list;
        nvars         : nvar list;
        ev_order      : (id, int) Hashtbl.t;
    }

    type edd =
        | DBot
        | DTop
        | DVal of int * (bool * int)   (* bool*int : new case ? iterations before widen ? *)
        | DChoice of int * id * (item * edd) list

    type edd_v = {
        ve            : varenv;
        root          : edd;
        leaves        : (int, ND.t) Hashtbl.t;
        (* add here eventual annotations *)
    }

    (*
      Utility functions for memoization

      memo : (('a -> 'b) -> 'a -> 'b) -> 'a -> 'b
                -> (int * 'a) -> (int * 'b)
      memo2 : (('a -> 'b -> 'c) -> 'a -> 'b -> 'c)
                -> 'a -> 'b -> 'c

      Where 'a = 'b = 'c = edd, but it can be adapted.
    *)
    let key = function
        | DBot -> 0
        | DTop -> 1
        | DVal (i, _) -> 2 * i + 2
        | DChoice(i, _, _) -> 2 * i + 3

    let memo f =
        let memo = Hashtbl.create 12 in
        let rec ff v =
          try Hashtbl.find memo (key v)
          with Not_found ->
            let r = f ff v in
            Hashtbl.add memo (key v) r; r
        in ff
          
    let memo2 f =
        let memo = Hashtbl.create 12 in
        let rec ff v1 v2 =
          try Hashtbl.find memo (key v1, key v2)
          with Not_found ->
            let r = f ff v1 v2 in
            Hashtbl.add memo (key v1, key v2) r; r
        in ff

    let edd_node_eq = function
        | DBot, DBot -> true
        | DTop, DTop -> true
        | DVal (i, _), DVal (j, _) when i = j -> true
        | DChoice(i, _, _), DChoice(j, _, _) when i = j -> true
        | _ -> false


    let new_node_fun () =
        let nc = ref 0 in
        let node_memo = Hashtbl.create 12 in
        fun v l ->
          let _, x0 = List.hd l in
          if List.exists (fun (_, x) -> not (edd_node_eq (x, x0))) l
            then begin
              let k = (v, List.map (fun (a, b) -> a, key b) l) in
              let n =
                try Hashtbl.find node_memo k
                with _ -> (incr nc; Hashtbl.add node_memo k !nc; !nc)
              in
              DChoice(n, v, l)
            end else x0

    let get_leaf_fun_st () =
        let leaves = Hashtbl.create 12 in
        let lc = ref 0 in
        let get_leaf st x = 
          if ND.is_top x then DTop else
          if ND.is_bot x then DBot else
            try
              Hashtbl.iter (fun i v -> if ND.eq v x then raise (Found_int i)) leaves;
              incr lc;
              Hashtbl.add leaves !lc x;
              DVal (!lc, st)
            with Found_int i -> DVal (i, st)
        in leaves, get_leaf

    let get_leaf_fun () =
        let leaves, get_leaf = get_leaf_fun_st () in
        leaves, get_leaf (false, 0)

    let rank ve = function
        | DChoice(_, x, _) -> Hashtbl.find ve.ev_order x
        | _ -> 10000000 (* HYPOTHESIS : program will never have more than
                              that number of variables *)

    (*
      edd_print : Format.formatter -> edd_v -> unit
    *)
    let edd_print fmt v =
        let max_v = ref 0 in

        let print_nodes = Queue.create () in
        let a = Hashtbl.create 12 in

        let node_pc = Hashtbl.create 12 in
        let f f_rec = function
          | DChoice(_, _, l) ->
            List.iter
              (fun (_, c) -> match c with
                  | DChoice(n, _, _) ->
                      begin try Hashtbl.add node_pc n (Hashtbl.find node_pc n + 1)
                      with Not_found -> Hashtbl.add node_pc n 1 end
                  | _ -> ())
            l;
            List.iter (fun (_, c) -> f_rec c) l
          | _ -> ()
        in memo f v.root;

        let rec print_n fmt = function
          | DBot -> Format.fprintf fmt "⊥";
          | DTop -> Format.fprintf fmt "⊤";
          | DVal (i, (s, _)) -> if i > !max_v then max_v := i;
                    Format.fprintf fmt "v%d%s" i (if s then "*" else "");
          | DChoice(_, v, l) ->
            match List.filter (fun (_, x) -> x <> DBot) l with
            | [(c, nn)] ->
              let aux fmt = function
                | DChoice(nn, _, _) as i when Hashtbl.find node_pc nn >= 2 ->
                  if Hashtbl.mem a nn then () else begin
                    Queue.push i print_nodes;
                    Hashtbl.add a nn ()
                  end;
                  Format.fprintf fmt "n%d" nn
                | x -> print_n fmt x
              in
              Format.fprintf fmt "%a = %s,@ %a" Formula_printer.print_id v c aux nn
            | _ ->
              Format.fprintf fmt "%a ? " Formula_printer.print_id v;
              let print_u fmt (c, i) =
                Format.fprintf fmt "%s → " c;
                match i with
                | DChoice(nn, v, l) ->
                  if Hashtbl.mem a nn then () else begin
                    Queue.push i print_nodes;
                    Hashtbl.add a nn ()
                  end;
                  Format.fprintf fmt "n%d" nn
                | _ -> Format.fprintf fmt "%a" print_n i
              in
              Format.fprintf fmt "@[<h>%a@]" (print_list print_u ", ") l;
        in
        Format.fprintf fmt "@[<hov>%a@]@." print_n v.root;
        while not (Queue.is_empty print_nodes) do
          match Queue.pop print_nodes with
          | DChoice(n, v, l) as x ->
            Format.fprintf fmt "n%d: @[<hov>%a@]@." n print_n x
          | _ -> assert false
        done;

        for id = 0 to !max_v do
          try let v = Hashtbl.find v.leaves id in
            Format.fprintf fmt "v%d: %a@." id ND.print v
          with Not_found -> ()
        done

    let edd_dump_graphviz v file =
        let o = open_out file in
        let fmt = Format.formatter_of_out_channel o in
        Format.fprintf fmt "digraph G {@[<v 4>@,";

        let nov = Hashtbl.create 12 in

        let f f_rec = function
          | DChoice(n, v, x) ->
            let aux fmt = function
            | DBot -> Format.fprintf fmt "bot"
            | DTop -> Format.fprintf fmt "top"
            | DVal(i, _) -> Format.fprintf fmt "v%d" i
            | DChoice(n, _, _) -> Format.fprintf fmt "n%d" n
            in
            let p = try Hashtbl.find nov v with _ -> [] in
            Hashtbl.replace nov v (n::p);
            Format.fprintf fmt "n%d [label=\"%s\"];@ " n v;
            List.iter (fun (i, c) ->
              if c <> DBot then Format.fprintf fmt "n%d -> %a [label=\"%s\"];@ " n aux c i;
              f_rec c) x
          | _ -> ()
        in memo f v.root;
        
        Hashtbl.iter (fun var nodes ->
            Format.fprintf fmt "{ rank = same; ";
            List.iter (Format.fprintf fmt "n%d; ") nodes;
            Format.fprintf fmt "}@ ")
          nov;

        Format.fprintf fmt "@]}@.";

        close_out o


    (*
      edd_bot : varenv -> edd_v
    *)
    let edd_bot ve = { ve; root = DBot; leaves = Hashtbl.create 1 }

    (*
      edd_top : evar list -> nvar list -> edd_v
    *)
    let edd_top ve = { ve; root = DTop; leaves = Hashtbl.create 1 }

    (*
      edd_of_cons : varenv -> enum_cons -> edd_v
    *)
    let edd_of_cons ve (op, vid, r) =
        let op = match op with | E_EQ -> (=) | E_NE -> (<>) in

        let leaves = Hashtbl.create 1 in

        let root = match r with
        | EItem x ->
          DChoice(0, vid,
              List.map (fun v -> if op v x then v, DTop else v, DBot)
                (List.assoc vid ve.evars))      
        | EIdent vid2 ->
            let a, b =
                if Hashtbl.find ve.ev_order vid < Hashtbl.find ve.ev_order vid2
                  then vid, vid2
                  else vid2, vid
            in
            let nc = ref 0 in
            let nb x =
              incr nc;
              DChoice(!nc, b,
                    List.map (fun v -> if op v x then v, DTop else v, DBot)
                      (List.assoc b ve.evars))
            in
            DChoice(0, a, List.map (fun x -> x, nb x) (List.assoc a ve.evars))
        in
        { ve; root; leaves }

    (*
      edd_join : edd_v -> edd_v -> edd_v
      edd_meet : edd_v -> edd_v -> edd_v
    *)
    let edd_join a b =
        if a.root = DBot then b else
        if b.root = DBot then a else
        if a.root = DTop || b.root = DTop then edd_top a.ve else begin
          let ve = a.ve in
          let leaves, get_leaf = get_leaf_fun () in
          let dq = new_node_fun () in

          let f f_rec na nb =
            match na, nb with
            | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
              let kl = List.map2
                  (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                    ta, f_rec ba bb)
                  la lb
              in
              dq va kl

            | DTop, _ | _, DTop -> DTop
            | DBot, DBot -> DBot

            | DChoice(_,va, la), _ when rank ve na < rank ve nb ->
              let kl = List.map (fun (k, ca) -> k, f_rec ca nb) la in
              dq va kl
            | _, DChoice(_, vb, lb) when rank ve nb < rank ve na ->
              let kl = List.map (fun (k, cb) -> k, f_rec na cb) lb in
              dq vb kl

            | DVal (u, _), DVal (v, _) ->
              let x = ND.join (Hashtbl.find a.leaves u) (Hashtbl.find b.leaves v) in
              get_leaf x
            | DVal(u, _), DBot ->
              get_leaf (Hashtbl.find a.leaves u)
            | DBot, DVal(v, _) ->
              get_leaf (Hashtbl.find b.leaves v)

            | _ -> assert false (* so that OCaml won't complain ; all cases ARE handled *)
          in
            { leaves; ve; root = time "join" (fun () -> memo2 f a.root b.root) }
        end

    let edd_meet a b =
        if a.root = DTop then b else
        if b.root = DTop then a else
        if a.root = DBot || b.root = DBot then edd_bot a.ve else begin
          let ve = a.ve in
          let leaves, get_leaf = get_leaf_fun () in
          let dq = new_node_fun () in

          let f f_rec na nb =
            match na, nb with
            | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
              let kl = List.map2
                  (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                    ta, f_rec ba bb)
                  la lb
              in
              dq va kl

            | DBot, _ | _, DBot -> DBot
            | DTop, DTop -> DTop

            | DChoice(_, va, la), _ when rank ve na < rank ve nb ->
              let kl = List.map (fun (k, ca) -> k, f_rec ca nb) la in
              dq va kl
            | _, DChoice(_, vb, lb) when rank ve nb < rank ve na ->
              let kl = List.map (fun (k, cb) -> k, f_rec na cb) lb in
              dq vb kl

            | DVal (u, _) , DVal (v, _) ->
              let x = ND.meet (Hashtbl.find a.leaves u) (Hashtbl.find b.leaves v) in
              get_leaf x
            | DVal(u, _), DTop ->
              get_leaf (Hashtbl.find a.leaves u)
            | DTop, DVal(v, _) ->
              get_leaf (Hashtbl.find b.leaves v)

            | _ -> assert false (* see above *)
          in
            { leaves; ve; root = time "meet" (fun () -> memo2 f a.root b.root) }
        end



    (*
      edd_num_apply : edd_v -> (ND.t -> ND.t) -> edd_v
      edd_apply_ncl : edd_v -> num_cons list -> edd_v
    *)
    let edd_num_apply v nfun =
        let ve = v.ve in

        let leaves, get_leaf = get_leaf_fun () in

        let dq = new_node_fun () in

        let f f_rec n =
          match n with
          | DBot -> DBot
          | DTop -> get_leaf (nfun (ND.top ve.nvars))
          | DVal (i, _) -> get_leaf (nfun (Hashtbl.find v.leaves i))
          | DChoice(n, var, l) ->
            let l = List.map (fun (k, v) -> k, f_rec v) l in
            dq var l
        in
        { leaves; ve; root = memo f v.root }

    let edd_apply_ncl v ncl =
        edd_num_apply v (fun n -> ND.apply_cl n ncl)

    (*
      edd_apply_ecl : edd_v -> enum_cons list -> edd_v
    *)
    let edd_apply_ecl v ec =
        let rec cl_k = function
          | [] -> edd_top v.ve
          | [a] -> edd_of_cons v.ve a
          | l ->
            let n = ref 0 in
            let la, lb = List.partition (fun _ -> incr n; !n mod 2 = 0) l in
            edd_meet (cl_k la) (cl_k lb)
        in
        let cons_edd = cl_k ec in
        edd_meet v cons_edd
        (*List.fold_left (fun v c -> edd_meet v (edd_of_cons v.ve c)) v ec*)

    (*
      edd_apply_cl : edd_v -> conslist -> edd_v
    *)
    let rec edd_apply_cl v (ec, nc, r) =
        let v = edd_apply_ecl v ec in
        match r with
        | CLTrue ->
          edd_apply_ncl v nc
        | CLFalse -> edd_bot v.ve
        | CLAnd (a, b) ->
          let v = edd_apply_cl v ([], nc, a) in
          if v.root = DBot then v else
            edd_apply_cl v ([], nc, b)
        | CLOr((eca, nca, ra), (ecb, ncb, rb)) ->
          edd_join (edd_apply_cl v (eca, nc@nca, ra))
                   (edd_apply_cl v (ecb, nc@ncb, rb))


    (*
      edd_extract_path : edd_v -> id -> edd_v
    *)
    let edd_extract_path v leaf_id =
        let ve = v.ve in

        let dq = new_node_fun () in

        let f f_rec n =
          match n with
          | DVal (i, _) when i = leaf_id -> DTop
          | DChoice(n, var, l) ->
            let l = List.map (fun (k, v) -> k, f_rec v) l in
            dq var l
          | _ -> DBot
        in
        { leaves = Hashtbl.create 1; ve; root = memo f v.root }

    (*
      edd_eq : edd_v -> edd_v -> bool
    *)
    let edd_eq a b =
        let f f_rec na nb =
          match na, nb with
          | DBot, DBot -> true
          | DTop, DTop -> true
          | DVal (i, _), DVal (j, _) ->
            ND.eq (Hashtbl.find a.leaves i) (Hashtbl.find b.leaves j)
          | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
            List.for_all2 (fun (ca, na) (cb, nb) -> assert (ca = cb); f_rec na nb)
                la lb
          | _ -> false
        in memo2 f a.root b.root


    (*
      edd_subset : edd_v -> edd_v -> bool
    *)
    let edd_subset a b =
        let rank = rank a.ve in
        let f f_rec na nb =
          match na, nb with
          | DBot, _ -> true
          | _, DTop -> true
          | DTop, DBot -> false

          | DVal(i, _), DBot -> ND.is_bot (Hashtbl.find a.leaves i)
          | DTop, DVal(i, _) -> ND.is_top (Hashtbl.find b.leaves i)

          | DVal(i, _), DVal(j, _) ->
            ND.subset (Hashtbl.find a.leaves i) (Hashtbl.find b.leaves j)

          | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
            List.for_all2 (fun (ca, na) (cb, nb) -> assert (ca = cb); f_rec na nb)
              la lb
          | DChoice(_, va, la), _ when rank na < rank nb ->
            List.for_all (fun (c, n) -> f_rec n nb) la
          | _, DChoice(_, vb, lb) when rank na > rank nb ->
            List.for_all (fun (c, n) -> f_rec na n) lb
          | _ -> assert false
        in memo2 f a.root b.root

    (*
      edd_forget_vars : edd_v -> id list -> edd_v
    *)
    let edd_forget_vars v vars =
        let ve = v.ve in

        let leaves, get_leaf = get_leaf_fun () in

        let nc = ref 0 in
        let memo = Hashtbl.create 12 in
        let node_memo = Hashtbl.create 12 in
        let rec f l =
            let kl = List.sort Pervasives.compare (List.map key l) in
            try Hashtbl.find memo kl
            with Not_found -> let r =
              try
                let cn, fn = List.fold_left
                  (fun (cn, fn) node -> match node with
                    | DBot -> cn, fn
                    | DTop -> raise Top
                    | DVal (i, _) -> cn, i::fn
                    | DChoice (n, v, l) -> (n, v, l)::cn, fn)
                  ([], []) l in
                let cn = List.sort
                  (fun (n, v1, _) (n, v2, _) -> Pervasives.compare
                      (Hashtbl.find ve.ev_order v1) (Hashtbl.find ve.ev_order v2))
                  cn in
                if cn = [] then
                  if fn = [] then DBot
                  else
                    let x = list_fold_op ND.join
                        (List.map (Hashtbl.find v.leaves) fn)
                    in get_leaf x
                else
                  let _, dv, cl = List.hd cn in
                  let d, nd = List.partition (fun (_, v, _) -> v = dv) cn in
                  let ch1 = List.map (fun (a, b, c) -> DChoice(a, b, c)) nd in
                  let ch2 = List.map (fun i -> DVal (i, (false, 0))) fn in
                  if List.mem dv vars then
                    (* Do union of all branches branching from nodes on variable dv *)
                    let ch3 = List.flatten
                      (List.map (fun (_, _, c) -> List.map snd c) d) in
                    f (ch1@ch2@ch3)
                  else
                    (* Keep disjunction on variable dv *)
                    let d, nd = List.partition (fun (_, v, _) -> v = dv) cn in
                    let cc = List.map
                      (fun (c, _) ->
                        let ch3 = List.map (fun (_, _, cl) -> List.assoc c cl) d in
                        c, f (ch1@ch2@ch3))
                      cl in
                    let _, x0 = List.hd cc in
                    if List.exists (fun (_, x) -> not (edd_node_eq (x, x0))) cc
                    then begin
                      let k = (dv, List.map (fun (a, b) -> a, key b) cc) in
                      let n =
                        try Hashtbl.find node_memo k
                        with _ -> (incr nc; Hashtbl.add node_memo k !nc; !nc)
                      in
                      DChoice(n, dv, cc)
                    end else x0
              with | Top -> DTop
            in Hashtbl.add memo kl r; r
        in
        { leaves; ve; root = f [v.root] }

    (*
      edd_eassign : edd_v -> (id * id) list -> edd_v
    *)
    let edd_eassign v ids =
      let v = edd_forget_vars v (List.map fst ids) in
      edd_apply_ecl v
        (List.map (fun (x, y) -> (E_EQ, x, EIdent y)) ids)


    (*
        Just a function to test EDDs
    *)
    let test () =
        let ve = {
            evars = ["x", ["tt"; "ff"]; "y", ["tt"; "ff"]; "z", ["tt"; "ff"]];
            nvars = [];
            ev_order = Hashtbl.create 2 } in
        Hashtbl.add ve.ev_order "x" 0;
        Hashtbl.add ve.ev_order "y" 1;
        Hashtbl.add ve.ev_order "z" 2;
        let u = edd_of_cons ve (E_EQ, "x", EIdent "y") in
        Format.printf "x = y : @[%a@]@." edd_print u;
        let v = edd_of_cons ve (E_NE, "y", EIdent "z") in
        Format.printf "y != z : @[%a@]@." edd_print v;
        let w = edd_meet u v in
        Format.printf "x = y && y != z : @[%a@]@." edd_print w;
        let t = edd_join u v in
        Format.printf "x = y || y != z : @[%a@]@." edd_print t;
        let e = edd_forget_vars w ["y"] in
        Format.printf "x = y && y != z ; forget y : @[%a@]@." edd_print e;
        let f = edd_forget_vars t ["y"] in
        Format.printf "x = y || y != z ; forget y : @[%a@]@." edd_print f


    (* ******************************
            Abstract interpret
      ******************************* *)
    
    type env = {
        rp            : rooted_prog;
        opt           : cmdline_opt;

        last_vars     : (bool * id * typ) list;
        all_vars      : (bool * id * typ) list;
        ve            : varenv;

        (* program expressions *)
        init_cl       : conslist;
        cl            : conslist;
        cl_g          : conslist;
        guarantees    : (id * bool_expr) list;

        (* abstract interpretation *)
        cycle         : (id * id * typ) list;     (* s'(x) = s(y) *)
        forget        : (id * typ) list;          (* s'(x) not specified *)
    }


    (*
      edd_find_starred : edd_v -> int option
      edd_unstar : edd_v -> int -> edd_v
    *)
    let edd_find_starred v =
        let f f_rec = function
          | DVal (i, (true, _)) -> raise (Found_int i)
          | DChoice(_, _, l) -> List.iter (fun (_, x) -> f_rec x) l
          | _ -> ()
        in
        try memo f v.root; None
        with Found_int i -> Some i

    let edd_unstar v i =
        let f f_rec = function
          | DChoice(a, b, l) -> DChoice(a, b, List.map (fun (c, x) -> c, f_rec x) l)
          | DVal(j, (s, n)) when i = j -> DVal(i, (false, n))
          | x -> x
        in
        { v with root = memo f v.root }


    (*
      edd_join_widen : edd_v -> edd_v -> edd_v
    *)
    let edd_widen (a:edd_v) (b:edd_v) =
        let ve = a.ve in
        let leaves, get_leaf = get_leaf_fun () in
        let dq = new_node_fun () in

        let f f_rec na nb =
          match na, nb with
          | DTop, _ | _, DTop -> DTop
          | DBot, DBot -> DBot

          | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
            let kl = List.map2
                (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                  ta, f_rec ba bb)
                la lb
            in
            dq va kl

          | DBot, DVal (i, _) ->
              get_leaf (Hashtbl.find b.leaves i)
          | DVal (i, _), DBot ->
              get_leaf (Hashtbl.find a.leaves i)
          | DVal (u, _), DVal (v, _) ->
            let p1, p2 = edd_extract_path a u, edd_extract_path b v in
            let widen = 
              if edd_eq p1 p2 then true else false
            in
            let x = (if widen then ND.widen else ND.join)
                 (Hashtbl.find a.leaves u) (Hashtbl.find b.leaves v) in
            get_leaf x

          | DChoice(_,va, la), _ when rank ve na < rank ve nb ->
            let kl = List.map (fun (k, ca) -> k, f_rec ca nb) la in
            dq va kl
          | _, DChoice(_, vb, lb) when rank ve nb < rank ve na ->
            let kl = List.map (fun (k, cb) -> k, f_rec na cb) lb in
            dq vb kl

          | _ -> assert false
        in
        { leaves; ve; root = time "join/W" (fun () -> memo2 f a.root b.root) }

    (*
      edd_accumulate : edd_v -> edd_v -> edd_v

      Sometimes do global widening.
    *)
    let edd_accumulate env (a:edd_v) (b:edd_v) =
        let ve = a.ve in
        let leaves, get_leaf = get_leaf_fun_st () in
        let dq = new_node_fun () in

        let f f_rec na nb =
          match na, nb with
          | DTop, _ | _, DTop -> DTop
          | DBot, DBot -> DBot

          | DChoice(_, va, la), DChoice(_, vb, lb) when va = vb ->
            let kl = List.map2
                (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                  ta, f_rec ba bb)
                la lb
            in
            dq va kl

          | DBot, DVal (i, _) ->
              get_leaf (true, 0) (Hashtbl.find b.leaves i)
          | DVal (i, s), DBot ->
              get_leaf s (Hashtbl.find a.leaves i)
          | DVal (u, (s1, i1)), DVal (v, _) ->
            let p1, p2 = edd_extract_path a u, edd_extract_path b v in
            let d1, d2 = Hashtbl.find a.leaves u, Hashtbl.find b.leaves v in
            let widen = edd_eq p1 p2 && i1 >= env.opt.widen_delay in
            let x = (if widen then ND.widen else ND.join) d1 d2 in
            get_leaf (s1, i1 + 1) x

          | DChoice(_,va, la), _ when rank ve na < rank ve nb ->
            let kl = List.map (fun (k, ca) -> k, f_rec ca nb) la in
            dq va kl
          | _, DChoice(_, vb, lb) when rank ve nb < rank ve na ->
            let kl = List.map (fun (k, cb) -> k, f_rec na cb) lb in
            dq vb kl

          | _ -> assert false
        in
        { leaves; ve; root = time "join/*" (fun () -> memo2 f a.root b.root) }

    (*
      edd_star_new : edd_v -> edd_v -> edd_v

      Star in s leaves that were not present in s0
    *)
    let edd_star_new s0 s =
        let f f_rec = function
          | DChoice(n, x, l) ->
            DChoice(n, x, List.map (fun (c, x) -> c, f_rec x) l)
          | DVal(i, (false, n)) when
              not (edd_subset (edd_meet (edd_extract_path s i) s) s0)
            ->
              DVal(i, (true, n))
          | x -> x
        in
        { s with root = memo f s.root }

    (*
      pass_cycle : env -> edd_v -> edd_v
      unpass_cycle : env -> edd_v -> edd_v
    *)
    let pass_cycle env v =
        let assign_e, assign_n = List.fold_left
          (fun (ae, an) (a, b, t) -> match t with
            | TEnum _ -> (a, b)::ae, an
            | TInt | TReal -> ae, (a, NIdent b)::an)
          ([], []) env.cycle in

        let v = edd_eassign v assign_e in
        let v = edd_num_apply v (fun nv -> ND.assign nv assign_n) in

        let ef, nf = List.fold_left
            (fun (ef, nf) (var, t) -> match t with
              | TEnum _ -> var::ef, nf
              | TReal | TInt -> ef, var::nf)
            ([], []) env.forget in

        let v = edd_forget_vars v ef in
        edd_num_apply v (fun nv -> List.fold_left ND.forgetvar nv nf)

    let unpass_cycle env v =
        let assign_e, assign_n = List.fold_left
          (fun (ae, an) (a, b, t) -> match t with
            | TEnum _ -> (b, a)::ae, an
            | TInt | TReal -> ae, (b, NIdent a)::an)
          ([], []) env.cycle in

        let v = edd_eassign v assign_e in
        let v = edd_num_apply v (fun nv -> ND.assign nv assign_n) in

        let ef, nf = List.fold_left
            (fun (ef, nf) (_, var, t) ->
              if var.[0] <> 'N' then
                match t with
                | TEnum _ -> var::ef, nf
                | TReal | TInt -> ef, var::nf
              else ef, nf)
            ([], []) env.rp.all_vars in

        let v = edd_forget_vars v ef in
        edd_num_apply v (fun nv -> List.fold_left ND.forgetvar nv nf)

    (*
      extract_linked_evars : conslist -> (id * id) list

      Extract all pairs of enum-type variable (x, y) appearing in an
      equation like x = y or x != y

      A couple may appear several times in the result.
    *)
    let rec extract_linked_evars_root (ecl, _, r) =
        let v_ecl = List.fold_left
            (fun c (_, x, v) -> match v with
                | EIdent y -> (x, y)::c
                | _ -> c)
            [] ecl
        in
        v_ecl

    let rec extract_const_vars_root (ecl, _, _) =
        List.fold_left
          (fun l (_, x, v) -> match v with 
                | EItem _ -> x::l
                | _ -> l)
          [] ecl

    (*
      scope_constrict : id list -> (id * id) list -> id list

      Orders the variable in the first argument such as to minimize the
      sum of the distance between the position of two variables appearing in
      a couple of the second list. (minimisation is approximate, this is
      an heuristic so that the EDD will not explode in size when expressing
      equations such as x = y && u = v && a != b)
    *)
    let scope_constrict vars cp_id =
        let var_i = Array.of_list vars in
        let n = Array.length var_i in

        let i_var = Hashtbl.create n in
        Array.iteri (fun i v -> Hashtbl.add i_var v i) var_i;

        let cp_i = List.map
          (fun (x, y) -> Hashtbl.find i_var x, Hashtbl.find i_var y)
          cp_id in

        let eval i =
          let r = Array.make n (-1) in
          Array.iteri (fun pos var -> r.(var) <- pos) i;
          Array.iteri (fun _ x -> assert (x <> (-1))) r;
          List.fold_left
            (fun s (x, y) -> s + abs (r.(x) - r.(y)))
            0 cp_i
        in

        let best = Array.init n (fun i -> i) in

        let usefull = ref true in
        Format.printf "SCA";
        while !usefull do
          Format.printf ".@?";

          usefull := false;
          let try_s x =
            if eval x < eval best then begin
              Array.blit x 0 best 0 n;
              usefull := true
            end
          in

          for i = 0 to n-1 do
            let tt = Array.copy best in
            (* move item i at beginning *)
            let temp = tt.(i) in
            for j = i downto 1 do tt.(j) <- tt.(j-1) done;
            tt.(0) <- temp;
            (* try all positions *)
            try_s tt;
            for j = 1 to n-1 do
              let temp = tt.(j-1) in
              tt.(j-1) <- tt.(j);
              tt.(j) <- temp;
              try_s tt
            done
          done
        done;
        Format.printf "@.";

        Array.to_list (Array.map (Array.get var_i) best)

    (*
      force_ordering : id list -> (float * id list) list -> id list

      Determine a good ordering for enumerate variables based on the FORCE algorithm
    *)
    let force_ordering vars groups =
        let var_i = Array.of_list vars in
        let n = Array.length var_i in

        let i_var = Hashtbl.create n in
        Array.iteri (fun i v -> Hashtbl.add i_var v i) var_i;
        Hashtbl.add i_var "#BEGIN" (-1);

        let ngroups = List.map
          (fun (w, l) -> w, List.map (Hashtbl.find i_var) l)
          groups in

        let ord = Array.init n (fun i -> i) in

        for iter = 0 to 500 do
            let rev = Array.make n (-1) in
            for i = 0 to n-1 do rev.(ord.(i)) <- i done;

            let bw = Array.make n 0. in
            let w = Array.make n 0. in

            let gfun (gw, l) =
              let sp = List.fold_left (+.) 0.
                (List.map
                  (fun i -> if i = -1 then -.gw else float_of_int (rev.(i))) l)
              in
              let b = sp /. float_of_int (List.length l) in
              List.iter (fun i -> if i >= 0 then begin
                          bw.(i) <- bw.(i) +. (gw *. b);
                          w.(i) <- w.(i) +. gw end)
                  l 
            in
            List.iter gfun ngroups;

            let b = Array.init n
              (fun i ->
                if w.(i) = 0. then
                    float_of_int i
                else bw.(i) /. w.(i)) in
                    
            let ol = List.sort
              (fun i j -> Pervasives.compare b.(i) b.(j))
              (Array.to_list ord) in
            Array.blit (Array.of_list ol) 0 ord 0 n
        done;
        List.map (Array.get var_i) (Array.to_list ord)

    (*
      init_env : cmdline_opt -> rooted_prog -> env
    *)
    let init_env opt rp =
        let init_f, f = Transform.f_of_prog rp false in
        let _, f_g = Transform.f_of_prog rp true in
        Format.printf "Init formula:@.%a@.@." Formula_printer.print_expr init_f;
        Format.printf "Cycle formula:@.%a@.@." Formula_printer.print_expr f;

        let cl = Formula.conslist_of_f f in
        let cl_g = Formula.conslist_of_f f_g in
        let init_cl = Formula.conslist_of_f init_f in
        Format.printf "Cycle conslist:@.%a@.@." Formula_printer.print_conslist cl;

        let guarantees = Transform.guarantees_of_prog rp in
        Format.printf "Guarantees:@.";
        List.iter (fun (id, f) ->
            Format.printf "  %s: %a@." id Formula_printer.print_expr f)
          guarantees;
        Format.printf "@.";
        
        (* add variables from LASTs *)
        let last_vars = uniq_sorted (List.sort compare (Transform.extract_last_vars f_g)) in
        let last_vars = List.map
          (fun id ->
            let (_, _, ty) = List.find (fun (_, u, _) -> id = "L"^u) rp.all_vars
              in false, id, ty)
          last_vars in
        let all_vars = last_vars @ rp.all_vars in

        Format.printf "Vars: @[<hov>%a@]@.@."
            (print_list Ast_printer.print_typed_var ", ")
            all_vars;

        let num_vars, enum_vars = List.fold_left
            (fun (nv, ev) (_, id, t) -> match t with
                | TEnum ch -> nv, (id, ch)::ev
                | TInt -> (id, false)::nv, ev
                | TReal -> (id, true)::nv, ev)
            ([], []) all_vars in


        (* calculate order for enumerated variables *)
        let evars = List.map fst enum_vars in

        let lv = extract_linked_evars_root cl_g in
        let lv = uniq_sorted
             (List.sort Pervasives.compare (List.map ord_couple lv)) in

        let lv_f = List.map (fun (a, b) -> (1.0, [a; b])) lv in
        let lv_f = lv_f @ (List.map (fun v -> (10.0, ["#BEGIN"; v]))
          (extract_const_vars_root cl)) in
        let lv_f = lv_f @ (List.map (fun v -> (5.0, ["#BEGIN"; v]))
          (List.filter (fun n -> is_suffix n "init") evars)) in
        let lv_f = lv_f @ (List.map (fun v -> (3.0, ["#BEGIN"; v]))
          (List.filter (fun n -> is_suffix n "state") evars)) in
        let lv_f = lv_f @
          (List.map (fun v -> (0.7, [v; "L"^v]))
            (List.filter (fun n -> List.mem ("L"^n) evars) evars)) in
        let evars_ord =
          if true then
            time "FORCE" (fun () -> force_ordering evars lv_f)
          else
            time "SCA" (fun () -> scope_constrict evars lv)
        in

        let evars_ord =
          if false then
            let va, vb = List.partition (fun n -> is_suffix n "init") evars_ord in
            let vb, vc = List.partition (fun n -> is_suffix n "state") vb in
            (List.rev va) @ vb @ vc
          else
            evars_ord
        in

        let ev_order = Hashtbl.create (List.length evars) in
        List.iteri (fun i x -> Hashtbl.add ev_order x i) evars_ord;
        let ve = { evars = enum_vars; nvars = num_vars; ev_order } in

        Format.printf "Order for variables: @[<hov>[%a]@]@."
          (print_list Formula_printer.print_id ", ") evars_ord;

        (* calculate cycle variables and forget variables *)
        let cycle = List.fold_left
          (fun q (_, id, ty) ->
              if id.[0] = 'L' then
                (id, String.sub id 1 (String.length id - 1), ty)::q
              else q)
          [] last_vars
        in
        let forget = List.map (fun (_, id, ty) -> (id, ty)) rp.all_vars in


        { rp; opt; ve; last_vars; all_vars;
          init_cl; cl; cl_g; guarantees;
          cycle; forget }

          


    let do_prog opt rp =
        let e = init_env opt rp in

        (* Do iterations until fixpoint is reached *)
        let rec ch_it n x =
          edd_dump_graphviz x (Format.sprintf "/tmp/graph-it%d.dot" n);
          match edd_find_starred x with
          | None ->
            Format.printf "It. %d : full iteration.@." n;

            let d2 = edd_apply_cl x e.cl in
            let dc = pass_cycle e d2 in
            if dc.root = DBot then begin
              Format.printf "@.WARNING: contradictory hypotheses!@.@.";
              x
            end else begin
              let y = edd_star_new x (edd_accumulate e x dc) in

              if e.opt.vverbose_ci then
                Format.printf "d2 %a@. dc %a@. y %a@."
                  edd_print d2 edd_print dc edd_print y;

              if e.opt.verbose_ci then
                Format.printf " -> %a@." edd_print y;

              if not (edd_eq x y) then ch_it (n+1) y else y
            end
          | Some i ->
            let path = edd_extract_path x i in
            let x = edd_unstar x i in
            Format.printf "It. %d: @[<hov>%a@]@." n edd_print path;

            let path_target = unpass_cycle e path in
            let start = edd_meet path x in

            let f i =
              let i = edd_meet path i in
              let i' = edd_meet i path_target in
              let j = edd_apply_cl i' e.cl in

              if e.opt.vverbose_ci then
                Format.printf "i %a@.i' %a@.j %a@."
                  edd_print i edd_print i' edd_print j;

              let q = edd_join start (pass_cycle e j) in
              edd_meet path q
            in

            let rec iter n i =
              let fi = f i in
              let j =
                if n < e.opt.widen_delay then
                  edd_join i fi
                else
                  edd_widen i fi
              in
              if edd_eq i j then j else iter (n+1) j
            in
            let y = iter 0 start in
            let z = fix edd_eq f y in


            let fj = pass_cycle e (edd_apply_cl z e.cl) in
            if fj.root = DBot then begin
              Format.printf "@.WARNING: contradictory hypotheses!@.@.";
              x
            end else begin
              let r = edd_star_new x (edd_accumulate e x fj) in

              if e.opt.verbose_ci then
                Format.printf " -> %a@." edd_print r;

              ch_it (n+1) r
            end
        in

        Format.printf "Calculating initial state...@.";
        let init_acc = edd_star_new (edd_bot e.ve)
            (pass_cycle e (edd_apply_cl (edd_top e.ve) e.init_cl)) in
        
        (* Dump final state *)
        let acc = ch_it 0 init_acc in
        edd_dump_graphviz acc "/tmp/graph-final0.dot";

        Format.printf "Finishing up...@.";
        let final = edd_apply_cl acc e.cl in
        edd_dump_graphviz final "/tmp/graph-final.dot";
        if e.opt.verbose_ci then
          Format.printf "@.Final:@.@[<hov>%a@]@." edd_print final;

        (* Check guarantees *)
        let check_guarantee (id, f) =
          let cl = Formula.conslist_of_f f in
          Format.printf "@[<hv 4>%s:@ %a ⇒ ⊥  @ "
            id Formula_printer.print_conslist cl;
          let z = edd_apply_cl final cl in
          if z.root = DBot then
            Format.printf "OK@]@ "
          else
            Format.printf "FAIL@]@ "
        in
        if e.guarantees <> [] then begin
          Format.printf "Guarantee @[<v 0>";
          List.iter check_guarantee e.guarantees;
          Format.printf "@]@."
        end;

        (* Examine probes *)
        if List.exists (fun (p, _, _) -> p) e.all_vars then begin
          let final_flat = edd_forget_vars final
            (List.fold_left
              (fun l (_, id, ty) -> match ty with
                | TInt | TReal -> l
                | TEnum _ -> id::l)
              [] e.all_vars) in
          let final_flat = match final_flat.root with
            | DTop -> ND.top e.ve.nvars
            | DBot -> ND.bottom e.ve.nvars
            | DVal(i, _) -> Hashtbl.find final_flat.leaves i
            | DChoice _ -> assert false
          in

          Format.printf "Probes: @[<v 0>";
          List.iter (fun (p, id, ty) ->
            if p then match ty with
            | TInt | TReal ->
              Format.printf "%a ∊ %a@." Formula_printer.print_id id
                ND.print_itv (ND.project final_flat id)
            | TEnum _ -> Format.printf "%a : enum variable@."
                Formula_printer.print_id id)
            e.all_vars;
          Format.printf "@]@."
        end

end