ext_arith.ml

(*  Copyright 2014 INRIA  *)

open Expr
open Node
open Mlproof
module LL = Llproof
module M = Map.Make(Expr)
module S = Simplex.Make(Expr)

open Arith

exception Incoherent_typing

(* Expression/Types manipulation *)
let equal = Expr.equal

(* Misc *)
let array_for_all2 p a a' =
    if Array.length a <> Array.length a' then
        raise (Invalid_argument "array_for_all2");
    try
        for i = 0 to Array.length a - 1 do
            if not (p a.(i) a'.(i)) then raise Exit
        done;
        true
    with Exit -> false

(* Nodes generated by the extension *)
let gnode g n = match n with
    | Node n -> Node { n with ngoal = g }
    | Stop -> Stop

let requal r r' = match r, r' with
    | Ext("arith", s, l), Ext("arith", s', l') ->
            s = s' && List.for_all2 equal l l'
    | _ -> false

let prio_eq p p' = match p, p' with
    | Prop, Prop | Arity, Arity | Arity_eq, Arity_eq -> true
    | Inst e, Inst e' -> equal e e'
    | _ -> false

let _nequal n n' =
    List.for_all2 equal n.nconc n'.nconc &&
    requal n.nrule n'.nrule &&
    prio_eq n.nprio n'.nprio &&
    n.ngoal = n'.ngoal &&
    array_for_all2 (List.for_all2 equal) n.nbranches n'.nbranches

let nequal n n' = match n, n' with
    | Node n, Node n' -> _nequal n n'
    | Stop, Stop -> true
    | _ -> false

let mk_node_const c e = (* e is a trivially false comparison of constants *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "const", [const (Q.to_string c); e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| |];
    }

let mk_node_eq a b e = (* e : a = b *)
    let a, b =
        try
            let expr, _, c = of_bexpr (lesseq a b) in
            to_nexpr expr, const (Q.to_string c)
        with NotaFormula -> a, b
    in
    Node {
        nconc = [e];
        nrule = Ext ("arith", "eq", [a; b; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [expr_norm (lesseq a b); expr_norm (greatereq a b)] |];
    }

let mk_node_neq a b e = (* e : a != b *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "neq", [a; b; e]);
        nprio = Prop;
        ngoal = 0;
        (* nbranches = [| [expr_norm (less a b)]; [expr_norm (greater a b)] |]; *)
        nbranches = [| [less a b]; [greater a b] |];
    }

let mk_node_tighten s x c e =
    Node {
        nconc = [e];
        nrule = Ext ("arith", "tighten_" ^ s, [x; const c; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [mk_bop s x (const c)] |];
    }

let mk_node_var e1 e2 e = (* e1 : v = expr, e2 : v {comp} const, e : expr {comp} const *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "var", [e1; e2; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [e1; e2] |];
    }

let mk_node_neg s a e = (* e : ~ {s} b *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "neg1_" ^ s, [a; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [mk_ubop (comp_neg s) a] |];
    }

let mk_node_neg2 s a b e = (* e : ~ a {s} b *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "neg2_" ^ s, [a; b; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [mk_bop (comp_neg s) a b ] |];
    }

let mk_node_int_lt a b e = (* e : a < b, => a <= b - 1 *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "int_lt", [a; b; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [expr_norm (lesseq a (minus_one b))] |];
    }

let mk_node_int_gt a b e = (* e : a > b, => a >= b + 1 *)
    Node {
        nconc = [e];
        nrule = Ext ("arith", "int_gt", [a; b; e]);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [expr_norm (greatereq a (plus_one b))] |];
    }

let mk_node_branch v e e' =
    Node {
        nconc = [];
        nrule = Ext ("arith", "simplex_branch", [e; e']);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [e]; [e']; |];
    }

let mk_node_lin e l =
    Node {
        nconc = l;
        nrule = Ext ("arith", "simplex_lin", e :: l);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [e] |];
    }

let mk_node_sim x e b res =
    Node {
        nconc = e :: b;
        nrule = Ext("arith", "simplex_bound", res :: x :: e :: b);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| [res] |];
    }

