open Ast
open Ast_util
open Formula
open Typing
open Cmdline

open Util
open Num_domain

open Varenv


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  
        **********************  *)

    (*
      This abstract domain is capable of representing values of a program using enumerated
      variables and numerical variables. The representation is a decision graph on
      enumerated variables, whose leaves are values for the numerical variables in a given
      numerical domain (relationnal or non-relationnal). This domain necessarily does the
      disjunction between all the cases that appear for the enumerated variables, and is
      not able to do a disjunction with respect to a condition on numerical variables if
      that condition is not bound to a boolean variable appearing in the program.

      Possible extensions :
      - use groups of variables, so that an EDD does not have to consider all the
        numerical and enumerated variables at once
      - add the possibility for numeric condition decision nodes (something like ghost
        boolean variables that would be bound to a numeric condition)
      - ...

      This domain is currently the most promising lead in our research on abstract
      interpretation of Scade programs.
    *)

    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;
    }

    (*
      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 dq = new_node_fun () in

        let 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 cc = List.map
                      (fun (c, _) ->
                        let ch3 = List.map (fun (_, _, cl) -> List.assoc c cl) d in
                        c, f (ch1@ch2@ch3))
                      cl in
                    dq dv cc
              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;
            last_vars = [];
            all_vars = [];
            cycle = [];
            forget = [];
            forget_inv = []; } 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;

        ve            : varenv;

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


    (*
      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_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 "widen" (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 "accumulate" (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 : varenv -> 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.ve.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.ve.forget_inv in

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



    (*
      init_env : cmdline_opt -> rooted_prog -> env
    *)
    let init_env opt rp =
        let f = Transform.f_of_prog_incl_init rp false in

        let f = simplify_k (get_root_true f) f in
        Format.printf "Complete formula:@.%a@.@." Formula_printer.print_expr f;

        (*
          Here we simplify the program formula so that uselessly redundant variables don't
          appear anymore. If an enumerated equation x = y appears at the root of the
          program, then we chose to remove either x or y.
        *)
        let facts = get_root_true f in
        let f, rp, repls = List.fold_left
          (fun (f, (rp : rooted_prog), repls) eq ->
            match eq with
            | BEnumCons(E_EQ, a, EIdent b)
                when a.[0] <> 'L' && b.[0] <> 'L' ->

              let a = try List.assoc a repls with Not_found -> a in
              let b = try List.assoc b repls with Not_found -> b in

              if a = b then
                f, rp, repls
              else begin
                let keep, repl =
                  if String.length a <= String.length b
                    then a, b
                    else b, a
                in
                Format.printf "Replacing %s with %s@." repl keep;
                let f = formula_replace_evars [repl, keep; "L"^repl, "L"^keep] f in
                let rp =
                  { rp with all_vars =
                      List.filter (fun (_, id, _) -> id <> repl) rp.all_vars } in
                let repls = [repl, keep; "L"^repl, "L"^keep]@
                  (List.map
                    (fun (r, k) -> r,
                      if k = repl then keep else
                      if k = "L"^repl then "L"^keep else k)
                    repls) in
                f, rp, repls
              end
            | _ -> f, rp, repls)
          (f, rp, []) facts in
        Format.printf "Complete formula after simpl:@.%a@.@."
            Formula_printer.print_expr f;

        (*
          Here we specialize the program formula for the two following cases :
          - L/must_reset = tt, this is the first instant of the program (global reset)
          - L/must_reset = ff, this is for all the rest of the time
        *)
        let init_f = simplify_k [BEnumCons(E_EQ, "L/must_reset", EItem bool_true)] f in
        let f = simplify_k [BEnumCons(E_NE, "L/must_reset", EItem bool_true)] f in

        let init_f = simplify_k (get_root_true init_f) init_f in
        let f = simplify_k (get_root_true f) f 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 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
        let guarantees = List.map
          (fun (id, f, v) -> id, formula_replace_evars repls f, v)
          guarantees in
        Format.printf "Guarantees:@.";
        List.iter (fun (id, f, _) ->
            Format.printf "  %s: %a@." id Formula_printer.print_expr f)
          guarantees;
        Format.printf "@.";
        
        let ve = mk_varenv rp f cl in

        { rp; opt; ve; init_cl; cl; guarantees; }



    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.ve 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.ve 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.ve (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.ve (edd_apply_cl (edd_top e.ve) e.init_cl)) in
        
        (* Iterate *)
        let acc = ch_it 0 init_acc in

        (* Dump final state *)
        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.ve.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.ve.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.ve.all_vars;
          Format.printf "@]@."
        end

end