Replace quadratic equiv

constraint-solver/src/rule_based_optimizer/environment.rs

@@ -4,7 +4,7 @@ use std::{
     hash::Hash,
 };
 
-use itertools::Itertools;
+use itertools::{EitherOrBoth, Itertools};
 use powdr_number::FieldElement;
 
 use crate::{
@@ -185,6 +185,64 @@ impl<T: FieldElement> Environment<T> {
         Some((self.insert(&l), self.insert(&r)))
     }
 
+    /// Returns Some(C) if `a - b = C' and both are affine.
+    pub fn constant_difference(&self, a: Expr, b: Expr) -> Option<T> {
+        let db = self.expressions.borrow();
+        let a = &db[a];
+        let b = &db[b];
+        (a.is_affine()
+            && b.is_affine()
+            && a.linear_components()
+                .zip(b.linear_components())
+                .all(|(x, y)| x == y))
+        .then(|| *a.constant_offset() - *b.constant_offset())
+    }
+
+    /// If this returns Some((v1, v2, factor)), then
+    /// a is obtained from b * factor by substituting v2 by v1.
+    pub fn differ_in_exactly_one_variable(&self, a_id: Expr, b_id: Expr) -> Option<(Var, Var, T)> {
+        let db = self.expressions.borrow();
+        let a = &db[a_id];
+        let b = &db[b_id];
+        if !a.is_affine()
+            || !b.is_affine()
+            || a.linear_components().len() != b.linear_components().len()
+            || a.linear_components().len() < 2
+        {
+            return None;
+        }
+        // First find the variables, ignoring the coefficients.
+        let (v1, v2) = a
+            .linear_components()
+            .merge_join_by(b.linear_components(), |(v1, _), (v2, _)| v1.cmp(v2))
+            .filter(|either| !matches!(either, EitherOrBoth::Both(_, _)))
+            .collect_tuple()?;
+        let (left_var, right_var, factor) = match (v1, v2) {
+            (EitherOrBoth::Left((lv, lc)), EitherOrBoth::Right((rv, rc)))
+            | (EitherOrBoth::Right((rv, rc)), EitherOrBoth::Left((lv, lc))) => {
+                (*lv, *rv, *lc / *rc)
+            }
+            _ => return None,
+        };
+        // Now verify that the other coefficients agree with the factor
+        if *a.constant_offset() != *b.constant_offset() * factor {
+            return None;
+        }
+        if !a
+            .linear_components()
+            .filter(|(v, _)| **v != left_var)
+            .map(|(_, c)| *c)
+            .eq(b
+                .linear_components()
+                .filter(|(v, _)| **v != right_var)
+                .map(|(_, bc)| *bc * factor))
+        {
+            return None;
+        }
+
+        Some((left_var, right_var, factor))
+    }
+
     /// Substitutes the variable `var` by the constant `value` in the expression `e`
     /// and returns the resulting expression.
     pub fn substitute_by_known(&self, e: Expr, var: Var, value: T) -> Expr {

constraint-solver/src/rule_based_optimizer/rules.rs

@@ -75,6 +75,13 @@ crepe! {
       Expression(e),
       for v in env.on_expr(e, (), |e, _| e.referenced_unknown_variables().cloned().collect_vec());
 
+    struct Product(Expr, Expr, Expr);
+    Product(e, l, r) <-
+      Expression(e),
+      Env(env),
+      let Some((l, r)) = env.try_as_single_product(e);
+    Product(e, r, l) <- Product(e, l, r);
+
     // AffineExpression(e, coeff, var, offset) => e = coeff * var + offset
     struct AffineExpression<T: FieldElement>(Expr, T, Var, T);
     AffineExpression(e, coeff, var, offset) <-
@@ -113,6 +120,55 @@ crepe! {
 
     struct Equivalence(Var, Var);
 
+    //------- quadratic equivalence -----
+
+    // QuadraticEquivalenceCandidate(E, expr, offset) =>
+    //   E = ((expr) * (expr + offset) = 0) is a constraint and
+    //   expr is affine with at least 2 variables.
+    struct QuadraticEquivalenceCandidate<T: FieldElement>(Expr, Expr, T);
+    QuadraticEquivalenceCandidate(e, r, offset) <-
+       Env(env),
+       AlgebraicConstraint(e),
+       Product(e, l, r), // note that this will always produce two facts for (l, r) and (r, l)
+       ({env.affine_var_count(l).unwrap_or(0) > 1 && env.affine_var_count(r).unwrap_or(0) > 1}),
+       let Some(offset) = env.constant_difference(l, r);
+
+    // QuadraticEquivalenceCandidatePair(expr1, expr2, offset1 / coeff, v1, v2) =>
+    //  (expr1) * (expr1 + offset1) = 0 and (expr2) * (expr2 + offset2) = 0 are constraints,
+    //  expr1 is affine with at least 2 variables and is obtained from
+    //  expr2 * factor by substituting v2 by v1 (factor != 0),
+    //  offset1 == offset2 * factor and coeff is the coefficient of v1 in expr1.
+    //
+    //  This means that v1 is always equal to (-expr1 / coeff) or equal to
+    //  (-(expr1 + offset1) / coeff) = (-expr1 / coeff - offset1 / coeff).
+    //  Because of the above, also v2 is equal to
+    //  (-expr1 / coeff) or equal to (-(expr1 + offset1) / coeff) [Yes, expr1!].
+    struct QuadraticEquivalenceCandidatePair<T: FieldElement>(Expr, Expr, T, Var, Var);
+    QuadraticEquivalenceCandidatePair(expr1, expr2, offset1 / coeff, v1, v2) <-
+      Env(env),
+      QuadraticEquivalenceCandidate(_, expr1, offset1),
+      QuadraticEquivalenceCandidate(_, expr2, offset2),
+      (expr1 < expr2),
+      let Some((v1, v2,factor)) = env.differ_in_exactly_one_variable(expr1, expr2),
+      (offset1 == offset2 * factor),
+      let coeff = env.on_expr(expr1, (), |e, _| *e.coefficient_of_variable_in_affine_part(&v1).unwrap());
+
+    // QuadraticEquivalence(v1, v2) => v1 and v2 are equal in all satisfying assignments.
+    // Because of QuadraticEquivalenceCandidatePair, v1 is equal to X or X + offset,
+    // where X is some value that depends on other variables. Similarly, v2 is equal to X or X + offset.
+    // Because of the range constraints of v1 and v2, these two "or"s are exclusive ors.
+    // This means depending on the value of X, it is either X or X + offset.
+    // Since this "decision" only depens on X, both v1 and v2 are either X or X + offset at the same time
+    // and thus equal.
+    struct QuadraticEquivalence(Var, Var);
+    QuadraticEquivalence(v1, v2) <-
+      QuadraticEquivalenceCandidatePair(_, _, offset, v1, v2),
+      RangeConstraintOnVar(v1, rc),
+      RangeConstraintOnVar(v2, rc),
+      (rc.is_disjoint(&rc.combine_sum(&RangeConstraint::from_value(offset))));
+
+    Equivalence(v1, v2) <- QuadraticEquivalence(v1, v2);
+
     ReplaceAlgebraicConstraintBy(e, env.substitute_by_known(e, v, val)) <-
       Env(env),
       Assignment(v, val),

constraint-solver/src/rule_based_optimizer/tests.rs

@@ -137,3 +137,87 @@ fn test_rule_based_optimization_simple_assignment() {
     );
     expect!["(y) * (y - 1) - 3 = 0"].assert_eq(&optimized_system.0.to_string());
 }
+
+#[test]
+fn add_with_carry() {
+    // This tests a case of equivalent constraints that appear in the
+    // way "add with carry" is performed in openvm.
+    // X and Y end up being equivalent because they are both either
+    // A or A - 256, depending on whether the value of A is between
+    // 0 and 255 or not.
+    // A is the result of an addition with carry.
+    let mut system = IndexedConstraintSystem::default();
+    system.add_algebraic_constraints([
+        assert_zero(
+            (v("X") * c(7) - v("A") * c(7) + c(256) * c(7)) * (v("X") * c(7) - v("A") * c(7)),
+        ),
+        assert_zero((v("Y") - v("A") + c(256)) * (v("Y") - v("A"))),
+    ]);
+    system.add_bus_interactions([bit_constraint("X", 8), bit_constraint("Y", 8)]);
+    let optimized_system = rule_based_optimization(
+        system,
+        NoRangeConstraints,
+        TestBusInteractionHandler,
+        &mut new_var(),
+        None,
+    );
+    // Y has been replaced by X
+    expect![[r#"
+        (7 * A - 7 * X - 1792) * (7 * A - 7 * X) = 0
+        (A - X - 256) * (A - X) = 0
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: X, 8 }
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: X, 8 }"#]]
+    .assert_eq(&optimized_system.0.to_string());
+}
+
+#[test]
+fn test_rule_based_optimization_quadratic_equality() {
+    let mut system = IndexedConstraintSystem::default();
+    system.add_algebraic_constraints([
+        assert_zero(
+            (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1")
+                - c(30720) * v("mem_ptr_limbs__0_1")
+                + c(737280))
+                * (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1")
+                    - c(30720) * v("mem_ptr_limbs__0_1")
+                    + c(737281)),
+        ),
+        assert_zero(
+            (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1")
+                - c(30720) * v("mem_ptr_limbs__0_2")
+                + c(737280))
+                * (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1")
+                    - c(30720) * v("mem_ptr_limbs__0_2")
+                    + c(737281)),
+        ),
+    ]);
+    system.add_bus_interactions([
+        bit_constraint("rs1_data__0_1", 8),
+        bit_constraint("rs1_data__1_1", 8),
+        BusInteraction {
+            bus_id: c(3),
+            multiplicity: c(1),
+            payload: vec![c(-503316480) * v("mem_ptr_limbs__0_1"), c(14)],
+        },
+        BusInteraction {
+            bus_id: c(3),
+            multiplicity: c(1),
+            payload: vec![c(-503316480) * v("mem_ptr_limbs__0_2"), c(14)],
+        },
+    ]);
+    let optimized_system = rule_based_optimization(
+        system,
+        NoRangeConstraints,
+        TestBusInteractionHandler,
+        &mut new_var(),
+        None,
+    );
+    // Note that in the system below, mem_ptr_limbs__0_2 has been eliminated
+    expect![[r#"
+        (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737280) * (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737281) = 0
+        (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737280) * (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737281) = 0
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: rs1_data__0_1, 8 }
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: rs1_data__1_1, 8 }
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: -(503316480 * mem_ptr_limbs__0_1), 14 }
+        BusInteraction { bus_id: 3, multiplicity: 1, payload: -(503316480 * mem_ptr_limbs__0_1), 14 }"#]].assert_eq(&optimized_system.0.to_string());
+}

constraint-solver/src/solver.rs

@@ -20,7 +20,6 @@ mod boolean_extractor;
 mod constraint_splitter;
 mod exhaustive_search;
 mod linearizer;
-mod quadratic_equivalences;
 mod var_transformation;
 
 /// Solve a constraint system, i.e. derive assignments for variables in the system.

constraint-solver/src/solver/base.rs

@@ -12,7 +12,7 @@ use crate::solver::boolean_extractor::BooleanExtractor;
 use crate::solver::constraint_splitter::try_split_constraint;
 use crate::solver::linearizer::Linearizer;
 use crate::solver::var_transformation::Variable;
-use crate::solver::{exhaustive_search, quadratic_equivalences, Error, Solver, VariableAssignment};
+use crate::solver::{exhaustive_search, Error, Solver, VariableAssignment};
 use crate::utils::possible_concrete_values;
 
 use std::collections::{BTreeSet, HashMap, HashSet};
@@ -326,8 +326,6 @@ where
             let mut progress = false;
             // Try solving constraints in isolation.
             progress |= self.solve_in_isolation()?;
-            // Try to find equivalent variables using quadratic constraints.
-            progress |= self.try_solve_quadratic_equivalences();
 
             if !progress {
                 // This might be expensive, so we only do it if we made no progress
@@ -372,18 +370,6 @@ where
         Ok(progress)
     }
 
-    /// Tries to find equivalent variables using quadratic constraints.
-    fn try_solve_quadratic_equivalences(&mut self) -> bool {
-        let equivalences = quadratic_equivalences::find_quadratic_equalities(
-            self.constraint_system.system().algebraic_constraints(),
-            &*self,
-        );
-        for (x, y) in &equivalences {
-            self.apply_assignment(y, &GroupedExpression::from_unknown_variable(x.clone()));
-        }
-        !equivalences.is_empty()
-    }
-
     /// Find groups of variables with a small set of possible assignments.
     /// For each group, performs an exhaustive search in the possible assignments
     /// to deduce new range constraints (also on other variables).