let mk_node_conflict e e' =
    Node {
        nconc = [e; e'];
        nrule = Ext("arith", "conflict", [e; e']);
        nprio = Prop;
        ngoal = 0;
        nbranches = [| |];
    }

let mk_node_inst e v = match e with
  | Eall (e', p, _) ->
          let t = get_type e' in
          let term = coerce t v in
          let n = Expr.substitute [(e', term)] p in
          Node {
              nconc = [e];
              nrule = Ext("arith", "All", [e; term]);
              nprio = Inst e;
              ngoal = 0;
              nbranches = [| [n] |];
          }
  | Eex (e', p, _) ->
          let t = get_type e' in
          let term = coerce t v in
          let n = Expr.substitute [(e', term)] (enot p) in
          let ne = enot e in
          Node {
              nconc = [ne];
              nrule = Ext("arith", "NotEx", [ne; term]);
              nprio = Inst e;
              ngoal = 0;
              nbranches = [| [n] |];
          }
  | _ -> assert false

(*
let mk_node_switch e a m =
    let new_branch k =
        let a' = eapp (tvar (newname ()) (get_type a), []) in
        let b = to_nexpr [(Q.of_int m), a'; (Q.of_int k), etrue] in
        [ Arith.arith_eq a b;
          Expr.substitute_expr (a, b) e]
    in
    let n = {
        nconc = [e];
        nrule = Ext("arith", "switch", [e; a; const (Q.to_string (Q.of_int m))]);
        nprio = Prop;
        ngoal = 0;
        nbranches = Array.init m new_branch;
    }
    in Node n
    *)

(* Helper around the simplex module *)
type simplex_state = {
    core : S.t;
    ignore : expr list;
    bindings : (expr * expr * expr option * expr option) list;
}

let simplex_empty () = {
    core = S.create ();
    ignore = [];
    bindings = [];
}

let simplex_copy s = {
    core = S.copy s.core;
    ignore = s.ignore;
    bindings = s.bindings;
}

let simplex_print b s =
    let fmt = Format.formatter_of_buffer b in
    let print_var fmt e = Format.fprintf fmt "%s" (Print.sexpr e) in
    Format.fprintf fmt "%a@." (S.print_debug print_var) s.core;
    Log.pp_list ~sep:"\n" (fun b (e, def, inf, upp) ->
        Printf.bprintf b "%a \tin\t %a \t<<\t%a"
        Print.pp_expr def
        (Log.pp_opt Print.pp_expr) inf
        (Log.pp_opt Print.pp_expr) upp
    ) b s.bindings

let bounds_of_comp s c = match s with
    | "$less" -> ((Q.minus_inf,false),(c,true))
    | "$lesseq" -> ((Q.minus_inf,false), (c,false))
    | "$greater" -> ((c,true), (Q.inf,false))
    | "$greatereq" -> ((c,false), (Q.inf,false))
    | _ -> ((Q.minus_inf,false), (Q.inf,false))

let low_binding test c d f low high =
    if test c d then
        f, high
    else
        low, high

let high_binding test c d f low high =
    if test c d then
        low, f
    else
        low, high

let new_bindings low high f = function
    | [c, x], "$less", c' ->
            if Q.sign c <= -1 then begin match low with
                | None -> f, high
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d >= 0 ->
                                low_binding Q.geq Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d <= -1 ->
                                low_binding Q.geq Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end else begin match high with
                | None -> low, f
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d >= 0 ->
                                high_binding Q.leq Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d <= -1 ->
                                high_binding Q.leq Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end
    | [c, x], "$lesseq", c' ->
            if Q.sign c <= -1 then begin match low with
                | None -> f, high
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d >= 0 ->
                                low_binding Q.gt Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d <= -1 ->
                                low_binding Q.gt Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end else begin match high with
                | None -> low, f
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d >= 0 ->
                                high_binding Q.lt Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d <= -1 ->
                                high_binding Q.lt Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end
    | [c, x], "$greater", c' ->
            if Q.sign c >= 0 then begin match low with
                | None -> f, high
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d >= 0 ->
                                low_binding Q.geq Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d <= -1 ->
                                low_binding Q.geq Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end else begin match high with
                | None -> low, f
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d >= 0 ->
                                high_binding Q.leq Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d <= -1 ->
                                high_binding Q.leq Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end
    | [c, x], "$greatereq", c' ->
            if Q.sign c >= 0 then begin match low with
                | None -> f, high
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d >= 0 ->
                                low_binding Q.gt Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d <= -1 ->
                                low_binding Q.gt Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end else begin match high with
                | None -> low, f
                | Some expr -> begin match (of_bexpr expr) with
                        | [d, y], ("$less"|"$lesseq"), d' when Q.sign d >= 0 ->
                                high_binding Q.lt Q.(c' / c) Q.(d' / d) f low high
                        | [d, y], ("$greater"|"$greatereq"), d' when Q.sign d <= -1 ->
                                high_binding Q.lt Q.(c' / c) Q.(d' / d) f low high
                        | _ -> assert false
                        end
            end
    | _ -> low, high

let translate_bound e r = match e with
    | None -> r ()
    | Some e -> begin match (of_bexpr e) with
        | [c, x], ("$less"|"$greater"), c' -> Q.div c' c, true
        | [c, x], ("$lesseq"|"$greatereq"), c' -> Q.div c' c, false
        | _ -> assert false
    end

let pop_option = function | Some a -> a | None -> assert false

let bound_of_expr is_high e bounds =
    if e <> [] && (List.for_all (fun (c,_) -> Q.equal Q.zero c) e) then
        [], Q.zero, false
    else begin
        assert (not (List.exists (fun (c,_) -> Q.equal Q.zero c) e));
        let xor a b = (a && not b) || (not a && b) in
        let rec aux = function
            | [] -> raise Exit
            | [c, x] ->
                    let v, _, einf, eupp = List.find (fun (y, _, _, _) -> equal x y) bounds in
                    if xor is_high (Q.sign c >= 0) then begin
                        let b,t = translate_bound einf (fun () -> raise Exit) in
                        [pop_option einf], Q.mul c b, t
                    end else begin
                        let b, t = translate_bound eupp (fun () -> raise Exit) in
                        [pop_option eupp], Q.mul c b, t
                    end
            | (c, x) :: r ->
                    let l, b, t = aux r in
                    let _, _, einf, eupp = List.find (fun (y, _, _, _) -> equal x y) bounds in
                    if xor is_high (Q.sign c >= 0) then
                        let b', t' = translate_bound einf (fun () -> raise Exit) in
                        (pop_option einf) :: l, Q.add b (Q.mul c b'), t || t'
                    else
                        let b', t' = translate_bound eupp (fun () -> raise Exit) in
                        (pop_option eupp) :: l, Q.add b (Q.mul c b'), t || t'
        in
        try aux e with Exit -> [], (if is_high then Q.inf else Q.minus_inf), false
    end


let bounds_of_clin v expr bounds =
    let _, _, einf, eupp = List.find (fun (y, _, _, _) -> equal y v) bounds in
    let inf, strict_inf = translate_bound einf (fun () -> Q.minus_inf, false) in
    let upp, strict_upp = translate_bound eupp (fun () -> Q.inf, false) in
    let l_bounds, low, strict_low = bound_of_expr false expr bounds in
    let h_bounds, high, strict_high = bound_of_expr true expr bounds in
    Log.debug 15 "arith -- Var bounds : %s %s %a %s %s"
            (Q.to_string inf) (if strict_inf then "<" else "<=") Print.pp_expr v
            (if strict_upp then "<" else "<=") (Q.to_string upp);
    Log.debug 15 "arith -- Deduced bounds : %s %s %a %s %s"
            (Q.to_string low) (if strict_low then "<" else "<=") Print.pp_expr v
            (if strict_high then "<" else "<=") (Q.to_string high);
    begin match strict_inf, strict_upp with
      | false, false      when Q.gt inf upp -> [], pop_option einf, pop_option eupp
      | true, _ | _, true when Q.geq inf upp -> [], pop_option einf, pop_option eupp
      | _ ->
        begin match strict_low, strict_upp with
          | false, false when Q.gt low upp (* low <= v && v <= upp && upp <  low *) -> l_bounds, greatereq v (const (Q.to_string low)), pop_option eupp
          | true, false when Q.geq low upp (* low <  v && v <= upp && upp <= low *) -> l_bounds, greater   v (const (Q.to_string low)), pop_option eupp
          | false, true when Q.geq low upp (* low <= v && v <  upp && upp <= low *) -> l_bounds, greatereq v (const (Q.to_string low)), pop_option eupp
          | true, true when Q.geq low upp  (* low <  v && v <  upp && upp <= low *) -> l_bounds, greater   v (const (Q.to_string low)), pop_option eupp
          | _ ->
            begin match strict_inf, strict_high with
              | false, false when Q.gt inf high -> h_bounds, lesseq v (const (Q.to_string high)), pop_option einf
              | true, false when Q.geq inf high -> h_bounds, lesseq v (const (Q.to_string high)), pop_option einf
              | false, true when Q.geq inf high -> h_bounds, less   v (const (Q.to_string high)), pop_option einf
              | true, true when Q.geq inf high  -> h_bounds, less   v (const (Q.to_string high)), pop_option einf
              | _ -> assert false
            end
        end
    end

let add_binding t x f (e, s, c) =
    let l1, l2 = List.partition (fun (y, _, _, _) -> equal x y) t.bindings in
    let _, def, low, high =
        if List.length l1 = 0 then begin x, x, None, None end
        else if List.length l1 = 1 then begin List.hd l1 end
        else assert false in
    let low, high = new_bindings low high (Some f) (e, s, c) in
    { t with bindings = (x, def, low, high) :: l2 }

let simplex_add t f (e, s, c) =
    match e with
    | []  -> assert false
    | [(c', x)] ->
            let b = Q.div c (Q.abs c') in
            let ((inf,strict_inf), (upp,strict_upp)) = bounds_of_comp s b in
            let ((inf,strict_low), (upp,strict_upp)) =
                if Q.sign c' <= -1 then ((Q.neg upp, false), (Q.neg inf,false)) else ((inf,strict_inf), (upp,strict_upp)) in
            Log.debug 7 "arith -- new bounds : %s %s %a %s %s"
            (Q.to_string inf) (if strict_low then "<" else "<=")
            Print.pp_expr x
            (if strict_upp then "<" else "<=") (Q.to_string upp);
            S.add_bounds t.core ~strict_lower:strict_low ~strict_upper:strict_upp (x, inf, upp);
            (add_binding t x f (e, s, c)), []
    | _ ->
            let expr = to_nexpr e in
            let v = tvar (newname ()) (get_type expr) in
            let e1 = Arith.arith_eq v expr in
            let e2 = mk_bop s v (const (Q.to_string c)) in
            Log.debug 7 "arith -- new variable : %a == %a" Print.pp_expr v Print.pp_expr expr;
            S.add_eq t.core (v, e);
            { core = t.core;
              ignore = e1 :: t.ignore;
              bindings = (v, e1, None, None) :: t.bindings;
            }, [f, mk_node_var e2 e1 f] (* The order (e2 before e1) is actually VERY important here !! *)

exception Internal_error
let nodes_of_tree s f t =
    let rec aux s f t = match !t with
    | None -> raise Internal_error
    | Some S.Branch (v, k, c, c') ->
            Log.debug 7 "arith -- branching: %a <= %s" Print.pp_expr v (Z.to_string k);
            let k = const (Z.to_string k) in
            let under = expr_norm (lesseq v k) in
            let above = expr_norm (greatereq v (plus_one k)) in
            (f, mk_node_branch v under above) :: (
                (aux (add_binding s v under (of_bexpr under)) under c) @
                (aux (add_binding s v above (of_bexpr above)) above c'))
    | Some S.Explanation (v, expr) ->
            let is_zero = expr <> [] && List.for_all (fun (c, _) -> Q.equal Q.zero c) expr in
            let expr = if is_zero then expr else sanitize expr in
            Log.debug 7 "arith -- simplex: %a = %a" Print.pp_expr v Print.pp_expr (to_nexpr expr);
            let l = v :: (List.map snd expr) in
            let relevant = List.map (fun (_, z, _, _) -> z)
                (List.filter (fun (y, y', _, _) -> not (equal y y') && List.exists (fun x -> equal x y) l) s.bindings) in
            let clin = expr_norm (Arith.arith_eq (to_nexpr expr) v) in
            let bounds, nb, conflict = bounds_of_clin v expr s.bindings in
            if bounds = [] && not is_zero then
                [f, mk_node_conflict nb conflict]
            else
                [f, mk_node_lin clin relevant;
                 clin, mk_node_sim v clin bounds nb;
                 nb, mk_node_conflict nb conflict]
    in
    aux s f t

let simplex_solve s e =
    let f = S.nsolve_incr s.core is_int in
    match f () with
    | None -> false, [] (* TODO: rerun f, or try other method ? *)
    | Some S.Solution _ -> false, []
    | Some S.Unsatisfiable cert ->
            Log.debug 5 "arith -- found unsat explanation.";
            Log.debug 10 "arith -- Simplex state :\n%a" simplex_print s;
            true, nodes_of_tree s e cert

(* Instanciation ordering (and substituting) *)

let subst_of_inst n = match n with
    | Node n -> begin match n.nrule with
        | Ext("arith", "All", [Eall(_) as e; v]) -> (e, v)
        | Ext("arith", "NotEx", [Enot(Eex(_) as e, _); v]) -> (e, v)
        | _ -> assert false
        end
    | Stop -> assert false

let subst_inst_rule subst = function
    | Ext("arith", "All", [e; v]) ->
            Ext("arith", "All", [substitute_meta subst e; v])
    | Ext("arith", "NotEx", [e; v]) ->
            Ext("arith", "NotEx", [substitute_meta subst e; v])
    | _ -> assert false

let subst_inst_aux subst n = match n with
    | Node n -> Node { n with
        nconc = List.map (substitute_meta subst) n.nconc;
        nrule = subst_inst_rule subst n.nrule;
        ngoal = 0;
        nbranches = Array.map (List.map (substitute_meta subst)) n.nbranches;
        }
    | Stop -> assert false

let subst_inst (e, n) (e', n') =
    let meta, v = subst_of_inst n in
    let subst = meta, v in
    let e''= substitute_meta subst e' in
    (e'', subst_inst_aux subst n')

let order_inst l =
    let l = List.sort (fun (e, _) (e', _) -> Stdlib.compare (Expr.size e) (Expr.size e')) l in
    let l = List.fold_left (fun acc x -> x :: (List.map (subst_inst x) acc)) [] l in
    l

(* Internal state *)
type state = {
    mutable global : (Node.node_item list) M.t;
    mutable solved : bool;
    stack : (expr * simplex_state * ((expr * Node.node_item) list)) Stack.t;
}

let empty_state =
    let st = Stack.create () in
    Stack.push (etrue, simplex_empty (), []) st;
    {
    global = M.empty;
    solved = false;
    stack = st;
}

let st_reset st =
    if Stack.length st.stack = 0 then
        Stack.push (etrue, simplex_empty (), []) st.stack
    else begin
        Log.debug 1 "%i states left in stack..." (Stack.length st.stack);
        assert false
    end

let st_solved st = st.solved <- true

let st_pop st =
    ignore (Stack.pop st.stack);
    Log.debug 7 "arith -- state stack pop (%i left)" (Stack.length st.stack);
    st.solved <- false

let st_head st =
    try
        let _, r, _ = Stack.top st.stack in
        r
    with Stack.Empty -> assert false

let st_push st (e, t, l) =
    let _, _, l' = Stack.pop st.stack in
    Stack.push (e, t, l @ l') st.stack;
    Log.debug 7 "arith -- state stack push (%i left)" (Stack.length st.stack)

let st_update st e =
    let _, t, l = Stack.pop st.stack in
    Stack.push (e, t, l) st.stack;
    Log.debug 11 "arith -- new stack head : %a" Print.pp_expr e

let st_branch st =
    try
        let e, t, _ = Stack.top st.stack in
        Stack.push (e, simplex_copy t, []) st.stack;
        Log.debug 7 "arith -- copying top stack state (%i left)" (Stack.length st.stack)
    with Stack.Empty ->
        assert false

let st_is_head st e =
    try
        let e', _, _ = Stack.top st.stack in
        equal e e'
    with Stack.Empty -> false

(* Internal state management *)
exception Found of (int -> node_item)

let is_coherent e = function
    | Stop -> true
    | Node n -> List.for_all (fun e' -> equal e e' || Index.member e') n.nconc

let ignore_expr, add_expr, add_branch, remove_expr, todo, set_global, reset, update =
    let st = empty_state in

    let reset () = st_reset st in

    let update e = st_update st e in

    let is_new e =
        try Stack.iter (fun (e', _, l) ->
            if (List.exists (fun (e', _) -> equal e e') l) then raise Exit) st.stack;
        true
        with Exit -> false

    and ignore_expr e =
        try
            Stack.iter (fun (_, t, l) ->
                if List.exists (equal e) t.ignore then raise Exit;
                if List.exists (fun (y, _) -> equal e y) l then raise Exit)
            st.stack;
            false
        with Exit -> true
    in

    let add e = (* try and compute a solution *)
        if is_new e && not st.solved && not (ignore_expr e) then begin
            try
                let (f, s, c) = of_bexpr e in
                (* if not (fis_tau f) then raise NotaFormula; *)
                let t = st_head st in
                if f <> [] then begin
                    let t', res = simplex_add t e (f, s, c) in
                    let b, res' = simplex_solve t' e in
                    let res'' = res @ res' in
                    st_push st (e, t', res'');
                    if b then st_solved st
                end
            with NotaFormula -> ()
        end
    in

    let add_branch () = st_branch st in

    let rec remove e = if st_is_head st e then begin st_pop st; remove e end in

    let todo e =
        let res = ref [] in
        let aux (e', n) = if equal e e' then res := n :: !res in
        M.iter (fun x l -> List.iter (fun n -> aux (x, n)) l) st.global;
        begin match !res with
        | [] -> Stack.iter (fun (_,_,l) -> List.iter aux l) st.stack
        | l -> res := List.rev !res
        end;
        List.filter (is_coherent e) !res
    in

    let set_global l =
        let res = ref false in
        let f (e, l) =
            try
                let l' = M.find e st.global in
                if not (List.for_all2 nequal l l') then begin
                    res := true;
                    Log.debug 9 "arith -- Got %i match for %a (previously : %i)"
                    (List.length l) Print.pp_expr e (List.length l');
                    List.iter (function
                        | Node n -> Log.debug 10 "arith -- match -> %a" Print.pp_mlrule n.nrule;
                        | Stop -> Log.debug 10 "arith -- match -> stop") l;
                    st.global <- M.add e l st.global
                end
            with Not_found ->
                res := true;
                Log.debug 9 "arith -- Got %i match for %a (previously : 0)"
                (List.length l) Print.pp_expr e;
                List.iter (function
                    | Node n -> Log.debug 10 "arith -- match -> %a" Print.pp_mlrule n.nrule
                    | Stop -> Log.debug 10 "arith -- match -> stop") l;
                st.global <- M.add e l st.global
        in
        List.iter f l;
        !res
    in
    ignore_expr, add, add_branch, remove, todo, set_global, reset, update

(* ML -> LL translation *)
let ssub s j = try String.sub s 0 j with Invalid_argument _ -> ""
let esub s j = try String.sub s j (String.length s - j) with Invalid_argument _ -> ""

let tr_rule f = function
    | Ext("arith", "const", [c; e]) ->
            LL.Rextension ("arith", "const", [f c], [f e], [[]])
    | Ext("arith", "eq", [a; b; e]) ->
            LL.Rextension("arith", "eq", [f a; f b], [f e], [[f (expr_norm (lesseq a b)); f (expr_norm (greatereq a b))]])
    | Ext("arith", "neq", [a; b; e]) ->
            LL.Rextension("arith", "neq", [f a; f b], [f e], [[f (less a b)];[f (greater a b)]])
    | Ext("arith", s, [x; c; e]) when ssub s 8 = "tighten_" ->
            LL.Rextension("arith", s, [f x; f c], [f e], [[f (mk_bop (esub s 8) x c)]])
    | Ext("arith", "var", [e1;e2;e]) ->
            LL.Rextension("arith", "var", [f e1;f e2], [f e], [[f e1;f e2]])
    | Ext("arith", s, [a; e]) when ssub s 5 = "neg1_" ->
            LL.Rextension("arith", s, [f a], [f e], [[f (mk_ubop (comp_neg (esub s 5)) a)]])
    | Ext("arith", s, [a; b; e]) when ssub s 5 = "neg2_" ->
            LL.Rextension("arith", s, [f a;f b], [f e], [[f (mk_bop (comp_neg (esub s 5)) a b)]])
    | Ext("arith", "int_lt", [a; b; e]) ->
            LL.Rextension("arith", "int_lt", [f a;f b], [f e], [[f (expr_norm (lesseq a (minus_one b)))]])
    | Ext("arith", "int_gt", [a; b; e]) ->
            LL.Rextension("arith", "int_gt", [f a;f b], [f e], [[f (expr_norm (greatereq a (plus_one b)))]])
    | Ext("arith", "simplex_branch", [e;e']) ->
            LL.Rextension("arith", "simplex_branch", [f e;f e'], [], [[f e];[f e']])
    | Ext("arith", "simplex_lin", e :: l) ->
            LL.Rextension("arith", "simplex_lin", List.map f (e :: l), List.map f l, [[f e]])
    | Ext("arith", "simplex_bound", res :: x :: e :: b) ->
            LL.Rextension("arith", "simplex_bound", List.map f (x :: e :: b), List.map f (e :: b), [[f res]])
    | Ext("arith", "conflict", [e;e']) ->
            LL.Rextension("arith", "conflict", [f e;f e'], [f e;f e'], [[]])
    | Ext("arith", "FM", [x;e;e';e'']) ->
            LL.Rextension("arith", "FM", [f x;f e;f e';f e''], [f e;f e'], [[f e'']])
    | Ext("arith", "All", [p; t]) ->
            LL.Rall (f p, f t)
    | Ext("arith", "NotEx", [Enot(p, _); t]) ->
            LL.Rnotex (f p, f t)
    | Ext("arith", s, _) ->
            Log.debug 0 "arith -- Unknow rule : '%s'" s;
            raise Exit
    | _ -> assert false

let mltoll f p hyps =
    let subproofs, subextras = List.split (Array.to_list hyps) in
    let extras = Expr.diff (List.concat subextras) p.mlconc in
    let nn = {
        LL.conc = List.map f (extras @ p.mlconc);
        LL.rule = tr_rule f p.mlrule;
        LL.hyps = subproofs;
    } in
    (nn, extras)


(* LL -> Coq translation *)
let get_bind = function
    | Eapp(Evar("=", _), [Evar(s, _) as s'; e], _) ->
            let f =
                if is_int s' then "Z.eq_refl"
                else if is_rat s' then "Qeq_refl"
                else "Req_refl"
            in
            s, f, e
    | _ -> assert false

let get_branch = function
    | Eapp(Evar("$lesseq", _), [e; Eapp(Evar(c, _), [], _)], _) -> e, c
    | _ -> assert false

let neg_comp_lemma = function
    | "$less" -> "lt"
    | "$lesseq" -> "leq"
    | "$greater" -> "gt"
    | "$greatereq" -> "geq"
    | _ -> assert false

let ll_p oc r =
    let pr fmt = Printf.fprintf oc fmt in
    match r with
    | LL.Rextension("arith", s, args, l, ll) ->
            pr "(* ARITH -- '%s' : " s;
            List.iter (fun e -> pr "%a : '%s', " Lltocoq.p_expr e (Print.sexpr (get_type e))) args;
            pr "\n * ->";
            List.iter (fun e -> pr "%a, " Lltocoq.p_expr e) l;
            List.iter (fun l ->
                pr "\n * |- ";
                List.iter (fun e -> pr "%a, " Lltocoq.p_expr e) l;
                ) ll;
            pr "*)\n"
    | _ -> assert false

let lltocoq oc r =
    let pr fmt = Printf.fprintf oc fmt in
    if Log.get_debug () >= 0 then ll_p oc r;
    match r with
    (* Constant comparison *)
    | LL.Rextension("arith", "const", _, [e], _) when is_rexpr e ->
            pr "real_const.\n"
    | LL.Rextension("arith", "const", _, [e], _) ->
            pr "arith_norm_in %s; arith_omega %s.\n" (Coqterm.getname e) (Coqterm.getname e)
    (* Equality rule *)
    | LL.Rextension("arith", "eq", [a; b], [e], [[less; greater]]) when is_rexpr a || is_rexpr b ->
            pr "apply (arith_refut _ _ (arith_real_eq %a %a)); [ intros (%s, %s) | real_simpl %a %s ].\n"
            Lltocoq.pp_expr (coqify_to_r a) Lltocoq.pp_expr (coqify_to_r b)
            (Coqterm.getname less) (Coqterm.getname greater)
            Lltocoq.pp_expr (coqify_to_r (norm_coef e)) (Coqterm.getname e)
    | LL.Rextension("arith", "eq", [a; b], [e], [[less; greater]]) ->
            pr "apply (arith_refut _ _ (arith_eq %a %a)); [ intros (%s, %s) | arith_simpl %a %s ].\n"
            Lltocoq.pp_expr (coqify_to_q a) Lltocoq.pp_expr (coqify_to_q b)
            (Coqterm.getname less) (Coqterm.getname greater)
            Lltocoq.pp_expr (coqify_to_q (norm_coef e)) (Coqterm.getname e)
    (* Non-equality rule *)
    | LL.Rextension("arith", "neq", [a; b], [e], [[less]; [greater]]) when is_rexpr a || is_rexpr b ->
            pr "apply (arith_refut _ _ (arith_real_neq %a %a)); [ intros [ %s | %s ] | z_eq_q %s ].\n"
            Lltocoq.pp_expr (coqify_to_r a) Lltocoq.pp_expr (coqify_to_r b)
            (Coqterm.getname less) (Coqterm.getname greater) (Coqterm.getname e)
    | LL.Rextension("arith", "neq", [a; b], [e], [[less]; [greater]]) ->
            pr "apply (arith_refut _ _ (arith_neq %a %a)); [ intros [ %s | %s ] | z_eq_q %s ].\n"
            Lltocoq.pp_expr (coqify_to_q a) Lltocoq.pp_expr (coqify_to_q b)
            (Coqterm.getname less) (Coqterm.getname greater) (Coqterm.getname e)
    (* Comparison tightening (on integers) *)
    | LL.Rextension("arith", "tighten_$lesseq", [x; c], [e], [[e']]) ->
            pr "pose proof (arith_tight_leq _ _ %s) as %s; unfold Qfloor, Zdiv in %s; simpl in %s.\n"
            (Coqterm.getname e) (Coqterm.getname e') (Coqterm.getname e') (Coqterm.getname e')
    | LL.Rextension("arith", "tighten_$greatereq", [x; c], [e], [[e']]) ->
            pr "pose proof (arith_tight_geq _ _ %s) as %s; unfold Qceiling, Zdiv in %s; simpl in %s.\n"
            (Coqterm.getname e) (Coqterm.getname e') (Coqterm.getname e') (Coqterm.getname e')
    (* Comparison Negation *)
    | LL.Rextension("arith", s, [a; b], [e], [[f]]) when ssub s 5 = "neg2_" && (is_rexpr a || is_rexpr b) ->
            pr "apply (arith_refut _ _ (arith_real_neg_%s %a %a)); [zenon_intro %s | exact %s].\n" (neg_comp_lemma (esub s 5))
            Lltocoq.pp_expr (coqify_to_r a) Lltocoq.pp_expr (coqify_to_r b) (Coqterm.getname f) (Coqterm.getname e)
    | LL.Rextension("arith", s, [a; b], [e], [[f]]) when ssub s 5 = "neg2_" ->
            pr "apply (arith_refut _ _ (arith_neg_%s %a %a)); [zenon_intro %s | exact %s].\n" (neg_comp_lemma (esub s 5))
            Lltocoq.pp_expr (coqify_to_q a) Lltocoq.pp_expr (coqify_to_q b) (Coqterm.getname f) (Coqterm.getname e)
    (* Strict integer inequalities *)
    | LL.Rextension("arith", "int_lt", [a; b], [e], [[f]]) ->
            pr "cut %a; [ zenon_intro %s | arith_omega %s ].\n" Lltocoq.p_expr f (Coqterm.getname f) (Coqterm.getname e)
    | LL.Rextension("arith", "int_gt", [a; b], [e], [[f]]) ->
            pr "cut %a; [ zenon_intro %s | arith_omega %s ].\n" Lltocoq.p_expr f (Coqterm.getname f) (Coqterm.getname e)
    (* Variable introduction by the simplex *)
    | LL.Rextension("arith", "var", _, [e], [[e1; e2]]) when is_rexpr e1 ->
            let v, b, expr = get_bind e2 in
            pr "pose (%s := %a).\n" v Lltocoq.pp_expr (coqify_term expr);
            pr "  pose proof (%s %s) as %s; change %s with %a in %s at 2.\n"
                b v (Coqterm.getname e2) v Lltocoq.pp_expr (coqify_term expr) (Coqterm.getname e2);
            pr "  cut %a; [zenon_intro %s | subst %s; real_const ].\n" Lltocoq.p_expr e1 (Coqterm.getname e1) v
    | LL.Rextension("arith", "var", _, [e], [[e1; e2]]) ->
            let v, b, expr = get_bind e2 in
            pr "pose (%s := %a).\n" v Lltocoq.pp_expr (coqify_term expr);
            pr "  pose proof (%s %s) as %s; change %s with %a in %s at 2.\n"
                b v (Coqterm.getname e2) v Lltocoq.pp_expr (coqify_term expr) (Coqterm.getname e2);
            pr "  cut %a; [zenon_intro %s | subst %s; arith_omega %s ].\n" Lltocoq.p_expr e1 (Coqterm.getname e1) v (Coqterm.getname e)
    (* Branching on integer variable value *)
    | LL.Rextension("arith", "simplex_branch", _, _, [[e]; [f]]) ->
            let expr, c = get_branch e in
            pr "destruct (arith_branch %a %s) as [ %s | %s ]; [ | ring_simplify in %s ].\n"
            Lltocoq.pp_expr (coqify_term expr) c (Coqterm.getname e) (Coqterm.getname f) (Coqterm.getname f)
    (* Simplex Linear combination *)
    | LL.Rextension("arith", "simplex_lin", _, l, [[e]]) when List.exists is_rexpr (e :: l) ->
            pr "cut %a; [ zenon_intro %s | %areal_norm; apply eq_refl ].\n" Lltocoq.p_expr e (Coqterm.getname e)
            (fun oc -> List.iter (fun e -> let s, _, _ = get_bind e in Printf.fprintf oc "subst %s; " s)) l
    | LL.Rextension("arith", "simplex_lin", _, l, [[e]]) ->
            pr "cut %a; [ zenon_intro %s | %aarith_norm; apply eq_refl ].\n" Lltocoq.p_expr e (Coqterm.getname e)
            (fun oc -> List.iter (fun e -> let s, _, _ = get_bind e in Printf.fprintf oc "subst %s; " s)) l
    (* New bound rule *)
    | LL.Rextension("arith", "simplex_bound", x :: _, e :: l, [[f]]) when List.exists is_rexpr (e :: l) ->
            let (b, _, _) = of_bexpr e in
            let k, b = fsep b x in
            let ee = Arith.arith_eq x (to_nexpr b) in
            pr "cut (%a); [ zenon_intro %s | real_simpl %a %s ].\n" Lltocoq.p_expr ee (Coqterm.getname ee)
            Lltocoq.pp_expr (coqify_to_r (const (Q.to_string (Q.inv k)))) (Coqterm.getname e);
            let b = List.map (fun (c, y) -> (c, y,
                List.find (fun e -> let b, _, _ = of_bexpr e in match b with
                | [(_, z)] -> equal y z | _ -> false) l)) b
            in
            List.iter (fun (c, y, e') ->
                pr "%s %a %s %s_%s.\n"
                (if Q.sign c >= 0 then "Rle_mult" else "Rle_mult_opp")
                Lltocoq.pp_expr (coqify_to_r @@ const @@ Q.to_string @@ Q.abs c)
                (Coqterm.getname e') (Coqterm.getname e') (Coqterm.getname e)) b;
            (* TODO: extend the function to mix le/lt inequalities *)
            let rec aux oc = function
                | [_, _, e1] -> pr "%s_%s" (Coqterm.getname e1) (Coqterm.getname e)
                | [_, _, e1; _, _, e2] -> Printf.fprintf oc "Rplus_le_compat _ _ _ _ %s_%s %s_%s"
                    (Coqterm.getname e1) (Coqterm.getname e) (Coqterm.getname e2) (Coqterm.getname e)
                | (_, _, e1) :: r -> pr "Rplus_le_compat _ _ _ _ (%a) %s_%s" aux r (Coqterm.getname e1) (Coqterm.getname e)
                | _ -> assert false
            in
            pr "pose proof (%a) as %s_pre.\n" aux b (Coqterm.getname f);
            pr "cut %a; [ zenon_intro %s | rewrite -> %s; real_simpl 1%%R %s_pre ].\n"
            Lltocoq.p_expr f (Coqterm.getname f) (Coqterm.getname ee) (Coqterm.getname f)
    | LL.Rextension("arith", "simplex_bound", x :: _, e :: l, [[f]]) ->
            let (b, _, _) = of_bexpr e in
            let k, b = fsep b x in
            let ee = Arith.arith_eq x (to_nexpr b) in
            pr "cut (%a); [ zenon_intro %s | arith_simpl %a %s ].\n" Lltocoq.p_expr ee (Coqterm.getname ee)
            Lltocoq.pp_expr (coqify_to_q (const (Q.to_string (Q.inv k)))) (Coqterm.getname e);
            let b = List.map (fun (c, y) -> (c, y,
                List.find (fun e -> let b, _, _ = of_bexpr e in match b with
                | [(_, z)] -> equal y z | _ -> false) l)) b
            in
            List.iter (fun (c, y, e') ->
                pr "%s %a %s %s_%s.\n"
                (if Q.sign c >= 0 then "Qle_mult" else "Qle_mult_opp")
                Lltocoq.pp_expr (coqify_to_q @@ const @@ Q.to_string @@ Q.abs c)
                (Coqterm.getname e') (Coqterm.getname e') (Coqterm.getname e)) b;
            (* TODO: extend the function to mix le/lt inequalities *)
            let rec aux oc = function
                | [_, _, e1] -> pr "%s_%s" (Coqterm.getname e1) (Coqterm.getname e)
                | [_, _, e1; _, _, e2] -> Printf.fprintf oc "Qplus_le_compat _ _ _ _ %s_%s %s_%s"
                    (Coqterm.getname e1) (Coqterm.getname e) (Coqterm.getname e2) (Coqterm.getname e)
                | (_, _, e1) :: r -> pr "Qplus_le_compat _ _ _ _ (%a) %s_%s" aux r (Coqterm.getname e1) (Coqterm.getname e)
                | _ -> assert false
            in
            pr "pose proof (%a) as %s_pre.\n" aux b (Coqterm.getname f);
            pr "cut %a; [ zenon_intro %s | rewrite -> %s; arith_simpl 1 %s_pre ].\n"
            Lltocoq.p_expr f (Coqterm.getname f) (Coqterm.getname ee) (Coqterm.getname f)
    (* Conflict rule *)
    | LL.Rextension("arith", "conflict", [e; e'], _, _) when is_rexpr e || is_rexpr e' ->
            pr "real_trans_simpl %s %s.\n" (Coqterm.getname e) (Coqterm.getname e')
    | LL.Rextension("arith", "conflict", [e; e'], _, _) ->
            pr "arith_trans_simpl %s %s.\n" (Coqterm.getname e) (Coqterm.getname e')
    (* Unknown rules *)
    | LL.Rextension("arith", s, _, _, _) ->
            pr "(* TODO unknown rule %s *)\n" s
    | _ -> pr "(* Don't know what to do *)\n"

(* Constants *)
let const_node e = (* comparison of constants *)
    let (f, s, c) = of_bexpr e in
    assert (f = []);
    begin match s with
    | "$less" when Q.geq Q.zero c -> [mk_node_const c e]
    | "$lesseq" when Q.gt Q.zero c -> [mk_node_const c e]
    | "$greater" when Q.leq Q.zero c -> [mk_node_const c e]
    | "$greatereq" when Q.lt Q.zero c -> [mk_node_const c e]
    | "=" when not (Q.equal Q.zero c) -> [mk_node_const c e]
    (*
    | "$is_int" when not (is_z c) -> [mk_node_const c e]
    | "$not_is_int" when is_z c -> [mk_node_const c e]
    | "$not_is_rat" -> [mk_node_const c e]
    *)
    | _ -> []
    end

let is_const e = try let (f, _, _) = of_bexpr e in f = [] with NotaFormula -> false

(* Extension functions *)
let add_formula e = match e with
    | Evar ("#branch", _) -> add_branch ()
    | _ -> ()

let remove_formula e = remove_expr e

let rec iter_open p =
    match ct_from_ml p with
    | None ->
            Log.debug 5 "arith -- empty tree received";
            false
    | Some t ->
        Log.debug 10 "arith -- Trying to solve tree:\n%s" (sctree t);
        let t = collapse t in
        Log.debug 7 "arith -- Simplified tree:\n%s" (sctree t);
        begin match solve_tree t with
        | Unsat ->
                Log.debug 5 "arith -- no solution found in choice tree";
                false
                (*
                begin try
                Log.debug 5 "arith -- switching to next instanciation ";
                iter_open (next_inst p)
                with EndReached -> false end
                *)
        | Abstract s ->
                Log.debug 5 "arith -- found a solution.";
                List.iter (fun (x, v) -> Log.debug 6 "arith -- %a <- %a" Print.pp_expr x Print.pp_expr v) s;
                let global = List.fold_left (fun acc (e, v) -> match e with
                    | Emeta(Eall(_) as e', _) ->
                            (e', mk_node_inst e' v) :: acc
                    | Emeta(Eex(_) as e', _) ->
                            (enot e', mk_node_inst e' v) :: acc
                    | _ -> assert false) [] s in
                let b = set_global (List.map (fun (e, x) -> e, [x]) (order_inst global)) in
                Log.debug 5 "arith -- global todo set (iter : %s)" (if b then "yes" else "no");
                reset ();
                b
        end

let newnodes e g _ =
    let res = match e with
    | _ when ignore_expr e -> []
    | Enot (Eapp (Evar("=",_), [a; b], _), _) when is_num a && is_num b ->
            [mk_node_neq a b e]
    | Enot (Eapp (Evar(("$less"|"$lesseq"|"$greater"|"$greatereq") as s,_), [a; b], _), _) ->
            [mk_node_neg2 s a b e]
    | Enot (Eapp (Evar(("$is_int"|"$is_rat") as s,_), [a], _), _) ->
            [mk_node_neg s a e]
    | _ when is_const e ->
            const_node e
    | Eapp (Evar("=",_), [a; b], _) when is_num a && is_num b ->
            [mk_node_eq a b e]
    | Eapp (Evar("$less",_), [a; b], _) when is_int a && is_int b ->
            [mk_node_int_lt a b e]
    | Eapp (Evar("$greater",_), [a; b], _) when is_int a && is_int b ->
            [mk_node_int_gt a b e]
    | _ -> begin try begin match (of_bexpr e) with
            | [(c', x)], ("$less"|"$lesseq"), c when is_int x && is_q (Q.div c c') ->
                    let c'' = Q.div c c' in
                    if Q.sign c' <= -1 then
                        [mk_node_tighten "$greatereq" x (Q.to_string (ceil c'')) e]
                    else
                        [mk_node_tighten "$lesseq" x (Q.to_string (floor c'')) e]
            | [(c', x)], ("$greater"|"$greatereq"), c when is_int x && is_q (Q.div c c') ->
                    let c'' = Q.div c c' in
                    if Q.sign c' <= -1 then
                        [mk_node_tighten "$lesseq" x (Q.to_string (floor c'')) e]
                    else
                        [mk_node_tighten "$greatereq" x (Q.to_string (ceil c'')) e]
            | _ -> add_expr e; []
            end with
            | NotaFormula -> []
        end
    in
    update e;
    List.map (gnode g) (res @ todo e)
    @ (if is_bexpr e then match e with | Enot(Eapp(Evar("=", _), _, _), _) | Eapp(Evar("=", _), _, _) -> [] | _ -> [Stop] else [])

let make_inst term g = assert false

let to_llproof = mltoll

let declare_context_coq oc =
    let pr fmt = Printf.fprintf oc fmt in
    pr "Require Import zenon_arith.\n";
    pr "Require Import zenon_arith_reals.\n";
    ()

let p_rule_coq = lltocoq

let predef () =
    [ "$less"; "$lesseq"; "$greater"; "$greatereq";
      "$sum"; "$difference"; "$product"; "$uminus";
      "$is_int"; "$not_is_int"; "$is_rat"; "$not_is_rat";
    ]
;;

let predecl () =
  []
;;

Extension.register {
  Extension.name = "arith";
  Extension.newnodes = newnodes;
  Extension.make_inst = make_inst;
  Extension.add_formula = add_formula;
  Extension.remove_formula = remove_formula;
  Extension.iter_open = iter_open;
  Extension.preprocess = (fun x -> x);
  Extension.add_phrase = (fun _ -> ());
  Extension.postprocess = (fun x -> x);
  Extension.to_llproof = to_llproof;
  Extension.declare_context_coq = declare_context_coq;
  Extension.p_rule_coq = p_rule_coq;
  Extension.predef = predef;
  Extension.predecl = predecl;
};;