From eb750ba6f97435354cbdfb96166144584261cc73 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Thu, 16 Apr 2020 03:38:55 -0700
Subject: [PATCH] [sledge] Avoid cyclic solutions to polynomial equations
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:
For non-linear polynomial equations, the solver can currently choose
to solve for a variable that occurs in a non-linear subterm, leading
to a non-idepotent solution substitution. For example, `x - y + (x ×
z) = 0` could be solved for `x`, leading to `x ↦ y - (x × z)` where
`x` is solved in terms of itself. This diff refines Term.solve_zero_eq
to avoid such cyclic solutions.

Reviewed By: jvillard

Differential Revision: D20637482

fbshipit-source-id: 6d7df85c3
---
 sledge/lib/equality_test.ml | 22 ++++++----------
 sledge/lib/sh_test.ml       | 15 +++--------
 sledge/lib/term.ml          | 51 ++++++++++++++++++++++++++++++++-----
 3 files changed, 55 insertions(+), 33 deletions(-)

diff --git a/sledge/lib/equality_test.ml b/sledge/lib/equality_test.ml
index 28960ad14..9a03bfa06 100644
--- a/sledge/lib/equality_test.ml
+++ b/sledge/lib/equality_test.ml
@@ -368,10 +368,9 @@ let%test_module _ =
       [%expect
         {|
       {sat= false;
-       rep= [[%x_5 ↦ (((u8) (%x_5 + -1)) + -2)];
-             [%y_6 ↦ (((u8) (%x_5 + -1)) + -3)];
-             [((u8) %y_6) ↦ (((u8) (%x_5 + -1)) + -5)];
-             [((u8) (%x_5 + -1)) ↦ ]]} |}]
+       rep= [[%x_5 ↦ (%y_6 + 1)];
+             [((u8) %y_6) ↦ (%y_6 + -2)];
+             [((u8) (%x_5 + -1)) ↦ (%y_6 + 3)]]} |}]
 
     let%test _ = is_false r16
 
@@ -383,10 +382,9 @@ let%test_module _ =
       [%expect
         {|
       {sat= false;
-       rep= [[%x_5 ↦ (((u8) %y_6) + 1)];
-             [%y_6 ↦ (((u8) %y_6) + 1)];
-             [((u8) %x_5) ↦ (((u8) %y_6) + 1)];
-             [((u8) %y_6) ↦ ]]} |}]
+       rep= [[%y_6 ↦ %x_5];
+             [((u8) %x_5) ↦ %x_5];
+             [((u8) %y_6) ↦ (%x_5 + -1)]]} |}]
 
     let%test _ = is_false r17
 
@@ -397,13 +395,9 @@ let%test_module _ =
       [%expect
         {|
         {sat= true;
-         rep= [[%y_6 ↦ (((u8) %y_6) + 1)];
-               [((u8) %x_5) ↦ %x_5];
-               [((u8) %y_6) ↦ ]]}
+         rep= [[((u8) %x_5) ↦ %x_5]; [((u8) %y_6) ↦ (%y_6 + -1)]]}
 
-          %x_5 = ((u8) %x_5)
-        ∧ ((u8) %y_6) = ((u8) (((u8) %y_6) + 1))
-        ∧ (((u8) %y_6) + 1) = %y_6 |}]
+          %x_5 = ((u8) %x_5) ∧ (%y_6 + -1) = ((u8) %y_6) |}]
 
     let r19 = of_eqs [(x, y + z); (x, !0); (y, !0)]
 
