Refactor free variable removal

autoprecompiles/src/constraint_optimizer.rs

@@ -315,7 +315,7 @@ fn remove_free_variables<T: FieldElement, V: Clone + Ord + Eq + Hash + Display>(
         .filter(|(variable, constraint)| match constraint {
             // Remove the algebraic constraint if we can solve for the variable.
             ConstraintRef::AlgebraicConstraint(constr) => {
-                can_always_be_satisfied_via_free_variable(*constr, variable)
+                constr.try_find_some_solution(variable).is_some()
             }
             ConstraintRef::BusInteraction(bus_interaction) => {
                 let bus_id = bus_interaction.bus_id.try_to_number().unwrap();
@@ -371,29 +371,6 @@ fn remove_free_variables<T: FieldElement, V: Clone + Ord + Eq + Hash + Display>(
     constraint_system
 }
 
-/// Returns true if the given constraint can always be made to be satisfied by setting the
-/// free variable, regardless of the values of other variables.
-fn can_always_be_satisfied_via_free_variable<
-    T: FieldElement,
-    V: Clone + Hash + Eq + Ord + Display,
->(
-    constraint: AlgebraicConstraint<&GroupedExpression<T, V>>,
-    free_variable: &V,
-) -> bool {
-    if constraint.try_solve_for(free_variable).is_some() {
-        true
-    } else if let Some((left, right)) = constraint.expression.try_as_single_product() {
-        // If either `left` or `right` can be set to 0, the constraint is satisfied.
-        can_always_be_satisfied_via_free_variable(AlgebraicConstraint::from(left), free_variable)
-            || can_always_be_satisfied_via_free_variable(
-                AlgebraicConstraint::from(right),
-                free_variable,
-            )
-    } else {
-        false
-    }
-}
-
 /// Removes any columns that are not connected to *stateful* bus interactions (e.g. memory),
 /// because those are the only way to interact with the rest of the zkVM (e.g. other
 /// instructions).

constraint-solver/src/algebraic_constraint/solve.rs

@@ -104,6 +104,20 @@ where
         Some(subtracted * (-coefficient.field_inverse()))
     }
 
+    /// Like `try_solve_for`, but also handles the case where the constraint is a
+    /// product of two expressions. In this case, the found solution (if any) might
+    /// not be unique.
+    pub fn try_find_some_solution(&self, variable: &V) -> Option<GroupedExpression<T, V>> {
+        if let Some((left, right)) = self.expression.try_as_single_product() {
+            // If either `left` or `right` can be set to 0, the constraint is satisfied.
+            return AlgebraicConstraint::from(left)
+                .try_find_some_solution(variable)
+                .or_else(|| AlgebraicConstraint::from(right).try_find_some_solution(variable));
+        }
+
+        self.try_solve_for(variable)
+    }
+
     /// Algebraically transforms the constraint such that `self = 0` is equivalent
     /// to `expr = result` and returns `result`.
     ///
@@ -596,6 +610,32 @@ mod tests {
         assert!(constr.try_solve_for(&"t").is_none());
     }
 
+    #[test]
+    fn try_find_some_solution() {
+        // Case 1: There is a unique solution for `a` (a = 5)
+        let expr = var("a") - constant(5);
+        let constr = AlgebraicConstraint::assert_zero(&expr);
+        assert_eq!(constr.try_solve_for(&"a").unwrap().to_string(), "5");
+        assert_eq!(
+            constr.try_find_some_solution(&"a").unwrap().to_string(),
+            "5"
+        );
+        assert!(constr.try_find_some_solution(&"t").is_none());
+
+        // Case 1: There are two solutions for `b` (b = 0 or b = 1)
+        let expr = var("b") * (constant(1) - var("b"));
+        let constr = AlgebraicConstraint::assert_zero(&expr);
+        assert!(
+            constr.try_solve_for(&"b").is_none(),
+            "Constraint does not have a unique solution"
+        );
+        assert_eq!(
+            constr.try_find_some_solution(&"b").unwrap().to_string(),
+            "0"
+        );
+        assert!(constr.try_find_some_solution(&"t").is_none());
+    }
+
     #[test]
     fn solve_for_expr() {
         let expr = var("w") + var("x") + constant(3) * var("y") + constant(5);