(* This file implement a specific version of elimination of epsilon,
   for epsilons that are generated by gnatprove for guards of function calls.
   For those, we want to introduce the axioms in a way which is as local as possible
   instead of globally, to protect as much as possible from proving false
   in case of irrealizable function definitions in context.

   What the transformation does in practice is eliminating
   epsilon constructs that
   1) occur in a goal/axiom
   2) have an obvious witness provided by a defining equality
      in their top conjuncts
   by replacing them with their witness, and introducing extra properties
   as local hypothesis/assumptions for goal/axioms,
   below top-level quantifier/premises.
 *)

open Term
open Decl

type witness = {
  witness : term;
  props   : term option;
}

(* Epsilon generated by SPARK are recognized by always having a witness.
   [find_witness x f] attempts to split out a defining equality for x from
   f's conjuncts. *)
let rec find_witness x f =
  match f.t_node with
  | Tapp (ls, [l; r]) when ls_equal ls ps_equ ->
    let test v t = vs_equal v x && not (Mvs.mem x (t_freevars Mvs.empty t)) in
    begin match l.t_node, r.t_node with
    | Tvar y, _ when test y r -> Some { witness = r; props = None }
    | _, Tvar y when test y l -> Some { witness = l; props = None }
    | _ -> None
    end
  | Tbinop (Tand, l, r) ->
    let join w l r =
      Some { w with props = Some (t_attr_copy f (t_and l r)) } in
    begin match find_witness x l with
    | Some ({ props = None; _ } as w) -> Some { w with props = Some r }
    | Some ({ props = Some p; _ } as w) -> join w p r
    | None -> match find_witness x r with
      | Some ({ props = None; _} as w) -> Some { w with props = Some l }
      | Some ({ props = Some p; _ } as w) -> join w l p
      | None -> None
    end
  | _ -> None

(* Eliminate epsilon that have a witness, and accumulate the properties
   stated over those witnesses in axml.
   Parameter topv lists variables that are to be treated as
   top-level constants, and as such should not re-quantified in
   generated axioms. *)
let rec lift_f_acc topv axml t0 =
  match t0.t_node with
  | Teps fb ->
    let vs, f = t_open_bound fb in
    let axml, t = match find_witness vs f with
      | Some { witness = w; props = None } -> lift_f_acc topv axml w
      | Some { witness = w; props = Some p } ->
        let axml, p = lift_f_acc topv axml p in
        let vl = Mvs.keys (Mvs.set_diff (t_vars t0) topv) in
        let p = t_forall_close_merge vl (t_subst_single vs w p) in
        p :: axml, w
      | None ->
        let axml, f = lift_f_acc topv axml f in
        axml, t_eps_close vs f
    in
    axml, t_attr_copy t0 t
  | _ -> let axml, t = t_map_fold (lift_f_acc topv) axml t0 in
    axml, t_attr_copy t0 t

let lift_f topv = lift_f_acc topv []

(* With pol giving polarity of formula, eliminate
   axioms and adds axioms as hypothesis/assumptions
   as nested in resulting formula as possible.
   Top-level bound variables are aggregated in topv to avoid
   re-quantifying over them in axiom formulas. *)
let rec lift_q pol topv t0 =
  let binop = if pol then Tand else Timplies in
  let wrap = List.fold_left (fun t ax -> t_binary binop ax t) in
  let t = match t0.t_node with
  | Tquant (Tforall, tq) ->
    let vsl, tr, t, cb = t_open_quant_cb tq in
    let axml_tr, tr =
      let ap = Lists.map_fold_left in ap (ap (lift_f_acc topv)) [] tr in
    let topv = List.fold_left Svs.add_left topv vsl in
    wrap (t_forall (cb vsl tr (lift_q pol topv t))) axml_tr
  | Tbinop (Tand, _, _)
  | Tbinop (Tor, _, _) -> t_map (lift_q pol topv) t0
  | Tbinop (Timplies, t1, t2) ->
    let axml, t1 = lift_f topv t1 in
    wrap (t_binary Timplies t1 (lift_q pol topv t2)) axml
  | Tlet (t1, bt2) ->
    let (x, t2) = t_open_bound bt2 in
    let axml, t1 = lift_f topv t1 in
    wrap (t_let t1 (t_close_bound x (lift_q pol topv t2))) axml
  | _ -> let axml, t = lift_f topv t0 in wrap t axml
  in
  t_attr_copy t0 t

let lift_q pol = lift_q pol Svs.empty

let lift_d d =
  [match d.d_node with
  (* For goals and axioms, introduce epsilon axioms
     after top-level quantifier and premises *)
  | Dprop (Pgoal, _, _) -> decl_map (lift_q false) d
  | Dprop (Paxiom, _, _) -> decl_map (lift_q true) d
  | _ -> d]

let eliminate_gnatprove_guard_epsilon = Trans.decl lift_d None

let () = Trans.register_transform "eliminate_gnatprove_guard_epsilon" eliminate_gnatprove_guard_epsilon
  ~desc:"Eliminate@ comprehension@ forms@ generated@ from@ gnatprove@ function@ call@ guards@ from@ axioms@ and@ goals"

