From 1b58d1e18a1eaeb63e6db29494cda9e908703234 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 18 Dec 2024 20:21:48 +0100 Subject: [PATCH] Introduce concept of assignment. (#2244) --- .../src/witgen/jit/block_machine_processor.rs | 49 ++++--- executor/src/witgen/jit/witgen_inference.rs | 130 +++++++++++++----- 2 files changed, 127 insertions(+), 52 deletions(-) diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index 499c1ee676..61254b9cec 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -47,34 +47,22 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> { let mut witgen = WitgenInference::new(self.fixed_data, self, known_variables); // In the latch row, set the RHS selector to 1. - witgen.assign( - &connection_rhs.selector, - self.latch_row as i32, - T::one().into(), - )?; + witgen.assign_constant(&connection_rhs.selector, self.latch_row as i32, T::one()); - // For each known argument, transfer the value to the expression in the connection's RHS. + // For each argument, connect the expression on the RHS with the formal parameter. for (index, expr) in connection_rhs.expressions.iter().enumerate() { - if known_args[index] { - let param_i = - AffineSymbolicExpression::from_known_symbol(Variable::Param(index), None); - witgen.assign(expr, self.latch_row as i32, param_i)?; - } + witgen.assign_variable(expr, self.latch_row as i32, Variable::Param(index)); } // Solve for the block witness. // Fails if any machine call cannot be completed. self.solve_block(&mut witgen)?; - // For each unknown argument, get the value from the expression in the connection's RHS. - // If assign() fails, it means that we weren't able to to solve for the operation's output. for (index, expr) in connection_rhs.expressions.iter().enumerate() { - if !known_args[index] { - let param_i = - AffineSymbolicExpression::from_unknown_variable(Variable::Param(index), None); - witgen - .assign(expr, self.latch_row as i32, param_i) - .map_err(|_| format!("Could not solve for params[{index}]"))?; + if !witgen.is_known(&Variable::Param(index)) { + return Err(format!( + "Unable to derive algorithm to compute output value \"{expr}\"" + )); } } @@ -265,6 +253,29 @@ params[2] = Add::c[0];" ); } + #[test] + fn unconstrained_output() { + let input = " + namespace Unconstrained(256); + col witness sel, a, b, c; + a + b = 0; + "; + let err_str = generate_for_block_machine( + input, + 1, + 0, + "Unconstrained::sel", + &["Unconstrained::a", "Unconstrained::b"], + &["Unconstrained::c"], + ) + .err() + .unwrap(); + assert_eq!( + err_str, + "Unable to derive algorithm to compute output value \"Unconstrained::c\"" + ); + } + #[test] // TODO: Currently fails, because the machine has a non-rectangular block shape. #[should_panic = "Incomplete machine calls"] diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index f126a1f6ad..bba058b567 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -39,6 +39,8 @@ pub struct WitgenInference<'a, T: FieldElement, FixedEval: FixedEvaluator> { fixed_evaluator: FixedEval, derived_range_constraints: HashMap>, known_variables: HashSet, + /// Internal equalities we were not able to solve yet. + assignments: Vec>, code: Vec>, } @@ -53,6 +55,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F fixed_evaluator, derived_range_constraints: Default::default(), known_variables: known_variables.into_iter().collect(), + assignments: Default::default(), code: Default::default(), } } @@ -61,11 +64,15 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.code } + pub fn is_known(&self, variable: &Variable) -> bool { + self.known_variables.contains(variable) + } + /// Process an identity on a certain row. pub fn process_identity(&mut self, id: &Identity, row_offset: i32) -> ProcessSummary { let result = match id { Identity::Polynomial(PolynomialIdentity { expression, .. }) => { - self.process_polynomial_identity(expression, row_offset) + self.process_equality_on_row(expression, row_offset, T::from(0).into()) } Identity::Lookup(LookupIdentity { id, left, right, .. @@ -89,38 +96,63 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.ingest_effects(result) } - /// Process the constraint that the expression evaluated at the given offset equals the given affine expression. - /// Note that either the expression or the value might contain unknown variables, but if we are not able to - /// solve the equation, we return an error. - pub fn assign( - &mut self, - expression: &Expression, - offset: i32, - value: AffineSymbolicExpression, - ) -> Result<(), String> { - let affine_expression = self - .evaluate(expression, offset) - .ok_or_else(|| format!("Expression is not affine: {expression}"))?; - let result = (affine_expression - value.clone()) - .solve() - .map_err(|err| format!("Could not solve ({expression} - {value}): {err}"))?; - match self.ingest_effects(result).complete { - true => Ok(()), - false => Err("Wasn't able to complete the assignment".to_string()), + /// Turns the given variable either to a known symbolic value or an unknown symbolic value + /// depending on if it is known or not. + /// If it is known to be range-constrained to a single value, that value is used. + fn variable_to_expression(&self, variable: Variable) -> AffineSymbolicExpression { + // If a variable is known and has a compile-time constant value, + // that value is stored in the range constraints. + let rc = self.range_constraint(variable.clone()); + if let Some(val) = rc.as_ref().and_then(|rc| rc.try_to_single_value()) { + val.into() + } else if self.known_variables.contains(&variable) { + AffineSymbolicExpression::from_known_symbol(variable, rc) + } else { + AffineSymbolicExpression::from_unknown_variable(variable, rc) } } - fn process_polynomial_identity( + /// Process the constraint that the expression evaluated at the given offset equals the given value. + /// This does not have to be solvable right away, but is always processed as soon as we have progress. + /// Note that all variables in the expression can be unknown and their status can also change over time. + pub fn assign_constant(&mut self, expression: &'a Expression, row_offset: i32, value: T) { + self.assignments.push(Assignment { + lhs: expression, + row_offset, + rhs: VariableOrValue::Value(value), + }); + self.process_assignments(); + } + + /// Process the constraint that the expression evaluated at the given offset equals the given formal variable. + /// This does not have to be solvable right away, but is always processed as soon as we have progress. + /// Note that all variables in the expression can be unknown and their status can also change over time. + pub fn assign_variable( + &mut self, + expression: &'a Expression, + row_offset: i32, + variable: Variable, + ) { + self.assignments.push(Assignment { + lhs: expression, + row_offset, + rhs: VariableOrValue::Variable(variable), + }); + self.process_assignments(); + } + + fn process_equality_on_row( &self, - expression: &Expression, + lhs: &Expression, offset: i32, + rhs: AffineSymbolicExpression, ) -> ProcessResult { - if let Some(r) = self.evaluate(expression, offset) { + if let Some(r) = self.evaluate(lhs, offset) { // TODO propagate or report error properly. // If solve returns an error, it means that the constraint is conflicting. // In the future, we might run this in a runtime-conditional, so an error // could just mean that this case cannot happen in practice. - r.solve().unwrap() + (r - rhs).solve().unwrap() } else { ProcessResult::empty() } @@ -180,6 +212,31 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F ProcessResult::empty() } + fn process_assignments(&mut self) { + loop { + let mut progress = false; + let new_assignments = std::mem::take(&mut self.assignments) + .into_iter() + .flat_map(|assignment| { + let rhs = match &assignment.rhs { + VariableOrValue::Variable(v) => self.variable_to_expression(v.clone()), + VariableOrValue::Value(v) => (*v).into(), + }; + let r = + self.process_equality_on_row(assignment.lhs, assignment.row_offset, rhs); + let summary = self.ingest_effects(r); + progress |= summary.progress; + // If it is not complete, queue it again. + (!summary.complete).then_some(assignment) + }) + .collect_vec(); + self.assignments.extend(new_assignments); + if !progress { + break; + } + } + } + fn ingest_effects(&mut self, process_result: ProcessResult) -> ProcessSummary { let mut progress = false; for e in process_result.effects { @@ -210,6 +267,9 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F Effect::Assertion(_) => self.code.push(e), } } + if progress { + self.process_assignments(); + } ProcessSummary { complete: process_result.complete, progress, @@ -248,16 +308,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.fixed_evaluator.evaluate(r, offset)?.into() } else { let variable = Variable::from_reference(r, offset); - // If a variable is known and has a compile-time constant value, - // that value is stored in the range constraints. - let rc = self.range_constraint(variable.clone()); - if let Some(val) = rc.as_ref().and_then(|rc| rc.try_to_single_value()) { - val.into() - } else if self.known_variables.contains(&variable) { - AffineSymbolicExpression::from_known_symbol(variable, rc) - } else { - AffineSymbolicExpression::from_unknown_variable(variable, rc) - } + self.variable_to_expression(variable) } } Expression::PublicReference(_) | Expression::Challenge(_) => { @@ -323,6 +374,19 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F } } +/// An equality constraint between an algebraic expression evaluated +/// on a certain row offset and a variable or fixed constant value. +struct Assignment<'a, T: FieldElement> { + lhs: &'a Expression, + row_offset: i32, + rhs: VariableOrValue, +} + +enum VariableOrValue { + Variable(V), + Value(T), +} + pub trait FixedEvaluator { fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { None