Transfer quadratic

autoprecompiles/src/constraint_optimizer/equal_zero_checks.rs

@@ -36,6 +36,7 @@ pub fn replace_equal_zero_checks<T: FieldElement, V: Clone + Ord + Hash + Displa
         + Clone,
     new_var: &mut impl FnMut() -> V,
 ) -> IndexedConstraintSystem<T, V> {
+    // println!("Starting equal zero check optimization...\n{constraint_system}");
     let binary_range_constraint = RangeConstraint::from_mask(1);
     // To keep performance reasonable, we split the system at stateful bus interactions
     // into smaller sub-systems and optimize each of them separately.
@@ -71,6 +72,7 @@ pub fn replace_equal_zero_checks<T: FieldElement, V: Clone + Ord + Hash + Displa
             }
         }
     }
+    println!("Final system:\n{constraint_system}");
     constraint_system
 }
 
@@ -163,7 +165,7 @@ fn try_replace_equal_zero_check<T: FieldElement, V: Clone + Ord + Hash + Display
         return;
     }
     // Here we know: `output = value` if and only if all variables in `inputs` are zero.
-    log::debug!(
+    println!(
         "Candidate found: {output} == {value} <=> all of {{{}}} are zero",
         inputs.iter().format(", ")
     );
@@ -212,16 +214,24 @@ fn try_replace_equal_zero_check<T: FieldElement, V: Clone + Ord + Hash + Display
     // Check that each input is non-negative and that the sum of inputs cannot wrap.
     // TODO we do not need this. If this property is false, we can still find
     // a better constraint that also models is-equal-zero.
-    let min_input = inputs
+    let (min_input, max_input) = inputs
         .iter()
-        .map(|v| solver.get(v))
+        .map(|v| {
+            let rc = solver.get(v);
+            println!("{v} : {rc}");
+            rc
+        })
         .reduce(|a, b| a.disjunction(&b))
         .unwrap()
-        .range()
-        .0;
+        .range();
+    println!("Inputs range: {min_input} - {max_input}");
     if min_input != T::from(0) {
         return;
     }
+    //println!("Max input range: {max_input}");
+    if max_input == T::from(1) {
+        println!("Binary inputs in {subsystem}");
+    }
     let sum_of_inputs = inputs
         .iter()
         .map(|v| solver.get(v))
@@ -252,6 +262,10 @@ fn try_replace_equal_zero_check<T: FieldElement, V: Clone + Ord + Hash + Display
 
     // If we remove at most two variables, it is not worth it.
     if variables_to_remove.len() <= 2 {
+        println!(
+            "Not optimizing equal-zero check for {output} since only {} variables would be removed\n{constraint_system}",
+            variables_to_remove.len()
+        );
         return;
     }
 
@@ -332,6 +346,7 @@ fn try_replace_equal_zero_check<T: FieldElement, V: Clone + Ord + Hash + Display
             GroupedExpression::one() - sum_inv.clone() * sum_of_inputs.clone(),
         ),
     ];
+    println!("Adding:\n{}", new_constraints.iter().join("\n"));
     constraint_system.add_algebraic_constraints(new_constraints.clone());
     constraint_system.add_derived_variable(
         sum_inv_var.clone(),

autoprecompiles/src/optimizer.rs

@@ -66,6 +66,7 @@ pub fn optimize<A: Adapter>(
             degree_bound,
         )?
         .clone();
+        println!("======================================================\n\n{constraint_system}\n======================================================");
         if stats == stats_logger::Stats::from(&constraint_system) {
             break;
         }

constraint-solver/src/algebraic_constraint/solve.rs

@@ -65,18 +65,26 @@ where
             return Err(Error::ConstraintUnsatisfiable(self.to_string()));
         }
 
-        if expression.is_quadratic() {
-            self.solve_quadratic(range_constraints)
+        let mut result = if expression.is_quadratic() {
+            self.solve_quadratic(range_constraints)?
         } else if let Some(k) = expression.try_to_known() {
             // If we know `expression` to be nonzero, we should have returned
             // Err already in the range constraint check above.
             assert!(k.is_zero());
             // TODO we could still process more information
             // and reach "unsatisfiable" here.
-            Ok(ProcessResult::complete(vec![]))
+            ProcessResult::complete(vec![])
         } else {
-            self.solve_affine(range_constraints)
+            self.solve_affine(range_constraints)?
+        };
+
+        if !result.complete {
+            // Try to transfer range constraints from the expression.
+            result
+                .effects
+                .extend(self.transfer_constraints(range_constraints));
         }
+        Ok(result)
     }
 
     /// Solves the constraint for `variable`. This is only possible if
@@ -89,6 +97,9 @@ where
     ///
     /// Returns the resulting solved grouped expression.
     pub fn try_solve_for(&self, variable: &V) -> Option<GroupedExpression<T, V>> {
+        if variable.to_string() == "a__0_3" {
+            println!("XX Trying to solve {self} for {variable}");
+        }
         let coefficient = self
             .expression
             .coefficient_of_variable_in_affine_part(variable)?;
@@ -101,6 +112,12 @@ where
             // we cannot solve for it.
             return None;
         }
