[sledge] Generalize Term.solve_zero_eq to specify solved-for subterm

Summary:
If `Term.solve_zero_eq` is passed `for_`, then that subterm is solved
for.

Reviewed By: ngorogiannis

Differential Revision: D19282647

fbshipit-source-id: 5d5b76af5

sledge/src/llair/term.ml

@@ -1019,11 +1019,22 @@ let is_false = function Integer {data} -> Z.is_false data | _ -> false
 
 (** Solve *)
 
-let solve_zero_eq = function
+let solve_zero_eq ?for_ e =
+  ( match e with
   | Add args ->
-      let c, q = Qset.min_elt_exn args in
+      let+ c, q =
+        match for_ with
+        | Some f ->
+            let q = Qset.count args f in
+            if Q.equal Q.zero q then None else Some (f, q)
+        | None -> Some (Qset.min_elt_exn args)
+      in
       let n = Sum.to_term (Qset.remove args c) in
       let d = rational (Q.neg q) in
       let r = div n d in
-      Some (c, r)
-  | _ -> None
+      (c, r)
+  | _ -> None )
+  |> check (fun soln ->
+         match (for_, soln) with
+         | Some f, Some (c, _) -> assert (equal f c)
+         | _ -> () )

sledge/src/llair/term.mli

@@ -248,6 +248,7 @@ val is_false : t -> bool
 
 (** Solve *)
 
-val solve_zero_eq : t -> (t * t) option
+val solve_zero_eq : ?for_:t -> t -> (t * t) option
 (** [solve_zero_eq d] is [Some (e, f)] if [d = 0] can be equivalently
-    expressed as [e = f] for some monomial subterm [e] of [d]. *)
+    expressed as [e = f] for some monomial subterm [e] of [d]. If [for_] is
+    passed, then the subterm [e] must be [for_]. *)