diff --git a/sledge/lib/sh_test.ml b/sledge/lib/sh_test.ml
index 303cb3eac..3ec54ce27 100644
--- a/sledge/lib/sh_test.ml
+++ b/sledge/lib/sh_test.ml
@@ -136,20 +136,11 @@ let%test_module _ =
       pp q' ;
       [%expect
         {|
-        ∃ %x_6 .
-          %x_6 = ((u8) %x_6)
-        ∧ ((u8) %y_7) = ((u8) (((u8) %y_7) + 1))
-        ∧ (((u8) %y_7) + 1) = %y_7
-        ∧ emp
+        ∃ %x_6 .   %x_6 = ((u8) %x_6) ∧ (%y_7 + -1) = ((u8) %y_7) ∧ emp
 
-          ((u8) %y_7) = ((u8) (((u8) %y_7) + 1))
-        ∧ (((u8) %y_7) + 1) = %y_7
-        ∧ (((u8) %y_7) = (%y_7 + -1))
-        ∧ emp
+          (%y_7 + -1) = ((u8) %y_7) ∧ (((u8) %y_7) = (%y_7 + -1)) ∧ emp
 
-          ((u8) %y_7) = ((u8) (((u8) %y_7) + 1))
-        ∧ (((u8) %y_7) + 1) = %y_7
-        ∧ emp |}]
+          (%y_7 + -1) = ((u8) %y_7) ∧ emp |}]
 
     let%expect_test _ =
       let q =
diff --git a/sledge/lib/term.ml b/sledge/lib/term.ml
index b8b38aa0c..df385eea6 100644
--- a/sledge/lib/term.ml
+++ b/sledge/lib/term.ml
@@ -1188,6 +1188,20 @@ let fold e ~init:s ~f =
   | Add args | Mul args -> Qset.fold ~f:(fun e _ s -> f e s) args ~init:s
   | Var _ | Label _ | Nondet _ | Float _ | Integer _ | Rational _ -> s
 
+let iter_terms e ~f =
+  let iter_terms_ iter_terms_ e =
+    ( match e with
+    | Ap1 (_, x) -> iter_terms_ x
+    | Ap2 (_, x, y) -> iter_terms_ x ; iter_terms_ y
+    | Ap3 (_, x, y, z) -> iter_terms_ x ; iter_terms_ y ; iter_terms_ z
+    | ApN (_, xs) | RecN (_, xs) -> IArray.iter ~f:iter_terms_ xs
+    | Add args | Mul args ->
+        Qset.iter args ~f:(fun arg _ -> iter_terms_ arg)
+    | Var _ | Label _ | Nondet _ | Float _ | Integer _ | Rational _ -> () ) ;
+    f e
+  in
+  fix iter_terms_ (fun _ -> ()) e
+
 let fold_terms e ~init ~f =
   let fold_terms_ fold_terms_ e s =
     let s =
@@ -1205,6 +1219,9 @@ let fold_terms e ~init ~f =
   in
   fix fold_terms_ (fun _ s -> s) e init
 
+let iter_vars e ~f =
+  iter_terms e ~f:(function Var _ as v -> f (v :> Var.t) | _ -> ())
+
 let fold_vars e ~init ~f =
   fold_terms e ~init ~f:(fun s -> function
     | Var _ as v -> f s (v :> Var.t) | _ -> s )
@@ -1231,18 +1248,38 @@ let height e =
 
 (** Solve *)
 
+let find_for ?for_ args =
+  let exists_var args ~f =
+    with_return (fun {return} ->
+        Qset.iter args ~f:(fun arg _ ->
+            iter_vars arg ~f:(fun v -> if f v then return true) ) ;
+        false )
+  in
+  let remove_if_non_occuring rejected args c q =
+    let args = Qset.remove args c in
+    let fv_c = fv c in
+    if exists_var ~f:(Var.Set.mem fv_c) args then None
+    else Some (c, q, Qset.union rejected args)
+  in
+  let rec find_for_ rejected args =
+    let* c, q = Qset.min_elt args in
+    remove_if_non_occuring rejected args c q
+    |> Option.or_else ~f:(fun () ->
+           find_for_ (Qset.add rejected c q) (Qset.remove args c) )
+  in
+  match for_ with
+  | Some c ->
+      let q = Qset.count args c in
+      if Q.equal Q.zero q then None
+      else remove_if_non_occuring Qset.empty args c q
+  | None -> find_for_ Qset.empty args
+
 let solve_zero_eq ?for_ e =
   [%Trace.call fun {pf} -> pf "%a%a" pp e (Option.pp " for %a" pp) for_]
   ;
   ( match e with
   | Add args ->
-      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+ c, q, args = find_for ?for_ 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