[sledge] Refactor QE in Sh.simplify to simplify Context interface

Summary:
The current implementation of quantifier elimination in Sh.simplify is
tightly coupled with the details of what the Context operations
support. In particular, successfully eliminating variables with
Context.elim effectively relies on being given a context that has been
transformed by Context.apply_subst. These operations are sound
independently, but achieving the desired result is delicate.

To simplify this situation, this diff refactors the tightly coupled
usage into a Context.apply_and_elim operation that hides the details
of the interaction inside the Context module. This enables an accurate
specification of apply_and_elim to be given much more simply than can
be done for the separate operations. This also simplifies the
implementation of Sh.simplify.

Reviewed By: jvillard

Differential Revision: D25756577

fbshipit-source-id: b344b3da6

sledge/src/fol/context.ml

@@ -1270,6 +1270,22 @@ let elim xs r =
   in
   (ks, {r with rep})
 
+let apply_and_elim ~wrt xs s r =
+  [%trace]
+    ~call:(fun {pf} -> pf "%a%a@ %a" Var.Set.pp_xs xs Subst.pp s pp_raw r)
+    ~retn:(fun {pf} (zs, r', ks) ->
+      pf "%a@ %a@ %a" Var.Set.pp_xs zs pp_raw r' Var.Set.pp_xs ks ;
+      assert (Var.Set.subset ks ~of_:xs) ;
+      assert (Var.Set.disjoint ks (fv r')) )
+  @@ fun () ->
+  if Subst.is_empty s then (Var.Set.empty, r, Var.Set.empty)
+  else
+    let zs, r = apply_subst wrt s r in
+    if is_unsat r then (Var.Set.empty, unsat, Var.Set.empty)
+    else
+      let ks, r = elim xs r in
+      (zs, r, ks)
+
 (*
  * Replay debugging
  *)
@@ -1286,7 +1302,7 @@ type call =
   | Normalize of t * Term.t
   | Apply_subst of Var.Set.t * Subst.t * t
   | Solve_for_vars of Var.Set.t list * t
-  | Elim of Var.Set.t * t
+  | Apply_and_elim of Var.Set.t * Var.Set.t * Subst.t * t
 [@@deriving sexp]
 
 let replay c =
@@ -1302,7 +1318,7 @@ let replay c =
   | Normalize (r, e) -> normalize r e |> ignore
   | Apply_subst (us, s, r) -> apply_subst us s r |> ignore
   | Solve_for_vars (vss, r) -> solve_for_vars vss r |> ignore
-  | Elim (ks, r) -> elim ks r |> ignore
+  | Apply_and_elim (wrt, xs, s, r) -> apply_and_elim ~wrt xs s r |> ignore
 
 (* Debug wrappers *)
 
@@ -1342,7 +1358,7 @@ let refutes_tmr = Timer.create "refutes" ~at_exit:report
 let normalize_tmr = Timer.create "normalize" ~at_exit:report
 let apply_subst_tmr = Timer.create "apply_subst" ~at_exit:report
 let solve_for_vars_tmr = Timer.create "solve_for_vars" ~at_exit:report
-let elim_tmr = Timer.create "elim" ~at_exit:report
+let apply_and_elim_tmr = Timer.create "apply_and_elim" ~at_exit:report
 
 let add us e r =
   wrap add_tmr (fun () -> add us e r) (fun () -> Add (us, e, r))
@@ -1381,4 +1397,7 @@ let solve_for_vars vss r =
     (fun () -> solve_for_vars vss r)
     (fun () -> Solve_for_vars (vss, r))
 
-let elim ks r = wrap elim_tmr (fun () -> elim ks r) (fun () -> Elim (ks, r))
+let apply_and_elim ~wrt xs s r =
+  wrap apply_and_elim_tmr
+    (fun () -> apply_and_elim ~wrt xs s r)
+    (fun () -> Apply_and_elim (wrt, xs, s, r))

sledge/src/fol/context.mli

@@ -108,12 +108,12 @@ val solve_for_vars : Var.Set.t list -> t -> Subst.t
     terms [e] with free variables contained in as short a prefix of [uss] as
     possible. *)
 
-val elim : Var.Set.t -> t -> Var.Set.t * t
-(** [elim vs x] is [(ks, y)] where [ks] is such that [vs ⊇ ks] and
-    [ks ∩ fv y = ∅] and [y] is such that [∃vs-ks. y] is equivalent to
-    [∃vs. x]. [elim] only removes terms from the existing representation,
-    without performing any additional solving. This means that if a variable
-    in [vs] occurs in an interpreted term in [x], it will not be eliminated. *)
+val apply_and_elim :
+  wrt:Var.Set.t -> Var.Set.t -> Subst.t -> t -> Var.Set.t * t * Var.Set.t
+(** Apply a solution substitution to eliminate the solved variables. That
+    is, [apply_and_elim ~wrt vs s x] is [(zs, x', ks)] where
+    [∃zs. r' ∧ ∃ks. s] is equivalent to [∃xs. r] where [zs] are
+    fresh with respect to [wrt] and [ks ⊆ xs] and is maximal. *)
 
 (**/**)
 

sledge/src/sh.ml

@@ -753,25 +753,19 @@ let remove_absent_xs ks q =
   let ks = Var.Set.inter ks q.xs in
   if Var.Set.is_empty ks then q
   else
-    let ks, ctx = Context.elim ks q.ctx in
-    if Var.Set.is_empty ks then q
-    else
-      let xs = Var.Set.diff q.xs ks in
-      let djns =
-        let rec trim_ks ks djns =
-          List.map djns ~f:(fun djn ->
-              List.map djn ~f:(fun sjn ->
-                  let ks, ctx = Context.elim ks sjn.ctx in
-                  if Var.Set.is_empty ks then sjn
-                  else
-                    { sjn with
-                      us= Var.Set.diff sjn.us ks
-                    ; ctx
-                    ; djns= trim_ks ks sjn.djns } ) )
-        in
-        trim_ks ks q.djns
+    let xs = Var.Set.diff q.xs ks in
+    let djns =
+      let rec trim_ks ks djns =
+        List.map_endo djns ~f:(fun djn ->
+            List.map_endo djn ~f:(fun sjn ->
+                let us = Var.Set.diff sjn.us ks in
+                let djns = trim_ks ks sjn.djns in
+                if us == sjn.us && djns == sjn.djns then sjn
+                else {sjn with us; djns} ) )
       in
-      {q with xs; ctx; djns}
+      trim_ks ks q.djns
+    in
+    {q with xs; djns}
 
 let rec simplify_ us rev_xss q =
   [%Trace.call fun {pf} -> pf "%a@ %a" pp_vss (List.rev rev_xss) pp_raw q]
@@ -788,31 +782,32 @@ let rec simplify_ us rev_xss q =
       )
   in
   (* try to solve equations in ctx for variables in xss *)
-  let subst = Context.solve_for_vars (us :: List.rev rev_xss) q.ctx in
-  let removed, q =
-    if Context.Subst.is_empty subst then (Var.Set.empty, q)
-    else
-      (* normalize wrt solutions *)
-      let q = norm subst q in
-      if is_false q then (Var.Set.empty, false_ q.us)
-      else
-        let removed, ctx =
-          Context.elim (Var.Set.union_list rev_xss) q.ctx
-        in
-        let q = {q with ctx} in
-        let removed =
-          Var.Set.diff removed (fv ~ignore_ctx:() (elim_exists q.xs q))
-        in
-        let keep, removed, _ =
-          Context.Subst.partition_valid removed subst
-        in
-        let q = and_subst keep q in
-        (removed, q)
+  let xss = List.rev rev_xss in
+  let subst = Context.solve_for_vars (us :: xss) q.ctx in
+  let union_xss = Var.Set.union_list rev_xss in
+  let wrt = Var.Set.union us union_xss in
+  let fresh, ctx, removed =
+    Context.apply_and_elim ~wrt union_xss subst q.ctx
   in
-  (* re-quantify existentials *)
-  let q = exists xs q in
-  (* remove the eliminated variables from xs and subformulas' us *)
-  remove_absent_xs removed q
+  ( if Context.is_unsat ctx then false_ q.us
+  else if Context.Subst.is_empty subst then exists xs q
+  else
+    (* normalize wrt solutions *)
+    let q = extend_us fresh q in
+    (* opt: ctx already normalized wrt subst, so just preserve it *)
+    let q = {(norm subst {q with ctx= Context.empty}) with ctx} in
+    if is_false q then false_ q.us
+    else
+      (* opt: removed already disjoint from ctx, so ignore it *)
+      let removed =
+        Var.Set.diff removed (fv ~ignore_ctx:() (elim_exists q.xs q))
+      in
+      let keep, removed, _ = Context.Subst.partition_valid removed subst in
+      let q = and_subst keep q in
+      (* (re)quantify existentials *)
+      let q = exists (Var.Set.union fresh xs) q in
+      (* remove the eliminated variables from xs and subformulas' us *)
+      remove_absent_xs removed q )
   |>
   [%Trace.retn fun {pf} q' ->
     pf "%a@ %a" Context.Subst.pp subst pp_raw q' ;