Merge remote-tracking branch 'origin/main' into seqz

constraint-solver/src/algebraic_constraint/solve.rs

@@ -50,26 +50,27 @@ where
     /// `variable` does not appear in the quadratic component and
     /// has a coefficient which is known to be not zero.
     ///
+    /// If the constraint has the form `A + k * x = 0` where `A` does not
+    /// contain the variable `x` and `k` is a non-zero runtime constant,
+    /// it returns `A * (-k^(-1))`.
+    ///
     /// Returns the resulting solved grouped expression.
     pub fn try_solve_for(&self, variable: &V) -> Option<GroupedExpression<T, V>> {
-        let expression = self.expression;
+        let coefficient = self
+            .expression
+            .coefficient_of_variable_in_affine_part(variable)?;
+        if !coefficient.is_known_nonzero() {
+            return None;
+        }
 
-        if expression
-            .quadratic
-            .iter()
-            .flat_map(|(l, r)| [l, r])
-            .flat_map(|c| c.referenced_unknown_variables())
-            .contains(variable)
-        {
-            // The variable is in the quadratic component, we cannot solve for it.
+        let subtracted = self.expression.clone()
+            - GroupedExpression::from_unknown_variable(variable.clone()) * coefficient.clone();
+        if subtracted.referenced_unknown_variables().contains(variable) {
+            // There is another occurrence of the variable in the quadratic component,
+            // we cannot solve for it.
             return None;
         }
-        if !expression.linear.get(variable)?.is_known_nonzero() {
-            return None;
-        }
-        let mut result = self.expression.clone();
-        let coefficient = result.linear.remove(variable)?;
-        Some(result * (-T::one().field_div(&coefficient)))
+        Some(subtracted * (-T::one().field_div(coefficient)))
     }
 
     /// Algebraically transforms the constraint such that `self = 0` is equivalent
@@ -95,10 +96,14 @@ where
         let normalization_factor = expr
             .referenced_unknown_variables()
             .find_map(|var| {
-                let coeff = expression.coefficient_of_variable(var)?;
+                let coeff = expression.coefficient_of_variable_in_affine_part(var)?;
                 // We can only divide if we know the coefficient is non-zero.
                 if coeff.is_known_nonzero() {
-                    Some(expr.coefficient_of_variable(var).unwrap().field_div(coeff))
+                    Some(
+                        expr.coefficient_of_variable_in_affine_part(var)
+                            .unwrap()
+                            .field_div(coeff),
+                    )
                 } else {
                     None
                 }
@@ -165,10 +170,10 @@ impl<
         &self,
         range_constraints: &impl RangeConstraintProvider<T::FieldType, V>,
     ) -> Result<ProcessResult<T, V>, Error> {
-        let expression = self.expression;
+        let (_, mut linear, constant) = self.expression.components();
 
-        Ok(if expression.linear.len() == 1 {
-            let (var, coeff) = expression.linear.iter().next().unwrap();
+        Ok(if linear.clone().count() == 1 {
+            let (var, coeff) = linear.next().unwrap();
             // Solve "coeff * X + self.constant = 0" by division.
             assert!(
                 !coeff.is_known_zero(),
@@ -176,16 +181,16 @@ impl<
             );
             if coeff.is_known_nonzero() {
                 // In this case, we can always compute a solution.
-                let value = expression.constant.field_div(&-coeff.clone());
+                let value = constant.field_div(&-coeff.clone());
                 ProcessResult::complete(vec![assignment_if_satisfies_range_constraints(
                     var.clone(),
                     value,
                     range_constraints,
                 )?])
-            } else if expression.constant.is_known_nonzero() {
+            } else if constant.is_known_nonzero() {
                 // If the offset is not zero, then the coefficient must be non-zero,
                 // otherwise the constraint is violated.
-                let value = expression.constant.field_div(&-coeff.clone());
+                let value = constant.field_div(&-coeff.clone());
                 ProcessResult::complete(vec![
                     Assertion::assert_is_nonzero(coeff.clone()),
                     assignment_if_satisfies_range_constraints(
@@ -214,13 +219,12 @@ impl<
         &self,
         range_constraints: &impl RangeConstraintProvider<T::FieldType, V>,
     ) -> Vec<Effect<T, V>> {
-        let expression = self.expression;
         // Solve for each of the variables in the linear component and
         // compute the range constraints.
-        assert!(!expression.is_quadratic());
-        expression
-            .linear
-            .iter()
+        assert!(!self.expression.is_quadratic());
+        self.expression
+            .components()
+            .1
             .filter_map(|(var, _)| {
                 let rc = self.try_solve_for(var)?.range_constraint(range_constraints);
                 Some((var, rc))
@@ -858,7 +862,7 @@ mod tests {
 
         expr.substitute_by_unknown(&"x", &subst);
         assert!(!expr.is_quadratic());
-        assert_eq!(expr.linear.len(), 3);
+        assert_eq!(expr.components().1.count(), 3);
         assert_eq!(expr.to_string(), "a + b + y");
     }
 
@@ -880,14 +884,17 @@ mod tests {
         );
         // Structural validation
         assert_eq!(quadratic.len(), 2);
-        assert_eq!(quadratic[0].0.to_string(), "(a) * (b) + 1");
-        assert_eq!(quadratic[0].0.quadratic[0].0.to_string(), "a");
-        assert_eq!(quadratic[0].0.quadratic[0].1.to_string(), "b");
-        assert!(quadratic[0].0.linear.is_empty());
-        assert_eq!(
-            quadratic[0].0.constant.try_to_number(),
-            Some(GoldilocksField::from(1)),
-        );
+        {
+            assert_eq!(quadratic[0].0.to_string(), "(a) * (b) + 1");
+            let (inner_quadratic, inner_linear, inner_constant) = quadratic[0].0.components();
+            assert_eq!(inner_quadratic[0].0.to_string(), "a");
+            assert_eq!(inner_quadratic[0].1.to_string(), "b");
+            assert!(inner_linear.count() == 0);
+            assert_eq!(
+                inner_constant.try_to_number(),
+                Some(GoldilocksField::from(1)),
+            );
+        }
         assert_eq!(quadratic[0].1.to_string(), "w");
         assert_eq!(quadratic[1].0.to_string(), "a");
         assert_eq!(quadratic[1].1.to_string(), "b");

constraint-solver/src/constraint_system.rs

@@ -196,7 +196,10 @@ impl<
                     // `k * var + e` is in range rc <=>
                     // `var` is in range `(rc - RC[e]) / k` = `rc / k + RC[-e / k]`
                     // If we solve `expr` for `var`, we get `-e / k`.
-                    let k = expr.coefficient_of_variable(var).unwrap().try_to_number()?;
+                    let k = expr
+                        .coefficient_of_variable_in_affine_part(var)
+                        .unwrap()
+                        .try_to_number()?;
                     let expr = AlgebraicConstraint::assert_zero(expr).try_solve_for(var)?;
                     let rc = rc
                         .multiple(T::FieldType::from(1) / k)

constraint-solver/src/grouped_expression.rs

@@ -39,12 +39,12 @@ pub struct GroupedExpression<T, V> {
     /// Quadratic terms of the form `a * X * Y`, where `a` is a (symbolically)
     /// known value and `X` and `Y` are grouped expressions that
     /// have at least one unknown.
-    pub(crate) quadratic: Vec<(Self, Self)>,
+    quadratic: Vec<(Self, Self)>,
     /// Linear terms of the form `a * X`, where `a` is a (symbolically) known
     /// value and `X` is an unknown variable.
-    pub(crate) linear: BTreeMap<V, T>,
+    linear: BTreeMap<V, T>,
     /// Constant term, a (symbolically) known value.
-    pub(crate) constant: T,
+    constant: T,
 }
 
 impl<F: FieldElement, T: RuntimeConstant<FieldType = F>, V> GroupedExpression<T, V> {
@@ -236,10 +236,13 @@ impl<T: RuntimeConstant, V: Ord + Clone + Eq> GroupedExpression<T, V> {
             .unwrap()
     }
 
-    /// Returns the coefficient of the variable `variable` if this is an affine expression.
-    /// Panics if the expression is quadratic.
-    pub fn coefficient_of_variable<'a>(&'a self, var: &V) -> Option<&'a T> {
-        assert!(!self.is_quadratic());
+    /// Returns the coefficient of the variable `variable` in the affine part of this
+    /// expression.
+    /// If the expression is affine, this is the actual coefficient of the variable
+    /// in the expression. Otherwise, the quadratic part of the expression could
+    /// also contain the variable and thus the actual coefficient might be different
+    /// (even zero).
+    pub fn coefficient_of_variable_in_affine_part<'a>(&'a self, var: &V) -> Option<&'a T> {
         self.linear.get(var)
     }
 

constraint-solver/src/solver/quadratic_equivalences.rs

@@ -128,7 +128,10 @@ impl<T: RuntimeConstant, V: Ord + Clone + Hash + Eq> QuadraticEqualityCandidate<
     /// Returns an equivalent candidate that is normalized
     /// such that `var` has a coefficient of `1`.
     fn normalized_for_var(&self, var: &V) -> Self {
-        let coefficient = self.expr.coefficient_of_variable(var).unwrap();
+        let coefficient = self
+            .expr
+            .coefficient_of_variable_in_affine_part(var)
+            .unwrap();
 
         // self represents
         // `(coeff * var + X) * (coeff * var + X + offset) = 0`