+        if variable.to_string() == "a__0_3" {
+            println!(
+                "   -> {}",
+                subtracted.clone() * (-coefficient.field_inverse())
+            );
+        }
         Some(subtracted * (-coefficient.field_inverse()))
     }
 
@@ -196,10 +213,7 @@ where
                     ProcessResult::empty()
                 }
             } else {
-                ProcessResult {
-                    effects: self.transfer_constraints(range_constraints),
-                    complete: false,
-                }
+                ProcessResult::empty()
             },
         )
     }
@@ -213,11 +227,13 @@ where
     ) -> Vec<Effect<T, V>> {
         // Solve for each of the variables in the linear component and
         // compute the range constraints.
-        assert!(!self.expression.is_quadratic());
         self.expression
             .linear_components()
             .filter_map(|(var, _)| {
                 let rc = self.try_solve_for(var)?.range_constraint(range_constraints);
+                if var.to_string() == "a__0_3" && rc.range_width() <= 3.into() {
+                    println!("Solving {self} for {var} -> {rc}");
+                }
                 Some((var, rc))
             })
             .filter(|(_, constraint)| !constraint.is_unconstrained())
@@ -666,4 +682,34 @@ mod tests {
             "-(3 * y)"
         );
     }
+
+    #[test]
+    fn rc_for_quadratic() {
+        let cmp_result_1 = var("cmp_result_1");
+        let cmp_result_2 = var("cmp_result_2");
+        let a_0_3 = var("a__0_3");
+        let expr = ((constant(2) * cmp_result_1.clone()) * cmp_result_2.clone())
+            - (constant(2) * cmp_result_1)
+            - (constant(2) * cmp_result_2)
+            + (constant(2) * a_0_3);
+        let rc = HashMap::from([
+            ("cmp_result_1", RangeConstraint::from_mask(1u64)),
+            ("cmp_result_2", RangeConstraint::from_mask(1u64)),
+        ]);
+        // Try it manually first.
+        let constr = AlgebraicConstraint::assert_zero(&expr);
+        let solved = constr.try_solve_for(&"a__0_3").unwrap();
+        println!("-> {solved}");
+        let inferred_rc = solved.range_constraint(&rc);
+        println!("Inferred RC for a__0_3: {inferred_rc}");
+        // should infer binary constraint for "a__0_3"
+        let results = constr.solve(&rc).unwrap();
+        for e in &results.effects {
+            match e {
+                Effect::RangeConstraint(var, rc) => println!("{var}: {rc}"),
+                Effect::Assignment(var, value) => println!("{var} = {value}"),
+                _ => {}
+            }
+        }
+    }
 }

constraint-solver/src/range_constraint.rs

@@ -274,6 +274,11 @@ impl<T: FieldElement> RangeConstraint<T> {
         }
     }
 
+    pub fn negate(&self) -> Self {
+        // TODO correct?
+        Self::from_range(-self.max, -self.min)
+    }
+
     /// If only a single value satisfies this condition, returns this value.
     pub fn try_to_single_value(&self) -> Option<T> {
         if self.min == self.max && self.min.to_integer() & self.mask == self.min.to_integer() {

constraint-solver/src/solver/base.rs

@@ -356,6 +356,17 @@ where
                         .solve(&self.range_constraints)
                         .map_err(Error::AlgebraicSolverError)?
                         .effects;
+                    for e in &effects {
+                        match e {
+                            Effect::Assignment(v, expr) => {
+                                println!("Solved {c} -> {v} := {expr}");
+                            }
+                            Effect::RangeConstraint(v, rc) => {
+                                println!("Solved {c} -> {v} in {rc}");
+                            }
+                            _ => {}
+                        }
+                    }
                     if let Some(components) = try_split_constraint(&c, &self.range_constraints) {
                         progress |= self.add_algebraic_constraints_if_new(components);
                     }
@@ -501,6 +512,9 @@ where
         variable: &V,
         range_constraint: RangeConstraint<T>,
     ) -> bool {
+        if variable.to_string() == "a__0_3" {
+            println!("Update {variable} : {range_constraint}");
+        }
         if self.range_constraints.update(variable, &range_constraint) {
             let new_rc = self.range_constraints.get(variable);
             if let Some(value) = new_rc.try_to_single_value() {

openvm/src/powdr_extension/trace_generator/mod.rs

@@ -183,6 +183,13 @@ impl PowdrTraceGenerator {
 
                 let evaluator = MappingRowEvaluator::new(row_slice, &apc_poly_id_to_index);
 
+                for constr in &self.apc.machine().constraints {
+                    let v = evaluator.eval_expr(&constr.expr);
+                    if !v.is_zero() {
+                        println!("Constraint not satisfied: {constr} = {v} != 0");
+                        panic!();
+                    }
+                }
                 // replay the side effects of this row on the main periphery
                 self.apc
                     .machine()