From aeec462f4d950ae9d892ddd3ed989e7f6410adf8 Mon Sep 17 00:00:00 2001 From: chriseth Date: Mon, 9 Dec 2024 14:40:31 +0000 Subject: [PATCH 01/44] Expressions and solving routines. --- .../witgen/jit/affine_symbolic_expression.rs | 393 ++++++++++++++++++ executor/src/witgen/jit/mod.rs | 3 + .../src/witgen/jit/symbolic_expression.rs | 262 ++++++++++++ executor/src/witgen/mod.rs | 1 + 4 files changed, 659 insertions(+) create mode 100644 executor/src/witgen/jit/affine_symbolic_expression.rs create mode 100644 executor/src/witgen/jit/mod.rs create mode 100644 executor/src/witgen/jit/symbolic_expression.rs diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs new file mode 100644 index 0000000000..8474648f66 --- /dev/null +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -0,0 +1,393 @@ +use std::{ + collections::BTreeMap, + fmt::{self, Display, Formatter}, + ops::{Add, Mul, Neg, Sub}, +}; + +use itertools::Itertools; +use num_traits::Zero; +use powdr_number::FieldElement; + +use crate::witgen::global_constraints::RangeConstraintSet; + +use super::{super::range_constraints::RangeConstraint, symbolic_expression::SymbolicExpression}; + +/// The effect of solving a symbolic equation. +pub enum Effect { + Assignment(V, SymbolicExpression), + RangeConstraint(V, RangeConstraint), + Code(String), +} + +/// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, +/// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time, +/// and the `x_i` are unknown variables. +#[derive(Debug, Clone)] +pub struct AffineSymbolicExpression { + coefficients: BTreeMap>, + offset: SymbolicExpression, +} + +impl Display for AffineSymbolicExpression { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + if self.coefficients.is_empty() { + write!(f, "{}", self.offset) + } else { + write!( + f, + "{}", + self.coefficients + .iter() + .map(|(var, coeff)| if coeff.is_known_one() { + var.to_string() + } else if coeff.is_known_minus_one() { + format!("-{var}") + } else { + format!("{coeff} * {var}") + }) + .join(" + ") + )?; + if !self.offset.is_known_zero() { + write!(f, " + {}", self.offset)?; + } + Ok(()) + } + } +} + +impl From> for AffineSymbolicExpression { + fn from(k: SymbolicExpression) -> Self { + AffineSymbolicExpression { + coefficients: Default::default(), + offset: k, + } + } +} + +impl AffineSymbolicExpression { + pub fn from_known_variable(var: &str) -> Self { + SymbolicExpression::from_var(var).into() + } + pub fn from_unknown_variable(var: V) -> Self { + AffineSymbolicExpression { + coefficients: [(var, T::from(1).into())].into_iter().collect(), + offset: SymbolicExpression::from(T::from(0)), + } + } + pub fn from_number(n: T) -> Self { + SymbolicExpression::from(n).into() + } + + pub fn is_known_zero(&self) -> bool { + self.coefficients.is_empty() && self.offset.is_known_zero() + } + + pub fn is_known_one(&self) -> bool { + self.coefficients.is_empty() && self.offset.is_known_one() + } + + pub fn is_known(&self) -> bool { + self.coefficients.is_empty() + } + + /// Tries to multiply this expression with another one. + /// Returns `None` if the result would be quadratic. + pub fn try_mul(&self, other: &Self) -> Option { + if self.is_known_zero() || other.is_known_zero() { + return Some(AffineSymbolicExpression::from_number(T::from(0))); + } + if !self.coefficients.is_empty() && !other.coefficients.is_empty() { + return None; + } + let (multiplier, coefficients, offset) = if self.coefficients.is_empty() { + (&self.offset, &other.coefficients, &other.offset) + } else { + (&other.offset, &self.coefficients, &self.offset) + }; + let coefficients = coefficients + .iter() + .map(|(var, coeff)| (var.clone(), coeff * multiplier)) + .collect(); + let offset = offset * multiplier; + Some(AffineSymbolicExpression { + coefficients, + offset, + }) + } + + /// Solves the equation `self = 0` and returns how to compute the solution. + pub fn solve(&self, range_constraints: &impl RangeConstraintSet) -> Vec> { + match self.coefficients.len() { + 0 => { + return if self.offset.is_known_zero() { + vec![] + } else { + // TODO this is a non-satisfiable constraint - should we panic? + vec![] + }; + } + 1 => { + let (var, coeff) = self.coefficients.iter().next().unwrap(); + assert!(!coeff.is_known_zero()); + let value = &self.offset / &coeff.neg(); + return vec![Effect::Assignment(var.clone(), value)]; + } + _ => {} + } + + let r = self.solve_through_constraints(range_constraints); + if !r.is_empty() { + return r; + } + let negated = -self; + let r = negated.solve_through_constraints(range_constraints); + if !r.is_empty() { + return r; + } + self.transfer_constraints(range_constraints) + .into_iter() + .chain(self.transfer_constraints(range_constraints)) + .collect() + } + + fn solve_through_constraints( + &self, + range_constraints: &impl RangeConstraintSet, + ) -> Vec> { + let constrained_coefficients = self + .coefficients + .iter() + .filter_map(|(var, coeff)| { + coeff.try_to_number().and_then(|c| { + range_constraints + .range_constraint(var.clone()) + .map(|rc| (var.clone(), c, rc)) + }) + }) + .collect_vec(); + + // All the coefficients need to have known range constraints. + if constrained_coefficients.len() != self.coefficients.len() { + return Vec::new(); + } + + // Check if they are mutually exclusive and compute assignments. + let mut covered_bits: ::Integer = 0.into(); + let mut assignments = vec![]; + for (var, coeff, constraint) in constrained_coefficients { + let mask = *constraint.multiple(coeff).mask(); + if !(mask & covered_bits).is_zero() { + // Overlapping range constraints. + return vec![]; + } else { + covered_bits |= mask; + } + let mask = T::from(mask); + let masked = &(-&self.offset) & &mask.into(); + let rhs = masked.integer_div(&coeff.into()); + assignments.push(Effect::Assignment(var.clone(), rhs)); + } + + if covered_bits >= T::modulus() { + return vec![]; + } + + // TODO we need to add an assertion that offset is covered by the masks. + // Otherwise the equation is not solvable. + // TODO is this really the case? + assignments + } + + fn transfer_constraints( + &self, + range_constraints: &impl RangeConstraintSet, + ) -> Vec> { + // We are looking for X = a * Y + b * Z + ... or -X = a * Y + b * Z + ... + // where X is least constrained. + + let Some((solve_for, solve_for_coefficient)) = self + .coefficients + .iter() + .filter(|(_var, coeff)| coeff.is_known_one() || coeff.is_known_minus_one()) + .max_by_key(|(var, _c)| { + // Sort so that we get the least constrained variable. + range_constraints + .range_constraint((*var).clone()) + .map(|c| c.range_width()) + .unwrap_or_else(|| T::modulus()) + }) + else { + return vec![]; + }; + + // This only works if the coefficients are all known. + let Some(summands) = self + .coefficients + .iter() + .filter(|(var, _)| *var != solve_for) + .map(|(var, coeff)| { + coeff.try_to_number().and_then(|coeff| { + range_constraints + .range_constraint(var.clone()) + .map(|constraint| constraint.multiple(coeff)) + }) + }) + .chain(std::iter::once(self.offset.range_constraint())) + .collect::>>() + else { + return vec![]; + }; + let Some(constraint) = summands.into_iter().reduce(|c1, c2| c1.combine_sum(&c2)) else { + return vec![]; + }; + let constraint = if solve_for_coefficient.is_known_one() { + -constraint + } else { + constraint + }; + vec![Effect::RangeConstraint(solve_for.clone(), constraint)] + } +} + +impl Add for &AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn add(self, rhs: Self) -> Self::Output { + // TODO could iterate in parallel. + let mut coefficients = self.coefficients.clone(); + for (var, coeff) in rhs.coefficients.iter() { + coefficients + .entry(var.clone()) + .and_modify(|f| *f = &*f + coeff) + .or_insert_with(|| coeff.clone()); + } + let coefficients = coefficients + .into_iter() + .filter(|(_, f)| !f.is_known_zero()) + .collect(); + let offset = &self.offset + &rhs.offset; + AffineSymbolicExpression { + coefficients, + offset, + } + } +} + +impl Sub for &AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn sub(self, rhs: Self) -> Self::Output { + self + &-rhs + } +} + +impl Neg for &AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn neg(self) -> Self::Output { + AffineSymbolicExpression { + coefficients: self + .coefficients + .iter() + .map(|(var, coeff)| (var.clone(), -coeff)) + .collect(), + offset: -&self.offset, + } + } +} + +impl TryFrom<&AffineSymbolicExpression> for SymbolicExpression { + type Error = (); + + fn try_from(value: &AffineSymbolicExpression) -> Result { + if value.coefficients.is_empty() { + Ok(value.offset.clone()) + } else { + Err(()) + } + } +} + +#[cfg(test)] +mod test { + use powdr_number::GoldilocksField; + + use super::*; + + #[derive(Default)] + struct SimpleRangeConstraintSet(BTreeMap<&'static str, RangeConstraint>); + impl RangeConstraintSet<&str, GoldilocksField> for SimpleRangeConstraintSet { + fn range_constraint(&self, var: &str) -> Option> { + self.0.get(var).cloned() + } + } + + type ASE = AffineSymbolicExpression; + + fn from_number(x: i32) -> ASE { + ASE::from_number(GoldilocksField::from(x)) + } + + fn mul(a: &ASE, b: &ASE) -> ASE { + a.try_mul(b).unwrap() + } + + #[test] + fn solve_simple_eq() { + let y = ASE::from_known_variable("y"); + let x = ASE::from_unknown_variable("X"); + // 2 * X + 7 * y - 10 = 0 + let two = from_number(2); + let seven = from_number(7); + let ten = from_number(10); + let constr = &(&mul(&two, &x) + &mul(&seven, &y)) - &ten; + let effects = constr.solve(&SimpleRangeConstraintSet::default()); + assert_eq!(effects.len(), 1); + let Effect::Assignment(var, expr) = &effects[0] else { + panic!("Expected assignment"); + }; + assert_eq!(var.to_string(), "X"); + assert_eq!(expr.to_string(), "(((y * 7) + -10) / -2)"); + } + + #[test] + fn solve_bit_decomposition() { + let a = ASE::from_unknown_variable("a"); + let b = ASE::from_unknown_variable("b"); + let c = ASE::from_unknown_variable("c"); + let z = ASE::from_known_variable("Z"); + let range_constraints = SimpleRangeConstraintSet( + ["a", "b", "c"] + .into_iter() + .map(|var| (var, RangeConstraint::from_mask(0xffu32))) + .collect(), + ); + // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 + Z = 0 + let ten = from_number(10); + let constr = &(&(&(&mul(&a, &from_number(0x100)) + &mul(&b, &from_number(0x10000))) + + &mul(&c, &from_number(0x1000000))) + + &ten) + + &z; + // Without range constraints, this is not solvable. + assert!(constr + .solve(&SimpleRangeConstraintSet::default()) + .is_empty()); + // With range constraints, it should be solvable. + let effects = constr + .solve(&range_constraints) + .into_iter() + .map(|effect| match effect { + Effect::Assignment(v, expr) => format!("{v} = {expr};\n"), + _ => panic!(), + }) + .format("") + .to_string(); + assert_eq!( + effects, + "a = ((-((10 + Z)) & 65280) // 256); +b = ((-((10 + Z)) & 16711680) // 65536); +c = ((-((10 + Z)) & 4278190080) // 16777216); +" + ); + } +} diff --git a/executor/src/witgen/jit/mod.rs b/executor/src/witgen/jit/mod.rs new file mode 100644 index 0000000000..08196d529c --- /dev/null +++ b/executor/src/witgen/jit/mod.rs @@ -0,0 +1,3 @@ +mod affine_symbolic_expression; +mod cell; +mod symbolic_expression; diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs new file mode 100644 index 0000000000..83601df668 --- /dev/null +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -0,0 +1,262 @@ +use std::{ + fmt::{self, Display, Formatter}, + ops::{Add, BitAnd, Div, Mul, Neg}, +}; + +use powdr_number::FieldElement; + +use crate::witgen::range_constraints::RangeConstraint; + +/// A value that is known at run-time, defined through a complex expression +/// involving known cells or variables and compile-time constants. +/// Each of the sub-expressions can have its own range constraint. +#[derive(Debug, Clone)] +pub enum SymbolicExpression { + /// A concrete constant value known at compile time. + Concrete(T), + /// A symbolic value known at run-time, referencing either a cell or a local variable. + Variable(String, Option>), + BinaryOperation( + Box>, + BinaryOperator, + Box>, + Option>, + ), + UnaryOperation( + UnaryOperator, + Box>, + Option>, + ), +} + +#[derive(Debug, Clone)] +pub enum BinaryOperator { + Add, + Sub, + Mul, + /// Finite field division. + Div, + /// Integer division, i.e. convert field elements to unsigned integer and divide. + IntegerDiv, + BitAnd, +} + +#[derive(Debug, Clone)] +pub enum UnaryOperator { + Neg, +} + +impl SymbolicExpression { + pub fn is_known_zero(&self) -> bool { + self.try_to_number().map_or(false, |n| n.is_zero()) + } + + pub fn is_known_one(&self) -> bool { + self.try_to_number().map_or(false, |n| n.is_one()) + } + + pub fn is_known_minus_one(&self) -> bool { + self.try_to_number().map_or(false, |n| n == -T::from(1)) + } + + pub fn range_constraint(&self) -> Option> { + match self { + SymbolicExpression::Concrete(v) => Some(RangeConstraint::from_value(*v)), + SymbolicExpression::Variable(.., rc) + | SymbolicExpression::BinaryOperation(.., rc) + | SymbolicExpression::UnaryOperation(.., rc) => rc.clone(), + } + } + + pub fn try_to_number(&self) -> Option { + match self { + SymbolicExpression::Concrete(n) => Some(*n), + SymbolicExpression::Variable(..) + | SymbolicExpression::BinaryOperation(..) + | SymbolicExpression::UnaryOperation(..) => None, + } + } +} + +impl Display for SymbolicExpression { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + SymbolicExpression::Concrete(n) => { + if n.is_in_lower_half() { + write!(f, "{n}") + } else { + write!(f, "-{}", -*n) + } + } + SymbolicExpression::Variable(name, _) => write!(f, "{name}"), + SymbolicExpression::BinaryOperation(lhs, op, rhs, _) => { + write!(f, "({lhs} {op} {rhs})") + } + SymbolicExpression::UnaryOperation(op, expr, _) => write!(f, "{op}({expr})"), + } + } +} + +impl Display for BinaryOperator { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + BinaryOperator::Add => write!(f, "+"), + BinaryOperator::Sub => write!(f, "-"), + BinaryOperator::Mul => write!(f, "*"), + BinaryOperator::Div => write!(f, "/"), + BinaryOperator::IntegerDiv => write!(f, "//"), + BinaryOperator::BitAnd => write!(f, "&"), + } + } +} + +impl Display for UnaryOperator { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + UnaryOperator::Neg => write!(f, "-"), + } + } +} + +impl SymbolicExpression { + pub fn from_var(name: &str) -> Self { + SymbolicExpression::Variable(name.to_string(), None) + } +} + +impl From for SymbolicExpression { + fn from(n: T) -> Self { + SymbolicExpression::Concrete(n) + } +} + +impl Add for &SymbolicExpression { + type Output = SymbolicExpression; + + fn add(self, rhs: Self) -> Self::Output { + if self.is_known_zero() { + return rhs.clone(); + } + if rhs.is_known_zero() { + return self.clone(); + } + match (self, rhs) { + (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) => { + SymbolicExpression::Concrete(*a + *b) + } + _ => SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::Add, + Box::new(rhs.clone()), + self.range_constraint() + .zip(rhs.range_constraint()) + .map(|(a, b)| a.combine_sum(&b)), + ), + } + } +} + +impl Neg for &SymbolicExpression { + type Output = SymbolicExpression; + + fn neg(self) -> Self::Output { + match self { + SymbolicExpression::Concrete(n) => SymbolicExpression::Concrete(-*n), + SymbolicExpression::UnaryOperation(UnaryOperator::Neg, expr, _) => *expr.clone(), + _ => SymbolicExpression::UnaryOperation( + UnaryOperator::Neg, + Box::new(self.clone()), + self.range_constraint().map(|rc| rc.multiple(-T::from(1))), + ), + } + } +} + +impl Mul for &SymbolicExpression { + type Output = SymbolicExpression; + + fn mul(self, rhs: Self) -> Self::Output { + if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { + SymbolicExpression::Concrete(*a * *b) + } else if self.is_known_zero() || rhs.is_known_zero() { + SymbolicExpression::Concrete(T::from(0)) + } else if self.is_known_one() { + rhs.clone() + } else if rhs.is_known_one() { + self.clone() + } else if self.is_known_minus_one() { + -rhs + } else if rhs.is_known_minus_one() { + -self + } else { + SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::Mul, + Box::new(rhs.clone()), + None, + ) + } + } +} + +impl Div for &SymbolicExpression { + type Output = SymbolicExpression; + + /// Field element division. See `integer_div` for integer division. + fn div(self, rhs: Self) -> Self::Output { + if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { + assert!(b != &T::from(0)); + SymbolicExpression::Concrete(*a / *b) + } else if self.is_known_zero() { + // TODO should we still detect division by zero in this case? + SymbolicExpression::Concrete(T::from(0)) + } else if rhs.is_known_one() { + self.clone() + } else if rhs.is_known_minus_one() { + -self + } else { + // TODO other simplifications like `-x / -y => x / y`, `-x / concrete => x / -concrete`, etc. + SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::Div, + Box::new(rhs.clone()), + None, + ) + } + } +} + +impl SymbolicExpression { + pub fn integer_div(&self, rhs: &Self) -> Self { + if rhs.is_known_one() { + self.clone() + } else { + SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::IntegerDiv, + Box::new(rhs.clone()), + None, + ) + } + } +} + +impl BitAnd for &SymbolicExpression { + type Output = SymbolicExpression; + + fn bitand(self, rhs: Self) -> Self::Output { + if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { + SymbolicExpression::Concrete(T::from(a.to_integer() & b.to_integer())) + } else { + // TODO simplifications? + SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::BitAnd, + Box::new(rhs.clone()), + self.range_constraint() + .zip(rhs.range_constraint()) + .map(|(a, b)| a.conjunction(&b)), + ) + } + } +} diff --git a/executor/src/witgen/mod.rs b/executor/src/witgen/mod.rs index 5640c38e8e..ea6c62014e 100644 --- a/executor/src/witgen/mod.rs +++ b/executor/src/witgen/mod.rs @@ -33,6 +33,7 @@ mod expression_evaluator; pub mod fixed_evaluator; mod global_constraints; mod identity_processor; +mod jit; mod machines; mod processor; mod query_processor; From 3749925f352b054a35ec94f5be6e94fb37c98c89 Mon Sep 17 00:00:00 2001 From: chriseth Date: Mon, 9 Dec 2024 16:32:09 +0000 Subject: [PATCH 02/44] Fixes and one more test. --- .../witgen/jit/affine_symbolic_expression.rs | 75 ++++++++++++++----- executor/src/witgen/jit/mod.rs | 1 - 2 files changed, 56 insertions(+), 20 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 8474648f66..8e91e7399c 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -1,7 +1,7 @@ use std::{ collections::BTreeMap, fmt::{self, Display, Formatter}, - ops::{Add, Mul, Neg, Sub}, + ops::{Add, Neg, Sub}, }; use itertools::Itertools; @@ -14,9 +14,10 @@ use super::{super::range_constraints::RangeConstraint, symbolic_expression::Symb /// The effect of solving a symbolic equation. pub enum Effect { + /// variable can be assigned a value. Assignment(V, SymbolicExpression), + /// we learnt a new range constraint on variable. RangeConstraint(V, RangeConstraint), - Code(String), } /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, @@ -77,15 +78,12 @@ impl AffineSymbolicExpression { pub fn from_number(n: T) -> Self { SymbolicExpression::from(n).into() } - pub fn is_known_zero(&self) -> bool { self.coefficients.is_empty() && self.offset.is_known_zero() } - pub fn is_known_one(&self) -> bool { self.coefficients.is_empty() && self.offset.is_known_one() } - pub fn is_known(&self) -> bool { self.coefficients.is_empty() } @@ -150,6 +148,7 @@ impl AffineSymbolicExpression { .collect() } + /// Tries to solve a bit-decomposition equation. fn solve_through_constraints( &self, range_constraints: &impl RangeConstraintSet, @@ -182,10 +181,11 @@ impl AffineSymbolicExpression { } else { covered_bits |= mask; } - let mask = T::from(mask); - let masked = &(-&self.offset) & &mask.into(); - let rhs = masked.integer_div(&coeff.into()); - assignments.push(Effect::Assignment(var.clone(), rhs)); + let masked = &(-&self.offset) & &T::from(mask).into(); + assignments.push(Effect::Assignment( + var.clone(), + masked.integer_div(&coeff.into()), + )); } if covered_bits >= T::modulus() { @@ -322,20 +322,20 @@ mod test { } } - type ASE = AffineSymbolicExpression; + type Ase = AffineSymbolicExpression; - fn from_number(x: i32) -> ASE { - ASE::from_number(GoldilocksField::from(x)) + fn from_number(x: i32) -> Ase { + Ase::from_number(GoldilocksField::from(x)) } - fn mul(a: &ASE, b: &ASE) -> ASE { + fn mul(a: &Ase, b: &Ase) -> Ase { a.try_mul(b).unwrap() } #[test] fn solve_simple_eq() { - let y = ASE::from_known_variable("y"); - let x = ASE::from_unknown_variable("X"); + let y = Ase::from_known_variable("y"); + let x = Ase::from_unknown_variable("X"); // 2 * X + 7 * y - 10 = 0 let two = from_number(2); let seven = from_number(7); @@ -352,10 +352,10 @@ mod test { #[test] fn solve_bit_decomposition() { - let a = ASE::from_unknown_variable("a"); - let b = ASE::from_unknown_variable("b"); - let c = ASE::from_unknown_variable("c"); - let z = ASE::from_known_variable("Z"); + let a = Ase::from_unknown_variable("a"); + let b = Ase::from_unknown_variable("b"); + let c = Ase::from_unknown_variable("c"); + let z = Ase::from_known_variable("Z"); let range_constraints = SimpleRangeConstraintSet( ["a", "b", "c"] .into_iter() @@ -387,6 +387,43 @@ mod test { "a = ((-((10 + Z)) & 65280) // 256); b = ((-((10 + Z)) & 16711680) // 65536); c = ((-((10 + Z)) & 4278190080) // 16777216); +" + ); + } + + #[test] + fn solve_constraint_transfer() { + let a = Ase::from_unknown_variable("a"); + let b = Ase::from_unknown_variable("b"); + let c = Ase::from_unknown_variable("c"); + let z = Ase::from_unknown_variable("Z"); + let range_constraints = SimpleRangeConstraintSet( + ["a", "b", "c"] + .into_iter() + .map(|var| (var, RangeConstraint::from_mask(0xffu32))) + .collect(), + ); + // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 - Z = 0 + let ten = from_number(10); + let constr = &(&(&(&mul(&a, &from_number(0x100)) + &mul(&b, &from_number(0x10000))) + + &mul(&c, &from_number(0x1000000))) + + &ten) + - &z; + let effects = constr + .solve(&range_constraints) + .into_iter() + .map(|effect| match effect { + Effect::RangeConstraint(v, rc) => format!("{v}: {rc};\n"), + _ => panic!(), + }) + .format("") + .to_string(); + // It appears twice because we solve the positive and the negated equation. + // Maybe it is enough to only solve one. + assert_eq!( + effects, + "Z: [10, 4294967050] & 0xffffff0a; +Z: [10, 4294967050] & 0xffffff0a; " ); } diff --git a/executor/src/witgen/jit/mod.rs b/executor/src/witgen/jit/mod.rs index 08196d529c..a6d9dabf39 100644 --- a/executor/src/witgen/jit/mod.rs +++ b/executor/src/witgen/jit/mod.rs @@ -1,3 +1,2 @@ mod affine_symbolic_expression; -mod cell; mod symbolic_expression; From 41d12702bd34ea6dc01d93bc98cf9a5995f0b476 Mon Sep 17 00:00:00 2001 From: chriseth Date: Mon, 9 Dec 2024 16:38:37 +0000 Subject: [PATCH 03/44] Remove comment. --- executor/src/witgen/jit/affine_symbolic_expression.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 8e91e7399c..bdcdf83c6b 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -253,7 +253,6 @@ impl Add for &AffineSymbolicExpression { type Output = AffineSymbolicExpression; fn add(self, rhs: Self) -> Self::Output { - // TODO could iterate in parallel. let mut coefficients = self.coefficients.clone(); for (var, coeff) in rhs.coefficients.iter() { coefficients From f5e5edd7defe2926ec5de0f98e3b4270117c0444 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 08:56:58 +0000 Subject: [PATCH 04/44] Add assertions. --- .../witgen/jit/affine_symbolic_expression.rs | 128 +++++++++++++++--- .../src/witgen/jit/symbolic_expression.rs | 37 +++-- 2 files changed, 137 insertions(+), 28 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index bdcdf83c6b..1019642c02 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -8,7 +8,7 @@ use itertools::Itertools; use num_traits::Zero; use powdr_number::FieldElement; -use crate::witgen::global_constraints::RangeConstraintSet; +use crate::witgen::{global_constraints::RangeConstraintSet, EvalError}; use super::{super::range_constraints::RangeConstraint, symbolic_expression::SymbolicExpression}; @@ -18,6 +18,40 @@ pub enum Effect { Assignment(V, SymbolicExpression), /// we learnt a new range constraint on variable. RangeConstraint(V, RangeConstraint), + /// a run-time assertion. If this fails, we have conflicting constraints. + Assertion(Assertion), +} + +/// A run-time assertion. If this fails, we have conflicting constraints. +pub struct Assertion { + pub lhs: SymbolicExpression, + pub rhs: SymbolicExpression, + /// If this is true, we assert that both sides are equal. + /// Otherwise, we assert that they are different. + pub expected_equal: bool, +} + +impl Assertion { + pub fn assert_is_zero(condition: SymbolicExpression) -> Effect { + Self::assert_eq(condition, SymbolicExpression::from(T::from(0))) + } + pub fn assert_is_nonzero(condition: SymbolicExpression) -> Effect { + Self::assert_neq(condition, SymbolicExpression::from(T::from(0))) + } + pub fn assert_eq(lhs: SymbolicExpression, rhs: SymbolicExpression) -> Effect { + Effect::Assertion(Assertion { + lhs, + rhs, + expected_equal: true, + }) + } + pub fn assert_neq(lhs: SymbolicExpression, rhs: SymbolicExpression) -> Effect { + Effect::Assertion(Assertion { + lhs, + rhs, + expected_equal: false, + }) + } } /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, @@ -65,7 +99,7 @@ impl From> for AffineSymbolicExpressio } } -impl AffineSymbolicExpression { +impl AffineSymbolicExpression { pub fn from_known_variable(var: &str) -> Self { SymbolicExpression::from_var(var).into() } @@ -114,38 +148,50 @@ impl AffineSymbolicExpression { } /// Solves the equation `self = 0` and returns how to compute the solution. - pub fn solve(&self, range_constraints: &impl RangeConstraintSet) -> Vec> { + pub fn solve( + &self, + range_constraints: &impl RangeConstraintSet, + ) -> Result>, EvalError> { match self.coefficients.len() { 0 => { return if self.offset.is_known_zero() { - vec![] + Ok(vec![]) } else { - // TODO this is a non-satisfiable constraint - should we panic? - vec![] + Err(EvalError::ConstraintUnsatisfiable(format!("{} != 0", self))) }; } 1 => { let (var, coeff) = self.coefficients.iter().next().unwrap(); assert!(!coeff.is_known_zero()); - let value = &self.offset / &coeff.neg(); - return vec![Effect::Assignment(var.clone(), value)]; + let value = self.offset.field_div(&coeff.neg()); + let assignment = Effect::Assignment(var.clone(), value); + if let Some(coeff) = coeff.try_to_number() { + assert!(!coeff.is_zero(), "Zero coefficient has not been removed."); + return Ok(vec![assignment]); + } else { + return Ok(vec![ + Assertion::assert_is_nonzero(coeff.clone()), + assignment, + ]); + } } _ => {} } let r = self.solve_through_constraints(range_constraints); if !r.is_empty() { - return r; + return Ok(r); } let negated = -self; let r = negated.solve_through_constraints(range_constraints); if !r.is_empty() { - return r; + return Ok(r); } - self.transfer_constraints(range_constraints) + Ok(self + .transfer_constraints(range_constraints) .into_iter() .chain(self.transfer_constraints(range_constraints)) - .collect() + .collect()) } /// Tries to solve a bit-decomposition equation. @@ -172,7 +218,7 @@ impl AffineSymbolicExpression { // Check if they are mutually exclusive and compute assignments. let mut covered_bits: ::Integer = 0.into(); - let mut assignments = vec![]; + let mut effects = vec![]; for (var, coeff, constraint) in constrained_coefficients { let mask = *constraint.multiple(coeff).mask(); if !(mask & covered_bits).is_zero() { @@ -182,7 +228,7 @@ impl AffineSymbolicExpression { covered_bits |= mask; } let masked = &(-&self.offset) & &T::from(mask).into(); - assignments.push(Effect::Assignment( + effects.push(Effect::Assignment( var.clone(), masked.integer_div(&coeff.into()), )); @@ -192,10 +238,16 @@ impl AffineSymbolicExpression { return vec![]; } - // TODO we need to add an assertion that offset is covered by the masks. - // Otherwise the equation is not solvable. - // TODO is this really the case? - assignments + // We need to assert that the masks cover the offset, + // otherwise the equation is not solvable. + // We assert offset & !masks == 0 <=> offset == offset | masks. + // We use the latter since we cannot properly bit-negate inside the field. + effects.push(Assertion::assert_eq( + self.offset.clone(), + &self.offset | &T::from(covered_bits).into(), + )); + + effects } fn transfer_constraints( @@ -331,6 +383,30 @@ mod test { a.try_mul(b).unwrap() } + #[test] + fn unsolveable() { + let r = from_number(10).solve(&SimpleRangeConstraintSet::default()); + assert!(r.is_err()); + } + + #[test] + fn unsolvable_with_vars() { + let x = Ase::from_known_variable("X"); + let y = Ase::from_known_variable("Y"); + let constr = &(&x + &y) - &from_number(10); + let r = constr.solve(&SimpleRangeConstraintSet::default()); + assert!(r.is_err()); + } + + #[test] + fn solveable_without_vars() { + let constr = &from_number(0); + assert!(constr + .solve(&SimpleRangeConstraintSet::default()) + .unwrap() + .is_empty()); + } + #[test] fn solve_simple_eq() { let y = Ase::from_known_variable("y"); @@ -340,7 +416,7 @@ mod test { let seven = from_number(7); let ten = from_number(10); let constr = &(&mul(&two, &x) + &mul(&seven, &y)) - &ten; - let effects = constr.solve(&SimpleRangeConstraintSet::default()); + let effects = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); assert_eq!(effects.len(), 1); let Effect::Assignment(var, expr) = &effects[0] else { panic!("Expected assignment"); @@ -370,13 +446,25 @@ mod test { // Without range constraints, this is not solvable. assert!(constr .solve(&SimpleRangeConstraintSet::default()) + .unwrap() .is_empty()); // With range constraints, it should be solvable. let effects = constr .solve(&range_constraints) + .unwrap() .into_iter() .map(|effect| match effect { Effect::Assignment(v, expr) => format!("{v} = {expr};\n"), + Effect::Assertion(Assertion { + lhs, + rhs, + expected_equal, + }) => { + format!( + "assert {lhs} {} {rhs};\n", + if expected_equal { "==" } else { "!=" } + ) + } _ => panic!(), }) .format("") @@ -386,6 +474,7 @@ mod test { "a = ((-((10 + Z)) & 65280) // 256); b = ((-((10 + Z)) & 16711680) // 65536); c = ((-((10 + Z)) & 4278190080) // 16777216); +assert (10 + Z) == ((10 + Z) | 4294967040); " ); } @@ -410,6 +499,7 @@ c = ((-((10 + Z)) & 4278190080) // 16777216); - &z; let effects = constr .solve(&range_constraints) + .unwrap() .into_iter() .map(|effect| match effect { Effect::RangeConstraint(v, rc) => format!("{v}: {rc};\n"), diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 83601df668..aed0ff0bd0 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -1,6 +1,6 @@ use std::{ fmt::{self, Display, Formatter}, - ops::{Add, BitAnd, Div, Mul, Neg}, + ops::{Add, BitAnd, BitOr, Mul, Neg}, }; use powdr_number::FieldElement; @@ -39,6 +39,7 @@ pub enum BinaryOperator { /// Integer division, i.e. convert field elements to unsigned integer and divide. IntegerDiv, BitAnd, + BitOr, } #[derive(Debug, Clone)] @@ -106,6 +107,7 @@ impl Display for BinaryOperator { BinaryOperator::Div => write!(f, "/"), BinaryOperator::IntegerDiv => write!(f, "//"), BinaryOperator::BitAnd => write!(f, "&"), + BinaryOperator::BitOr => write!(f, "|"), } } } @@ -199,16 +201,14 @@ impl Mul for &SymbolicExpression { } } -impl Div for &SymbolicExpression { - type Output = SymbolicExpression; - +impl SymbolicExpression { /// Field element division. See `integer_div` for integer division. - fn div(self, rhs: Self) -> Self::Output { + /// If you use this, you must ensure that the divisor is not zero. + pub fn field_div(&self, rhs: &Self) -> Self { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { assert!(b != &T::from(0)); SymbolicExpression::Concrete(*a / *b) } else if self.is_known_zero() { - // TODO should we still detect division by zero in this case? SymbolicExpression::Concrete(T::from(0)) } else if rhs.is_known_one() { self.clone() @@ -224,9 +224,8 @@ impl Div for &SymbolicExpression { ) } } -} -impl SymbolicExpression { + /// Integer division, i.e. convert field elements to unsigned integer and divide. pub fn integer_div(&self, rhs: &Self) -> Self { if rhs.is_known_one() { self.clone() @@ -247,8 +246,9 @@ impl BitAnd for &SymbolicExpression { fn bitand(self, rhs: Self) -> Self::Output { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { SymbolicExpression::Concrete(T::from(a.to_integer() & b.to_integer())) + } else if self.is_known_zero() || rhs.is_known_zero() { + SymbolicExpression::Concrete(T::from(0)) } else { - // TODO simplifications? SymbolicExpression::BinaryOperation( Box::new(self.clone()), BinaryOperator::BitAnd, @@ -260,3 +260,22 @@ impl BitAnd for &SymbolicExpression { } } } + +impl BitOr for &SymbolicExpression { + type Output = SymbolicExpression; + + fn bitor(self, rhs: Self) -> Self::Output { + if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { + let v = a.to_integer() | b.to_integer(); + assert!(v <= T::modulus()); + SymbolicExpression::Concrete(T::from(v)) + } else { + SymbolicExpression::BinaryOperation( + Box::new(self.clone()), + BinaryOperator::BitOr, + Box::new(rhs.clone()), + None, + ) + } + } +} From 09773df396cd682c1b03ac1042b0c6f17a00ceb0 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 09:03:44 +0000 Subject: [PATCH 05/44] clippy --- executor/src/witgen/jit/affine_symbolic_expression.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 1019642c02..e3101f3b18 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -157,7 +157,7 @@ impl AffineSymbolicExpression { return if self.offset.is_known_zero() { Ok(vec![]) } else { - Err(EvalError::ConstraintUnsatisfiable(format!("{} != 0", self))) + Err(EvalError::ConstraintUnsatisfiable(format!("{self} != 0"))) }; } 1 => { From cb5638d6695ebfeec5d437cc51d0501f60444112 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 09:04:33 +0000 Subject: [PATCH 06/44] Only mention constraint in error message. --- executor/src/witgen/jit/affine_symbolic_expression.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index e3101f3b18..b9081c7007 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -157,7 +157,7 @@ impl AffineSymbolicExpression { return if self.offset.is_known_zero() { Ok(vec![]) } else { - Err(EvalError::ConstraintUnsatisfiable(format!("{self} != 0"))) + Err(EvalError::ConstraintUnsatisfiable(self.to_string())) }; } 1 => { From e348031113406c6252ee4ced84ab9677db8deb13 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 13:06:26 +0000 Subject: [PATCH 07/44] witgen inference. --- executor/src/witgen/jit/cell.rs | 59 ++++ executor/src/witgen/jit/mod.rs | 4 +- executor/src/witgen/jit/witgen_inference.rs | 289 ++++++++++++++++++++ executor/src/witgen/range_constraints.rs | 8 + 4 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 executor/src/witgen/jit/cell.rs create mode 100644 executor/src/witgen/jit/witgen_inference.rs diff --git a/executor/src/witgen/jit/cell.rs b/executor/src/witgen/jit/cell.rs new file mode 100644 index 0000000000..090e277048 --- /dev/null +++ b/executor/src/witgen/jit/cell.rs @@ -0,0 +1,59 @@ +use std::{ + fmt::{self, Display, Formatter}, + hash::{Hash, Hasher}, +}; + +use powdr_ast::analyzed::AlgebraicReference; + +/// The identifier of a witness cell in the trace table. +/// The `row_offset` is relative to a certain "zero row" defined +/// by the component that uses this data structure. +#[derive(Debug, Clone, Eq)] +pub struct Cell { + /// Name of the column, used only for display purposes. + pub column_name: String, + pub id: u64, + pub row_offset: i32, +} + +impl Hash for Cell { + fn hash(&self, state: &mut H) { + self.id.hash(state); + self.row_offset.hash(state); + } +} + +impl PartialEq for Cell { + fn eq(&self, other: &Self) -> bool { + self.id == other.id && self.row_offset == other.row_offset + } +} + +impl Ord for Cell { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + (self.id, self.row_offset).cmp(&(other.id, other.row_offset)) + } +} + +impl PartialOrd for Cell { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Cell { + pub fn from_reference(r: &AlgebraicReference, row_offset: i32) -> Self { + assert!(r.is_witness()); + Self { + column_name: r.name.clone(), + id: r.poly_id.id, + row_offset: r.next as i32 + row_offset, + } + } +} + +impl Display for Cell { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "{}[{}]", self.column_name, self.row_offset) + } +} diff --git a/executor/src/witgen/jit/mod.rs b/executor/src/witgen/jit/mod.rs index a6d9dabf39..c3151f1638 100644 --- a/executor/src/witgen/jit/mod.rs +++ b/executor/src/witgen/jit/mod.rs @@ -1,2 +1,4 @@ -mod affine_symbolic_expression; +pub(crate) mod affine_symbolic_expression; +mod cell; mod symbolic_expression; +mod witgen_inference; diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs new file mode 100644 index 0000000000..c395654dd4 --- /dev/null +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -0,0 +1,289 @@ +use std::{ + collections::{BTreeSet, HashMap, HashSet}, + iter::once, +}; + +use itertools::Itertools; +use powdr_ast::{ + analyzed::{ + AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, + AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, + LookupIdentity, PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, + SelectedExpressions, + }, + indent, + parsed::visitor::AllChildren, +}; +use powdr_number::FieldElement; + +use crate::witgen::global_constraints::RangeConstraintSet; + +use super::{ + super::{machines::MachineParts, range_constraints::RangeConstraint, FixedData}, + affine_symbolic_expression::{AffineSymbolicExpression, Effect}, + cell::Cell, + symbolic_expression::SymbolicExpression, +}; + +/// This component can generate code that solves identities. +/// It needs a driver that tells it which identities to process on which rows. +pub struct WitgenInference<'a, T: FieldElement> { + fixed_data: &'a FixedData<'a, T>, + range_constraints: HashMap>, + known_cells: HashSet, + code: Vec, // TODO make this a proper expression +} + +impl<'a, T: FieldElement> WitgenInference<'a, T> { + pub fn new( + fixed_data: &'a FixedData<'a, T>, + known_cells: impl IntoIterator, + ) -> Self { + Self { + fixed_data, + range_constraints: Default::default(), + known_cells: known_cells.into_iter().collect(), + code: Default::default(), + } + } + + fn cell_at_row(&self, id: u64, row_offset: i32) -> Cell { + let poly_id = PolyID { + id, + ptype: PolynomialType::Committed, + }; + Cell { + column_name: self.fixed_data.column_name(&poly_id).to_string(), + id, + row_offset, + } + } + + fn known_cells(&self) -> &HashSet { + &self.known_cells + } + + pub fn process_identity(&mut self, id: &Identity, row_offset: i32) { + let effects = match id { + Identity::Polynomial(PolynomialIdentity { expression, .. }) => { + self.process_polynomial_identity(expression, row_offset) + } + Identity::Lookup(LookupIdentity { + id, + source: _, + left, + right, + }) + | Identity::PhantomLookup(PhantomLookupIdentity { + id, + source: _, + left, + right, + multiplicity: _, + }) => { + // TODO multiplicity? + self.process_lookup(*id, left, right, row_offset) + } + _ => { + // TODO + vec![] + } + }; + self.ingest_effects(effects); + } + + fn process_polynomial_identity( + &self, + expression: &'a Expression, + offset: i32, + ) -> Vec> { + if let Some(r) = self.evaluate(expression, offset) { + // TODO remove unwrap + r.solve(self).unwrap() + } else { + vec![] + } + } + + fn process_lookup( + &self, + lookup_id: u64, + left: &SelectedExpressions, + right: &SelectedExpressions, + offset: i32, + ) -> Vec> { + // TODO: In the future, call the 'mutable state' to check if the + // lookup can always be answered. + + // If the RHS is fully fixed columns... + if right.expressions.iter().all(|e| match e { + Expression::Reference(r) => r.is_fixed(), + Expression::Number(_) => true, + _ => false, + }) { + // and the selector is known to be 1 and all except one expression is known on the LHS. + if self + .evaluate(&left.selector, offset) + .map(|s| s.is_known_one()) + == Some(true) + { + if let Some(inputs) = left + .expressions + .iter() + .map(|e| self.evaluate(e, offset)) + .collect::>>() + { + if inputs.iter().filter(|i| i.is_known()).count() == inputs.len() - 1 { + let mut var_decl = String::new(); + let mut output_var = String::new(); + let query = inputs + .iter() + .enumerate() + .map(|(i, e)| { + if e.is_known() { + format!("LookupCell::Input(&({e}))") + } else { + let var_name = format!("lookup_{lookup_id}_{i}"); + output_var = var_name.clone(); + var_decl.push_str(&format!( + "let mut {var_name}: FieldElement = Default::default();" + )); + format!("LookupCell::Output(&mut {var_name})") + } + }) + .format(", "); + let machine_call = format!( + "assert!(process_lookup(mutable_state, {lookup_id}, &mut [{query}]));" + ); + // TODO range constraints? + let output_expr = inputs.iter().find(|i| !i.is_known()).unwrap(); + return once(Effect::Code(var_decl)) + .chain(once(Effect::Code(machine_call))) + .chain( + (output_expr + - &KnownValue::from_known_local_var(&output_var).into()) + .solve(self), + ) + .collect(); + } + } + } + } + vec![] + } + + fn ingest_effects(&mut self, effects: Vec>) { + for e in effects { + match e { + Effect::Assignment(cell, assignment) => { + // TODO also use raneg constraint? + self.known_cells.insert(cell.clone()); + self.code.push(assignment); + } + Effect::RangeConstraint(cell, rc) => { + self.add_range_constraint(cell, rc); + } + Effect::Code(code) => { + self.code.push(code); + } + } + } + } + + fn add_range_constraint(&mut self, cell: Cell, rc: RangeConstraint) { + let rc = self + .range_constraint(cell.clone()) + .map_or(rc.clone(), |existing_rc| existing_rc.conjunction(&rc)); + // TODO if the conjuntion results in a single value, make the cell known. + self.range_constraints.insert(cell, rc); + } + + fn evaluate( + &self, + expr: &Expression, + offset: i32, + ) -> Option> { + Some(match expr { + Expression::Reference(r) => { + if r.is_fixed() { + todo!() + // let mut row = self.latch_row as i64 + offset as i64; + // while row < 0 { + // row += self.block_size as i64; + // } + // // TODO at some point we should check that all of the fixed columns are periodic. + // // TODO We can only do this for block machines. + // // For dynamic machines, fixed columns are "known but symbolic" + // let v = self.fixed_data.fixed_cols[&r.poly_id].values_max_size()[row as usize]; + // EvalResult::from_number(v) + } else { + let cell = Cell::from_reference(r, offset); + // If a cell is known and has a compile-time constant value, + // that value is stored in the range constraints. + if let Some(v) = self + .range_constraint(cell.clone()) + .and_then(|rc| rc.try_to_single_value()) + { + AffineSymbolicExpression::from_number(v) + } else if self.known_cells.contains(&cell) { + AffineSymbolicExpression::from_known_variable(cell) + } else { + AffineSymbolicExpression::from_unknown_variable(cell) + } + } + } + Expression::PublicReference(_) => return None, // TODO + Expression::Challenge(_) => return None, // TODO + Expression::Number(n) => AffineSymbolicExpression::from_number(*n), + Expression::BinaryOperation(op) => self.evaulate_binary_operation(op, offset)?, + Expression::UnaryOperation(op) => self.evaluate_unary_operation(op, offset)?, + }) + } + + fn evaulate_binary_operation( + &self, + op: &AlgebraicBinaryOperation, + offset: i32, + ) -> Option> { + let left = self.evaluate(&op.left, offset)?; + let right = self.evaluate(&op.right, offset)?; + match op.op { + AlgebraicBinaryOperator::Add => Some(&left + &right), + AlgebraicBinaryOperator::Sub => Some(&left - &right), + AlgebraicBinaryOperator::Mul => left.try_mul(&right), + _ => todo!(), + } + } + + fn evaluate_unary_operation( + &self, + op: &AlgebraicUnaryOperation, + offset: i32, + ) -> Option> { + let expr = self.evaluate(&op.expr, offset)?; + match op.op { + AlgebraicUnaryOperator::Minus => Some(-&expr), + } + } +} + +impl RangeConstraintSet for WitgenInference<'_, T> { + // TODO would be nice to use &Cell, but this leads to lifetime trouble + // in the solve() function. + fn range_constraint(&self, cell: Cell) -> Option> { + self.fixed_data + .global_range_constraints + .range_constraint(&AlgebraicReference { + name: Default::default(), + poly_id: PolyID { + id: cell.id, + ptype: PolynomialType::Committed, + }, + next: false, + }) + .iter() + .chain(self.range_constraints.get(&cell)) + .cloned() + .reduce(|gc, rc| gc.conjunction(&rc)) + } +} diff --git a/executor/src/witgen/range_constraints.rs b/executor/src/witgen/range_constraints.rs index 8e6d9748a8..15b99e9eca 100644 --- a/executor/src/witgen/range_constraints.rs +++ b/executor/src/witgen/range_constraints.rs @@ -157,6 +157,14 @@ impl RangeConstraint { mask: mask.unwrap_or_else(|| Self::from_range(min, max).mask), } } + + pub fn try_to_single_value(&self) -> Option { + if self.min == self.max { + Some(self.min) + } else { + None + } + } } /// The number of elements in an (inclusive) min/max range. From fe7908eb2acfe689f8b1e30ff16af5d52db1ee09 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 15:31:50 +0000 Subject: [PATCH 08/44] Generic vars for symbolic expression. --- .../witgen/jit/affine_symbolic_expression.rs | 37 +++++++------- .../src/witgen/jit/symbolic_expression.rs | 48 +++++++++---------- 2 files changed, 43 insertions(+), 42 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index b9081c7007..a6eb191cf4 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -15,37 +15,40 @@ use super::{super::range_constraints::RangeConstraint, symbolic_expression::Symb /// The effect of solving a symbolic equation. pub enum Effect { /// variable can be assigned a value. - Assignment(V, SymbolicExpression), + Assignment(V, SymbolicExpression), /// we learnt a new range constraint on variable. RangeConstraint(V, RangeConstraint), /// a run-time assertion. If this fails, we have conflicting constraints. - Assertion(Assertion), + Assertion(Assertion), } /// A run-time assertion. If this fails, we have conflicting constraints. -pub struct Assertion { - pub lhs: SymbolicExpression, - pub rhs: SymbolicExpression, +pub struct Assertion { + pub lhs: SymbolicExpression, + pub rhs: SymbolicExpression, /// If this is true, we assert that both sides are equal. /// Otherwise, we assert that they are different. pub expected_equal: bool, } -impl Assertion { - pub fn assert_is_zero(condition: SymbolicExpression) -> Effect { +impl Assertion { + pub fn assert_is_zero(condition: SymbolicExpression) -> Effect { Self::assert_eq(condition, SymbolicExpression::from(T::from(0))) } - pub fn assert_is_nonzero(condition: SymbolicExpression) -> Effect { + pub fn assert_is_nonzero(condition: SymbolicExpression) -> Effect { Self::assert_neq(condition, SymbolicExpression::from(T::from(0))) } - pub fn assert_eq(lhs: SymbolicExpression, rhs: SymbolicExpression) -> Effect { + pub fn assert_eq(lhs: SymbolicExpression, rhs: SymbolicExpression) -> Effect { Effect::Assertion(Assertion { lhs, rhs, expected_equal: true, }) } - pub fn assert_neq(lhs: SymbolicExpression, rhs: SymbolicExpression) -> Effect { + pub fn assert_neq( + lhs: SymbolicExpression, + rhs: SymbolicExpression, + ) -> Effect { Effect::Assertion(Assertion { lhs, rhs, @@ -59,8 +62,8 @@ impl Assertion { /// and the `x_i` are unknown variables. #[derive(Debug, Clone)] pub struct AffineSymbolicExpression { - coefficients: BTreeMap>, - offset: SymbolicExpression, + coefficients: BTreeMap>, + offset: SymbolicExpression, } impl Display for AffineSymbolicExpression { @@ -90,8 +93,8 @@ impl Display for AffineSymbolicExpression { } } -impl From> for AffineSymbolicExpression { - fn from(k: SymbolicExpression) -> Self { +impl From> for AffineSymbolicExpression { + fn from(k: SymbolicExpression) -> Self { AffineSymbolicExpression { coefficients: Default::default(), offset: k, @@ -100,7 +103,7 @@ impl From> for AffineSymbolicExpressio } impl AffineSymbolicExpression { - pub fn from_known_variable(var: &str) -> Self { + pub fn from_known_variable(var: V) -> Self { SymbolicExpression::from_var(var).into() } pub fn from_unknown_variable(var: V) -> Self { @@ -347,7 +350,9 @@ impl Neg for &AffineSymbolicExpression { } } -impl TryFrom<&AffineSymbolicExpression> for SymbolicExpression { +impl TryFrom<&AffineSymbolicExpression> + for SymbolicExpression +{ type Error = (); fn try_from(value: &AffineSymbolicExpression) -> Result { diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index aed0ff0bd0..defc54c127 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -11,22 +11,18 @@ use crate::witgen::range_constraints::RangeConstraint; /// involving known cells or variables and compile-time constants. /// Each of the sub-expressions can have its own range constraint. #[derive(Debug, Clone)] -pub enum SymbolicExpression { +pub enum SymbolicExpression { /// A concrete constant value known at compile time. Concrete(T), /// A symbolic value known at run-time, referencing either a cell or a local variable. - Variable(String, Option>), + Variable(V, Option>), BinaryOperation( - Box>, + Box, BinaryOperator, - Box>, - Option>, - ), - UnaryOperation( - UnaryOperator, - Box>, + Box, Option>, ), + UnaryOperation(UnaryOperator, Box, Option>), } #[derive(Debug, Clone)] @@ -47,7 +43,7 @@ pub enum UnaryOperator { Neg, } -impl SymbolicExpression { +impl SymbolicExpression { pub fn is_known_zero(&self) -> bool { self.try_to_number().map_or(false, |n| n.is_zero()) } @@ -79,7 +75,7 @@ impl SymbolicExpression { } } -impl Display for SymbolicExpression { +impl Display for SymbolicExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { SymbolicExpression::Concrete(n) => { @@ -120,20 +116,20 @@ impl Display for UnaryOperator { } } -impl SymbolicExpression { - pub fn from_var(name: &str) -> Self { - SymbolicExpression::Variable(name.to_string(), None) +impl SymbolicExpression { + pub fn from_var(name: V) -> Self { + SymbolicExpression::Variable(name, None) } } -impl From for SymbolicExpression { +impl From for SymbolicExpression { fn from(n: T) -> Self { SymbolicExpression::Concrete(n) } } -impl Add for &SymbolicExpression { - type Output = SymbolicExpression; +impl Add for &SymbolicExpression { + type Output = SymbolicExpression; fn add(self, rhs: Self) -> Self::Output { if self.is_known_zero() { @@ -158,8 +154,8 @@ impl Add for &SymbolicExpression { } } -impl Neg for &SymbolicExpression { - type Output = SymbolicExpression; +impl Neg for &SymbolicExpression { + type Output = SymbolicExpression; fn neg(self) -> Self::Output { match self { @@ -174,8 +170,8 @@ impl Neg for &SymbolicExpression { } } -impl Mul for &SymbolicExpression { - type Output = SymbolicExpression; +impl Mul for &SymbolicExpression { + type Output = SymbolicExpression; fn mul(self, rhs: Self) -> Self::Output { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { @@ -201,7 +197,7 @@ impl Mul for &SymbolicExpression { } } -impl SymbolicExpression { +impl SymbolicExpression { /// Field element division. See `integer_div` for integer division. /// If you use this, you must ensure that the divisor is not zero. pub fn field_div(&self, rhs: &Self) -> Self { @@ -240,8 +236,8 @@ impl SymbolicExpression { } } -impl BitAnd for &SymbolicExpression { - type Output = SymbolicExpression; +impl BitAnd for &SymbolicExpression { + type Output = SymbolicExpression; fn bitand(self, rhs: Self) -> Self::Output { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { @@ -261,8 +257,8 @@ impl BitAnd for &SymbolicExpression { } } -impl BitOr for &SymbolicExpression { - type Output = SymbolicExpression; +impl BitOr for &SymbolicExpression { + type Output = SymbolicExpression; fn bitor(self, rhs: Self) -> Self::Output { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { From 5014950ae8fe812a160581718a7c0b04c52c495b Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 16:07:32 +0000 Subject: [PATCH 09/44] Also implement operations for non-ref. --- .../witgen/jit/affine_symbolic_expression.rs | 52 ++++++++++++++----- .../src/witgen/jit/symbolic_expression.rs | 37 +++++++++++++ 2 files changed, 76 insertions(+), 13 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index a6eb191cf4..8bc114d3c4 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -230,7 +230,7 @@ impl AffineSymbolicExpression { } else { covered_bits |= mask; } - let masked = &(-&self.offset) & &T::from(mask).into(); + let masked = -&self.offset & T::from(mask).into(); effects.push(Effect::Assignment( var.clone(), masked.integer_div(&coeff.into()), @@ -327,6 +327,14 @@ impl Add for &AffineSymbolicExpression { } } +impl Add for AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn add(self, rhs: Self) -> Self::Output { + &self + &rhs + } +} + impl Sub for &AffineSymbolicExpression { type Output = AffineSymbolicExpression; @@ -335,6 +343,14 @@ impl Sub for &AffineSymbolicExpression { } } +impl Sub for AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn sub(self, rhs: Self) -> Self::Output { + &self - &rhs + } +} + impl Neg for &AffineSymbolicExpression { type Output = AffineSymbolicExpression; @@ -350,6 +366,14 @@ impl Neg for &AffineSymbolicExpression { } } +impl Neg for AffineSymbolicExpression { + type Output = AffineSymbolicExpression; + + fn neg(self) -> Self::Output { + -&self + } +} + impl TryFrom<&AffineSymbolicExpression> for SymbolicExpression { @@ -396,9 +420,9 @@ mod test { #[test] fn unsolvable_with_vars() { - let x = Ase::from_known_variable("X"); - let y = Ase::from_known_variable("Y"); - let constr = &(&x + &y) - &from_number(10); + let x = &Ase::from_known_variable("X"); + let y = &Ase::from_known_variable("Y"); + let constr = x + y - from_number(10); let r = constr.solve(&SimpleRangeConstraintSet::default()); assert!(r.is_err()); } @@ -420,7 +444,7 @@ mod test { let two = from_number(2); let seven = from_number(7); let ten = from_number(10); - let constr = &(&mul(&two, &x) + &mul(&seven, &y)) - &ten; + let constr = mul(&two, &x) + mul(&seven, &y) - ten; let effects = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); assert_eq!(effects.len(), 1); let Effect::Assignment(var, expr) = &effects[0] else { @@ -444,10 +468,11 @@ mod test { ); // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 + Z = 0 let ten = from_number(10); - let constr = &(&(&(&mul(&a, &from_number(0x100)) + &mul(&b, &from_number(0x10000))) - + &mul(&c, &from_number(0x1000000))) - + &ten) - + &z; + let constr = mul(&a, &from_number(0x100)) + + mul(&b, &from_number(0x10000)) + + mul(&c, &from_number(0x1000000)) + + ten + + z; // Without range constraints, this is not solvable. assert!(constr .solve(&SimpleRangeConstraintSet::default()) @@ -498,10 +523,11 @@ assert (10 + Z) == ((10 + Z) | 4294967040); ); // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 - Z = 0 let ten = from_number(10); - let constr = &(&(&(&mul(&a, &from_number(0x100)) + &mul(&b, &from_number(0x10000))) - + &mul(&c, &from_number(0x1000000))) - + &ten) - - &z; + let constr = mul(&a, &from_number(0x100)) + + mul(&b, &from_number(0x10000)) + + mul(&c, &from_number(0x1000000)) + + ten + - z; let effects = constr .solve(&range_constraints) .unwrap() diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index defc54c127..780f54793c 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -154,6 +154,13 @@ impl Add for &SymbolicExpression { } } +impl Add for SymbolicExpression { + type Output = SymbolicExpression; + fn add(self, rhs: Self) -> Self::Output { + &self + &rhs + } +} + impl Neg for &SymbolicExpression { type Output = SymbolicExpression; @@ -170,6 +177,13 @@ impl Neg for &SymbolicExpression { } } +impl Neg for SymbolicExpression { + type Output = SymbolicExpression; + fn neg(self) -> Self::Output { + -&self + } +} + impl Mul for &SymbolicExpression { type Output = SymbolicExpression; @@ -197,6 +211,13 @@ impl Mul for &SymbolicExpression { } } +impl Mul for SymbolicExpression { + type Output = SymbolicExpression; + fn mul(self, rhs: Self) -> Self { + &self * &rhs + } +} + impl SymbolicExpression { /// Field element division. See `integer_div` for integer division. /// If you use this, you must ensure that the divisor is not zero. @@ -257,6 +278,14 @@ impl BitAnd for &SymbolicExpression { } } +impl BitAnd for SymbolicExpression { + type Output = SymbolicExpression; + + fn bitand(self, rhs: Self) -> Self::Output { + &self & &rhs + } +} + impl BitOr for &SymbolicExpression { type Output = SymbolicExpression; @@ -275,3 +304,11 @@ impl BitOr for &SymbolicExpression { } } } + +impl BitOr for SymbolicExpression { + type Output = SymbolicExpression; + + fn bitor(self, rhs: Self) -> Self::Output { + &self | &rhs + } +} From e3e1fdb5035e6c43248799c6914a795f371460e9 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 16:48:08 +0000 Subject: [PATCH 10/44] witgen inference. --- .../witgen/jit/affine_symbolic_expression.rs | 29 +++- executor/src/witgen/jit/witgen_inference.rs | 125 ++++++++---------- 2 files changed, 80 insertions(+), 74 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 8bc114d3c4..cdaf4d68ef 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -20,6 +20,8 @@ pub enum Effect { RangeConstraint(V, RangeConstraint), /// a run-time assertion. If this fails, we have conflicting constraints. Assertion(Assertion), + /// a lookup / call to a different machine. + Lookup(u64, Vec>), } /// A run-time assertion. If this fails, we have conflicting constraints. @@ -57,6 +59,11 @@ impl Assertion { } } +pub enum LookupArgument { + Known(SymbolicExpression), + Unknown(AffineSymbolicExpression), +} + /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, /// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time, /// and the `x_i` are unknown variables. @@ -116,13 +123,20 @@ impl AffineSymbolicExpression { SymbolicExpression::from(n).into() } pub fn is_known_zero(&self) -> bool { - self.coefficients.is_empty() && self.offset.is_known_zero() + self.try_to_known().map_or(false, |k| k.is_known_zero()) } pub fn is_known_one(&self) -> bool { - self.coefficients.is_empty() && self.offset.is_known_one() + self.try_to_known().map_or(false, |k| k.is_known_one()) } pub fn is_known(&self) -> bool { - self.coefficients.is_empty() + self.try_to_known().is_some() + } + pub fn try_to_known(&self) -> Option<&SymbolicExpression> { + if self.coefficients.is_empty() { + Some(&self.offset) + } else { + None + } } /// Tries to multiply this expression with another one. @@ -150,6 +164,15 @@ impl AffineSymbolicExpression { }) } + /// If this expression contains a single unknown variable, returns it. + pub fn single_unknown_variable(&self) -> Option<&V> { + if self.coefficients.len() == 1 { + self.coefficients.keys().next() + } else { + None + } + } + /// Solves the equation `self = 0` and returns how to compute the solution. pub fn solve( &self, diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index c395654dd4..c944b55958 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -1,28 +1,21 @@ -use std::{ - collections::{BTreeSet, HashMap, HashSet}, - iter::once, -}; +use std::collections::{HashMap, HashSet}; use itertools::Itertools; -use powdr_ast::{ - analyzed::{ - AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, - AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, - LookupIdentity, PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, - SelectedExpressions, - }, - indent, - parsed::visitor::AllChildren, +use powdr_ast::analyzed::{ + AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, + AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, LookupIdentity, + PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, SelectedExpressions, }; use powdr_number::FieldElement; -use crate::witgen::global_constraints::RangeConstraintSet; +use crate::witgen::{ + global_constraints::RangeConstraintSet, jit::affine_symbolic_expression::LookupArgument, +}; use super::{ - super::{machines::MachineParts, range_constraints::RangeConstraint, FixedData}, + super::{range_constraints::RangeConstraint, FixedData}, affine_symbolic_expression::{AffineSymbolicExpression, Effect}, cell::Cell, - symbolic_expression::SymbolicExpression, }; /// This component can generate code that solves identities. @@ -31,7 +24,7 @@ pub struct WitgenInference<'a, T: FieldElement> { fixed_data: &'a FixedData<'a, T>, range_constraints: HashMap>, known_cells: HashSet, - code: Vec, // TODO make this a proper expression + code: Vec>, } impl<'a, T: FieldElement> WitgenInference<'a, T> { @@ -47,20 +40,12 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { } } - fn cell_at_row(&self, id: u64, row_offset: i32) -> Cell { - let poly_id = PolyID { - id, - ptype: PolynomialType::Committed, - }; - Cell { - column_name: self.fixed_data.column_name(&poly_id).to_string(), - id, - row_offset, - } + pub fn known_cells(&self) -> &HashSet { + &self.known_cells } - fn known_cells(&self) -> &HashSet { - &self.known_cells + pub fn code(self) -> Vec> { + self.code } pub fn process_identity(&mut self, id: &Identity, row_offset: i32) { @@ -98,7 +83,10 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { offset: i32, ) -> Vec> { if let Some(r) = self.evaluate(expression, offset) { - // TODO remove unwrap + // 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(self).unwrap() } else { vec![] @@ -121,50 +109,33 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { Expression::Number(_) => true, _ => false, }) { - // and the selector is known to be 1 and all except one expression is known on the LHS. + // and the selector is known to be 1... if self .evaluate(&left.selector, offset) .map(|s| s.is_known_one()) == Some(true) { - if let Some(inputs) = left + if let Some(lhs) = left .expressions .iter() .map(|e| self.evaluate(e, offset)) .collect::>>() { - if inputs.iter().filter(|i| i.is_known()).count() == inputs.len() - 1 { - let mut var_decl = String::new(); - let mut output_var = String::new(); - let query = inputs - .iter() - .enumerate() - .map(|(i, e)| { - if e.is_known() { - format!("LookupCell::Input(&({e}))") - } else { - let var_name = format!("lookup_{lookup_id}_{i}"); - output_var = var_name.clone(); - var_decl.push_str(&format!( - "let mut {var_name}: FieldElement = Default::default();" - )); - format!("LookupCell::Output(&mut {var_name})") - } - }) - .format(", "); - let machine_call = format!( - "assert!(process_lookup(mutable_state, {lookup_id}, &mut [{query}]));" - ); - // TODO range constraints? - let output_expr = inputs.iter().find(|i| !i.is_known()).unwrap(); - return once(Effect::Code(var_decl)) - .chain(once(Effect::Code(machine_call))) - .chain( - (output_expr - - &KnownValue::from_known_local_var(&output_var).into()) - .solve(self), - ) - .collect(); + // and all except one expression is known on the LHS. + let unknown = lhs.iter().filter(|e| !e.is_known()).collect_vec(); + if unknown.len() == 1 && unknown[0].single_unknown_variable().is_some() { + return vec![Effect::Lookup( + lookup_id, + lhs.into_iter() + .map(|e| { + if let Some(val) = e.try_to_known() { + LookupArgument::Known(val.clone()) + } else { + LookupArgument::Unknown(e) + } + }) + .collect(), + )]; } } } @@ -172,20 +143,31 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { vec![] } - fn ingest_effects(&mut self, effects: Vec>) { + fn ingest_effects(&mut self, effects: Vec>) { for e in effects { - match e { + match &e { Effect::Assignment(cell, assignment) => { - // TODO also use raneg constraint? + if let Some(rc) = assignment.range_constraint() { + // If the cell was determined to be a constant, we add this + // as a range constraint, so we can use it in future evaluations. + self.add_range_constraint(cell.clone(), rc); + } self.known_cells.insert(cell.clone()); - self.code.push(assignment); + self.code.push(e); } Effect::RangeConstraint(cell, rc) => { - self.add_range_constraint(cell, rc); + self.add_range_constraint(cell.clone(), rc.clone()); } - Effect::Code(code) => { - self.code.push(code); + Effect::Lookup(_, arguments) => { + for arg in arguments { + if let LookupArgument::Unknown(expr) = arg { + let cell = expr.single_unknown_variable().unwrap(); + self.known_cells.insert(cell.clone()); + } + } + self.code.push(e); } + Effect::Assertion(_) => self.code.push(e), } } } @@ -195,6 +177,7 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { .range_constraint(cell.clone()) .map_or(rc.clone(), |existing_rc| existing_rc.conjunction(&rc)); // TODO if the conjuntion results in a single value, make the cell known. + // TODO but do we also need to generate code for the assignment? self.range_constraints.insert(cell, rc); } From d21c061fe1188e1db0010157635d4937c28296c3 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 21:18:58 +0000 Subject: [PATCH 11/44] Move "from_var" fn. --- executor/src/witgen/jit/symbolic_expression.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 780f54793c..55a44fecf3 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -44,6 +44,10 @@ pub enum UnaryOperator { } impl SymbolicExpression { + pub fn from_var(name: V) -> Self { + SymbolicExpression::Variable(name, None) + } + pub fn is_known_zero(&self) -> bool { self.try_to_number().map_or(false, |n| n.is_zero()) } @@ -116,12 +120,6 @@ impl Display for UnaryOperator { } } -impl SymbolicExpression { - pub fn from_var(name: V) -> Self { - SymbolicExpression::Variable(name, None) - } -} - impl From for SymbolicExpression { fn from(n: T) -> Self { SymbolicExpression::Concrete(n) From 6bbaf54026cceb158807b4ed4967d7c3e0f9c329 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 21:27:14 +0000 Subject: [PATCH 12/44] Use rc. --- .../src/witgen/jit/symbolic_expression.rs | 37 ++++++++++--------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 55a44fecf3..4989ae5273 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -1,6 +1,7 @@ use std::{ fmt::{self, Display, Formatter}, ops::{Add, BitAnd, BitOr, Mul, Neg}, + rc::Rc, }; use powdr_number::FieldElement; @@ -17,12 +18,12 @@ pub enum SymbolicExpression { /// A symbolic value known at run-time, referencing either a cell or a local variable. Variable(V, Option>), BinaryOperation( - Box, + Rc, BinaryOperator, - Box, + Rc, Option>, ), - UnaryOperation(UnaryOperator, Box, Option>), + UnaryOperation(UnaryOperator, Rc, Option>), } #[derive(Debug, Clone)] @@ -141,9 +142,9 @@ impl Add for &SymbolicExpression { SymbolicExpression::Concrete(*a + *b) } _ => SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::Add, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), self.range_constraint() .zip(rhs.range_constraint()) .map(|(a, b)| a.combine_sum(&b)), @@ -165,10 +166,12 @@ impl Neg for &SymbolicExpression { fn neg(self) -> Self::Output { match self { SymbolicExpression::Concrete(n) => SymbolicExpression::Concrete(-*n), - SymbolicExpression::UnaryOperation(UnaryOperator::Neg, expr, _) => *expr.clone(), + SymbolicExpression::UnaryOperation(UnaryOperator::Neg, expr, _) => { + expr.as_ref().clone() + } _ => SymbolicExpression::UnaryOperation( UnaryOperator::Neg, - Box::new(self.clone()), + Rc::new(self.clone()), self.range_constraint().map(|rc| rc.multiple(-T::from(1))), ), } @@ -200,9 +203,9 @@ impl Mul for &SymbolicExpression { -self } else { SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::Mul, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), None, ) } @@ -232,9 +235,9 @@ impl SymbolicExpression { } else { // TODO other simplifications like `-x / -y => x / y`, `-x / concrete => x / -concrete`, etc. SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::Div, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), None, ) } @@ -246,9 +249,9 @@ impl SymbolicExpression { self.clone() } else { SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::IntegerDiv, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), None, ) } @@ -265,9 +268,9 @@ impl BitAnd for &SymbolicExpression { SymbolicExpression::Concrete(T::from(0)) } else { SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::BitAnd, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), self.range_constraint() .zip(rhs.range_constraint()) .map(|(a, b)| a.conjunction(&b)), @@ -294,9 +297,9 @@ impl BitOr for &SymbolicExpression { SymbolicExpression::Concrete(T::from(v)) } else { SymbolicExpression::BinaryOperation( - Box::new(self.clone()), + Rc::new(self.clone()), BinaryOperator::BitOr, - Box::new(rhs.clone()), + Rc::new(rhs.clone()), None, ) } From 1070ce9f35f241cf1422090f31eaf1a964e185c1 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 21:32:45 +0000 Subject: [PATCH 13/44] use retain. --- executor/src/witgen/jit/affine_symbolic_expression.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 8bc114d3c4..4df0d4eb28 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -309,16 +309,13 @@ impl Add for &AffineSymbolicExpression { fn add(self, rhs: Self) -> Self::Output { let mut coefficients = self.coefficients.clone(); - for (var, coeff) in rhs.coefficients.iter() { + for (var, coeff) in &rhs.coefficients { coefficients .entry(var.clone()) .and_modify(|f| *f = &*f + coeff) .or_insert_with(|| coeff.clone()); } - let coefficients = coefficients - .into_iter() - .filter(|(_, f)| !f.is_known_zero()) - .collect(); + coefficients.retain(|_, f| !f.is_known_zero()); let offset = &self.offset + &rhs.offset; AffineSymbolicExpression { coefficients, From ff6a2f74e1973af1ccbca65320fabd48f3f0a8a3 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 21:55:39 +0000 Subject: [PATCH 14/44] Fix bug and check for zero in "bitor" --- executor/src/witgen/jit/symbolic_expression.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 4989ae5273..b6e689b570 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -293,8 +293,12 @@ impl BitOr for &SymbolicExpression { fn bitor(self, rhs: Self) -> Self::Output { if let (SymbolicExpression::Concrete(a), SymbolicExpression::Concrete(b)) = (self, rhs) { let v = a.to_integer() | b.to_integer(); - assert!(v <= T::modulus()); + assert!(v < T::modulus()); SymbolicExpression::Concrete(T::from(v)) + } else if self.is_known_zero() { + rhs.clone() + } else if rhs.is_known_zero() { + self.clone() } else { SymbolicExpression::BinaryOperation( Rc::new(self.clone()), From e6843ba1ea3671a7c793b40aae4e6b690285e349 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 22:06:58 +0000 Subject: [PATCH 15/44] Simplify known. --- .../witgen/jit/affine_symbolic_expression.rs | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 4df0d4eb28..2dd7436757 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -115,22 +115,19 @@ impl AffineSymbolicExpression { pub fn from_number(n: T) -> Self { SymbolicExpression::from(n).into() } - pub fn is_known_zero(&self) -> bool { - self.coefficients.is_empty() && self.offset.is_known_zero() - } - pub fn is_known_one(&self) -> bool { - self.coefficients.is_empty() && self.offset.is_known_one() - } - pub fn is_known(&self) -> bool { - self.coefficients.is_empty() + + /// If this expression does not contain unknown variables, returns the symbolic expression. + pub fn try_to_known(&self) -> Option<&SymbolicExpression> { + if self.coefficients.is_empty() { + Some(&self.offset) + } else { + None + } } /// Tries to multiply this expression with another one. /// Returns `None` if the result would be quadratic. pub fn try_mul(&self, other: &Self) -> Option { - if self.is_known_zero() || other.is_known_zero() { - return Some(AffineSymbolicExpression::from_number(T::from(0))); - } if !self.coefficients.is_empty() && !other.coefficients.is_empty() { return None; } From eb62678bd00f6cc4083367ef86f32e1765b2312f Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 10 Dec 2024 22:25:29 +0000 Subject: [PATCH 16/44] Add test. --- executor/src/witgen/jit/witgen_inference.rs | 58 +++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 30a3d109bf..2458e6e516 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -273,3 +273,61 @@ impl RangeConstraintSet for WitgenInference<'_, T> { .reduce(|gc, rc| gc.conjunction(&rc)) } } + +#[cfg(test)] +mod test { + use powdr_ast::analyzed::Analyzed; + use powdr_number::GoldilocksField; + + use crate::witgen::{jit::affine_symbolic_expression::Assertion, FixedData}; + + use super::*; + + fn format_code(effects: &[Effect]) -> String { + effects + .iter() + .map(|effect| match effect { + Effect::Assignment(v, expr) => format!("{v} = {expr};"), + Effect::Assertion(Assertion { + lhs, + rhs, + expected_equal, + }) => { + format!( + "assert {lhs} {} {rhs};", + if *expected_equal { "==" } else { "!=" } + ) + } + Effect::Lookup(id, args) => { + format!( + "lookup({id}, [{}]);", + args.iter() + .map(|arg| match arg { + LookupArgument::Known(k) => format!("Known({k})"), + LookupArgument::Unknown(u) => format!("Unknown({u})"), + }) + .join(", ") + ) + } + Effect::RangeConstraint(..) => { + panic!("Range constraints should not be part of the code.") + } + }) + .join("\n") + } + + #[test] + fn simple_polynomial_solving() { + let input = "let X; let Y; let Z; X = 1; Y = X + 1; Z * Y = X + 10;"; + let analyzed: Analyzed = + powdr_pil_analyzer::analyze_string(input).unwrap(); + let fixed_data = FixedData::new(&analyzed, &[], &[], Default::default(), 0); + let mut witgen = WitgenInference::new(&fixed_data, vec![]); + for id in analyzed.identities.iter() { + witgen.process_identity(id, 0); + } + assert_eq!(witgen.known_cells().len(), 3); + let code = format_code(&witgen.code()); + assert_eq!(code, "X[0] = 1;\nY[0] = 2;\nZ[0] = -9223372034707292155;"); + } +} From 53dddb47c6ef599e793dca6e7d01bf5cb41bcb34 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 09:01:18 +0000 Subject: [PATCH 17/44] Add completeness flag. --- .../witgen/jit/affine_symbolic_expression.rs | 87 +++++++++++++------ executor/src/witgen/jit/witgen_inference.rs | 71 +++++++++++---- 2 files changed, 113 insertions(+), 45 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index e15f839265..06ecb4b8db 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -64,6 +64,27 @@ pub enum LookupArgument { Unknown(AffineSymbolicExpression), } +#[derive(Default)] +pub struct ProcessResult { + pub effects: Vec>, + pub complete: bool, +} + +impl ProcessResult { + pub fn empty() -> Self { + Self { + effects: vec![], + complete: false, + } + } + pub fn complete(effects: Vec>) -> Self { + Self { + effects, + complete: true, + } + } +} + /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, /// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time, /// and the `x_i` are unknown variables. @@ -167,12 +188,14 @@ impl AffineSymbolicExpression { pub fn solve( &self, range_constraints: &impl RangeConstraintSet, - ) -> Result>, EvalError> { + ) -> Result, EvalError> { match self.coefficients.len() { 0 => { return if self.offset.is_known_zero() { - Ok(vec![]) + Ok(ProcessResult::complete(vec![])) } else { + // TODO this does not always mean it is unsatisfiable, since + // the offset could be a known variable. Err(EvalError::ConstraintUnsatisfiable(self.to_string())) }; } @@ -183,38 +206,46 @@ impl AffineSymbolicExpression { let assignment = Effect::Assignment(var.clone(), value); if let Some(coeff) = coeff.try_to_number() { assert!(!coeff.is_zero(), "Zero coefficient has not been removed."); - return Ok(vec![assignment]); + return Ok(ProcessResult::complete(vec![assignment])); } else { - return Ok(vec![ + // TODO if the coefficient is known and zero, it does not mean + // that the equation is unsatisfiable. So this should actually + // be a conditional assignment. + return Ok(ProcessResult::complete(vec![ Assertion::assert_is_nonzero(coeff.clone()), assignment, - ]); + ])); } } _ => {} } let r = self.solve_through_constraints(range_constraints); - if !r.is_empty() { + if r.complete { return Ok(r); } let negated = -self; let r = negated.solve_through_constraints(range_constraints); - if !r.is_empty() { + if r.complete { return Ok(r); } - Ok(self + + let effects = self .transfer_constraints(range_constraints) .into_iter() .chain(self.transfer_constraints(range_constraints)) - .collect()) + .collect(); + Ok(ProcessResult { + effects, + complete: false, + }) } /// Tries to solve a bit-decomposition equation. fn solve_through_constraints( &self, range_constraints: &impl RangeConstraintSet, - ) -> Vec> { + ) -> ProcessResult { let constrained_coefficients = self .coefficients .iter() @@ -229,7 +260,7 @@ impl AffineSymbolicExpression { // All the coefficients need to have known range constraints. if constrained_coefficients.len() != self.coefficients.len() { - return Vec::new(); + return ProcessResult::empty(); } // Check if they are mutually exclusive and compute assignments. @@ -239,7 +270,7 @@ impl AffineSymbolicExpression { let mask = *constraint.multiple(coeff).mask(); if !(mask & covered_bits).is_zero() { // Overlapping range constraints. - return vec![]; + return ProcessResult::empty(); } else { covered_bits |= mask; } @@ -251,7 +282,7 @@ impl AffineSymbolicExpression { } if covered_bits >= T::modulus() { - return vec![]; + return ProcessResult::empty(); } // We need to assert that the masks cover the offset, @@ -263,7 +294,7 @@ impl AffineSymbolicExpression { &self.offset | &T::from(covered_bits).into(), )); - effects + ProcessResult::complete(effects) } fn transfer_constraints( @@ -443,6 +474,7 @@ mod test { assert!(constr .solve(&SimpleRangeConstraintSet::default()) .unwrap() + .effects .is_empty()); } @@ -455,9 +487,10 @@ mod test { let seven = from_number(7); let ten = from_number(10); let constr = mul(&two, &x) + mul(&seven, &y) - ten; - let effects = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); - assert_eq!(effects.len(), 1); - let Effect::Assignment(var, expr) = &effects[0] else { + let result = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); + assert!(result.complete); + assert_eq!(result.effects.len(), 1); + let Effect::Assignment(var, expr) = &result.effects[0] else { panic!("Expected assignment"); }; assert_eq!(var.to_string(), "X"); @@ -484,14 +517,13 @@ mod test { + ten + z; // Without range constraints, this is not solvable. - assert!(constr - .solve(&SimpleRangeConstraintSet::default()) - .unwrap() - .is_empty()); + let result = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); + assert!(!result.complete && result.effects.is_empty()); // With range constraints, it should be solvable. - let effects = constr - .solve(&range_constraints) - .unwrap() + let result = constr.solve(&range_constraints).unwrap(); + assert!(result.complete); + let effects = result + .effects .into_iter() .map(|effect| match effect { Effect::Assignment(v, expr) => format!("{v} = {expr};\n"), @@ -538,9 +570,10 @@ assert (10 + Z) == ((10 + Z) | 4294967040); + mul(&c, &from_number(0x1000000)) + ten - z; - let effects = constr - .solve(&range_constraints) - .unwrap() + let result = constr.solve(&range_constraints).unwrap(); + assert!(!result.complete); + let effects = result + .effects .into_iter() .map(|effect| match effect { Effect::RangeConstraint(v, rc) => format!("{v}: {rc};\n"), diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 2458e6e516..658373a6f3 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -14,7 +14,7 @@ use crate::witgen::{ use super::{ super::{range_constraints::RangeConstraint, FixedData}, - affine_symbolic_expression::{AffineSymbolicExpression, Effect}, + affine_symbolic_expression::{AffineSymbolicExpression, Effect, ProcessResult}, cell::Cell, }; @@ -48,8 +48,11 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { self.code } - pub fn process_identity(&mut self, id: &Identity, row_offset: i32) { - let effects = match id { + /// Process an identity on a certain row. + /// Returns true if this identity/row pair was fully processed and + /// should not be considered again. + pub fn process_identity(&mut self, id: &Identity, row_offset: i32) -> bool { + let result = match id { Identity::Polynomial(PolynomialIdentity { expression, .. }) => { self.process_polynomial_identity(expression, row_offset) } @@ -71,25 +74,28 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { } _ => { // TODO - vec![] + ProcessResult::empty() } }; - self.ingest_effects(effects); + self.ingest_effects(result.effects); + result.complete } fn process_polynomial_identity( &self, expression: &'a Expression, offset: i32, - ) -> Vec> { + ) -> ProcessResult { if let Some(r) = self.evaluate(expression, 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. + // TODO this can also happen if we solve a constraint where all values are + // known but symbolic. r.solve(self).unwrap() } else { - vec![] + ProcessResult::empty() } } @@ -99,7 +105,7 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { left: &SelectedExpressions, right: &SelectedExpressions, offset: i32, - ) -> Vec> { + ) -> ProcessResult { // TODO: In the future, call the 'mutable state' to check if the // lookup can always be answered. @@ -127,7 +133,7 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { .filter(|e| e.try_to_known().is_none()) .collect_vec(); if unknown.len() == 1 && unknown[0].single_unknown_variable().is_some() { - return vec![Effect::Lookup( + let effects = vec![Effect::Lookup( lookup_id, lhs.into_iter() .map(|e| { @@ -139,11 +145,12 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { }) .collect(), )]; + return ProcessResult::complete(effects); } } } } - vec![] + ProcessResult::empty() } fn ingest_effects(&mut self, effects: Vec>) { @@ -316,18 +323,46 @@ mod test { .join("\n") } - #[test] - fn simple_polynomial_solving() { - let input = "let X; let Y; let Z; X = 1; Y = X + 1; Z * Y = X + 10;"; + fn solve_on_rows(input: &str, rows: &[i32], known_cells: Vec<(&str, i32)>) -> String { let analyzed: Analyzed = powdr_pil_analyzer::analyze_string(input).unwrap(); let fixed_data = FixedData::new(&analyzed, &[], &[], Default::default(), 0); - let mut witgen = WitgenInference::new(&fixed_data, vec![]); - for id in analyzed.identities.iter() { - witgen.process_identity(id, 0); + let known_cells = known_cells.iter().map(|(name, row_offset)| { + let id = fixed_data.try_column_by_name(name).unwrap().id; + Cell { + column_name: name.to_string(), + id, + row_offset: *row_offset, + } + }); + let mut witgen = WitgenInference::new(&fixed_data, known_cells); + let mut complete = HashSet::new(); + for _ in 0..4 { + for row in rows { + for id in analyzed.identities.iter() { + if !complete.contains(&(id.id(), row)) && witgen.process_identity(id, *row) { + complete.insert((id.id(), row)); + } + } + } } - assert_eq!(witgen.known_cells().len(), 3); - let code = format_code(&witgen.code()); + format_code(&witgen.code()) + } + + #[test] + fn simple_polynomial_solving() { + let input = "let X; let Y; let Z; X = 1; Y = X + 1; Z * Y = X + 10;"; + let code = solve_on_rows(input, &[0], vec![]); assert_eq!(code, "X[0] = 1;\nY[0] = 2;\nZ[0] = -9223372034707292155;"); } + + #[test] + fn fib() { + let input = "let X; let Y; X' = Y; Y' = X + Y;"; + let code = solve_on_rows(input, &[0, 1], vec![("X", 0), ("Y", 0)]); + assert_eq!( + code, + "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);\nX[2] = Y[1];\nY[2] = (X[1] + Y[1]);" + ); + } } From 67e238e7726a6ba2bb4f75bafb28d8969270a3d0 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 10:03:31 +0000 Subject: [PATCH 18/44] Change completeness condition. --- executor/src/witgen/jit/witgen_inference.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 658373a6f3..69ad807da5 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -337,11 +337,11 @@ mod test { }); let mut witgen = WitgenInference::new(&fixed_data, known_cells); let mut complete = HashSet::new(); - for _ in 0..4 { + while complete.len() != analyzed.identities.len() * rows.len() { for row in rows { for id in analyzed.identities.iter() { - if !complete.contains(&(id.id(), row)) && witgen.process_identity(id, *row) { - complete.insert((id.id(), row)); + if !complete.contains(&(id.id(), *row)) && witgen.process_identity(id, *row) { + complete.insert((id.id(), *row)); } } } From 9a8b3ade497763e2e048abfcd3df060f0785f1ba Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 10:09:20 +0000 Subject: [PATCH 19/44] pow --- executor/src/witgen/jit/witgen_inference.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 69ad807da5..25792ed41a 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -6,7 +6,7 @@ use powdr_ast::analyzed::{ AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, LookupIdentity, PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, SelectedExpressions, }; -use powdr_number::FieldElement; +use powdr_number::{FieldElement, LargeInt}; use crate::witgen::{ global_constraints::RangeConstraintSet, jit::affine_symbolic_expression::LookupArgument, @@ -244,7 +244,13 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { AlgebraicBinaryOperator::Add => Some(&left + &right), AlgebraicBinaryOperator::Sub => Some(&left - &right), AlgebraicBinaryOperator::Mul => left.try_mul(&right), - _ => todo!(), + AlgebraicBinaryOperator::Pow => { + let result = left + .try_to_known()? + .try_to_number()? + .pow(right.try_to_known()?.try_to_number()?.to_integer()); + Some(AffineSymbolicExpression::from_number(result)) + } } } From 236f056f18a0449f223fdc33c22de242f2c681a9 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 10:10:05 +0000 Subject: [PATCH 20/44] remove todo --- executor/src/witgen/jit/witgen_inference.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 25792ed41a..e1fcb0cb4a 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -267,8 +267,6 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { } impl RangeConstraintSet for WitgenInference<'_, T> { - // TODO would be nice to use &Cell, but this leads to lifetime trouble - // in the solve() function. fn range_constraint(&self, cell: Cell) -> Option> { self.fixed_data .global_range_constraints From cad7f37b3b3e02686aab364bc7482b89310bf707 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 11:01:51 +0000 Subject: [PATCH 21/44] Simplify and fix solving and remove unused conversion. --- .../witgen/jit/affine_symbolic_expression.rs | 73 +++++++------------ 1 file changed, 28 insertions(+), 45 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 2dd7436757..f89ec896e6 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -152,46 +152,43 @@ impl AffineSymbolicExpression { &self, range_constraints: &impl RangeConstraintSet, ) -> Result>, EvalError> { - match self.coefficients.len() { + Ok(match self.coefficients.len() { 0 => { - return if self.offset.is_known_zero() { - Ok(vec![]) + if self.offset.is_known_zero() { + vec![] } else { - Err(EvalError::ConstraintUnsatisfiable(self.to_string())) - }; + return Err(EvalError::ConstraintUnsatisfiable(self.to_string())); + } } 1 => { let (var, coeff) = self.coefficients.iter().next().unwrap(); assert!(!coeff.is_known_zero()); - let value = self.offset.field_div(&coeff.neg()); - let assignment = Effect::Assignment(var.clone(), value); - if let Some(coeff) = coeff.try_to_number() { - assert!(!coeff.is_zero(), "Zero coefficient has not been removed."); - return Ok(vec![assignment]); + if coeff.try_to_number().is_some() { + let value = self.offset.field_div(&-coeff); + vec![Effect::Assignment(var.clone(), value)] } else { - return Ok(vec![ - Assertion::assert_is_nonzero(coeff.clone()), - assignment, - ]); + // We can only solve this if we know that the coefficient cannot be zero. + // TODO We can either do this by adding a condition or we can check the range constraint. + vec![] } } - _ => {} - } - - let r = self.solve_through_constraints(range_constraints); - if !r.is_empty() { - return Ok(r); - } - let negated = -self; - let r = negated.solve_through_constraints(range_constraints); - if !r.is_empty() { - return Ok(r); - } - Ok(self - .transfer_constraints(range_constraints) - .into_iter() - .chain(self.transfer_constraints(range_constraints)) - .collect()) + _ => { + let r = self.solve_through_constraints(range_constraints); + if !r.is_empty() { + return Ok(r); + } + let negated = -self; + let r = negated.solve_through_constraints(range_constraints); + if !r.is_empty() { + r + } else { + self.transfer_constraints(range_constraints) + .into_iter() + .chain(self.transfer_constraints(range_constraints)) + .collect() + } + } + }) } /// Tries to solve a bit-decomposition equation. @@ -368,20 +365,6 @@ impl Neg for AffineSymbolicExpression { } } -impl TryFrom<&AffineSymbolicExpression> - for SymbolicExpression -{ - type Error = (); - - fn try_from(value: &AffineSymbolicExpression) -> Result { - if value.coefficients.is_empty() { - Ok(value.offset.clone()) - } else { - Err(()) - } - } -} - #[cfg(test)] mod test { use powdr_number::GoldilocksField; From ef33a568695cf1347611026eb34600067eab6858 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 11:54:46 +0000 Subject: [PATCH 22/44] Make use of range constraints in division check. --- .../src/witgen/jit/affine_symbolic_expression.rs | 9 ++++++--- executor/src/witgen/jit/symbolic_expression.rs | 11 +++++++++++ executor/src/witgen/range_constraints.rs | 14 ++++++++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index f89ec896e6..89c92848ce 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -162,13 +162,16 @@ impl AffineSymbolicExpression { } 1 => { let (var, coeff) = self.coefficients.iter().next().unwrap(); - assert!(!coeff.is_known_zero()); - if coeff.try_to_number().is_some() { + assert!( + !coeff.is_known_zero(), + "Zero coefficient has not been removed." + ); + if coeff.is_known_nonzero() { let value = self.offset.field_div(&-coeff); vec![Effect::Assignment(var.clone(), value)] } else { // We can only solve this if we know that the coefficient cannot be zero. - // TODO We can either do this by adding a condition or we can check the range constraint. + // TODO We could do this by adding a runtime condition vec![] } } diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index b6e689b570..60b4dd43d5 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -61,6 +61,17 @@ impl SymbolicExpression { self.try_to_number().map_or(false, |n| n == -T::from(1)) } + pub fn is_known_nonzero(&self) -> bool { + // Only checking range constraint is enough since if this is a known + // fixed value, we will get a range constraint with just a single value. + if let Some(rc) = self.range_constraint() { + !rc.allows_value(0.into()) + } else { + // unknown + false + } + } + pub fn range_constraint(&self) -> Option> { match self { SymbolicExpression::Concrete(v) => Some(RangeConstraint::from_value(*v)), diff --git a/executor/src/witgen/range_constraints.rs b/executor/src/witgen/range_constraints.rs index 8e6d9748a8..231f54939b 100644 --- a/executor/src/witgen/range_constraints.rs +++ b/executor/src/witgen/range_constraints.rs @@ -82,6 +82,20 @@ impl RangeConstraint { range_width(self.min, self.max) } + /// Returns true if `v` is an allowed value for this range constraint. + pub fn allows_value(&self, v: T) -> bool { + if self.min <= self.max { + if v < self.min || v > self.max { + return false; + } + } else { + if v < self.min && v > self.max { + return false; + } + } + v.to_integer() & self.mask == v.to_integer() + } + /// The range constraint of the sum of two expressions. pub fn combine_sum(&self, other: &Self) -> Self { // TODO we could use "add_with_carry" to see if this created an overflow. From b3e81e33925137b00159b19ae95be6d81bf97a29 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 11:55:50 +0000 Subject: [PATCH 23/44] Typos --- executor/src/witgen/jit/affine_symbolic_expression.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 89c92848ce..5233db6b71 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -393,7 +393,7 @@ mod test { } #[test] - fn unsolveable() { + fn unsolvable() { let r = from_number(10).solve(&SimpleRangeConstraintSet::default()); assert!(r.is_err()); } @@ -408,7 +408,7 @@ mod test { } #[test] - fn solveable_without_vars() { + fn solvable_without_vars() { let constr = &from_number(0); assert!(constr .solve(&SimpleRangeConstraintSet::default()) From 84c5d7575d52dc71905190a6e92c8a2b134bb928 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 12:00:07 +0000 Subject: [PATCH 24/44] Comment display and remove parentheses for unary. --- executor/src/witgen/jit/affine_symbolic_expression.rs | 7 ++++--- executor/src/witgen/jit/symbolic_expression.rs | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 5233db6b71..71c58e6911 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -66,6 +66,7 @@ pub struct AffineSymbolicExpression { offset: SymbolicExpression, } +/// Display for affine symbolic expressions, for informational purposes only. impl Display for AffineSymbolicExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { if self.coefficients.is_empty() { @@ -481,9 +482,9 @@ mod test { .to_string(); assert_eq!( effects, - "a = ((-((10 + Z)) & 65280) // 256); -b = ((-((10 + Z)) & 16711680) // 65536); -c = ((-((10 + Z)) & 4278190080) // 16777216); + "a = ((-(10 + Z) & 65280) // 256); +b = ((-(10 + Z) & 16711680) // 65536); +c = ((-(10 + Z) & 4278190080) // 16777216); assert (10 + Z) == ((10 + Z) | 4294967040); " ); diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 60b4dd43d5..08f2f8fa6d 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -91,6 +91,7 @@ impl SymbolicExpression { } } +/// Display for affine symbolic expressions, for informational purposes only. impl Display for SymbolicExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { @@ -105,7 +106,7 @@ impl Display for SymbolicExpression { SymbolicExpression::BinaryOperation(lhs, op, rhs, _) => { write!(f, "({lhs} {op} {rhs})") } - SymbolicExpression::UnaryOperation(op, expr, _) => write!(f, "{op}({expr})"), + SymbolicExpression::UnaryOperation(op, expr, _) => write!(f, "{op}{expr}"), } } } From 99cbb3da88107b0c2179a2443c97b99f6206748c Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 12:22:55 +0000 Subject: [PATCH 25/44] extend documentation. --- executor/src/witgen/jit/affine_symbolic_expression.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 71c58e6911..fd8c312583 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -58,8 +58,9 @@ impl Assertion { } /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`, -/// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time, -/// and the `x_i` are unknown variables. +/// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time +/// (which can still include variables or symbols, which are only known at run-time), +/// and the `x_i` are variables that are unknown at this point. #[derive(Debug, Clone)] pub struct AffineSymbolicExpression { coefficients: BTreeMap>, From 7c2f994842438bb493d0e076bf0ebb9783e7def5 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 12:23:55 +0000 Subject: [PATCH 26/44] clippy --- executor/src/witgen/range_constraints.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/executor/src/witgen/range_constraints.rs b/executor/src/witgen/range_constraints.rs index 231f54939b..f19fb55a65 100644 --- a/executor/src/witgen/range_constraints.rs +++ b/executor/src/witgen/range_constraints.rs @@ -88,10 +88,8 @@ impl RangeConstraint { if v < self.min || v > self.max { return false; } - } else { - if v < self.min && v > self.max { - return false; - } + } else if v < self.min && v > self.max { + return false; } v.to_integer() & self.mask == v.to_integer() } From a21ca0c84f5a496d77622d1483b1cf12e6a8a9a5 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 12:30:46 +0000 Subject: [PATCH 27/44] simplify "allows_value" --- executor/src/witgen/range_constraints.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/executor/src/witgen/range_constraints.rs b/executor/src/witgen/range_constraints.rs index f19fb55a65..a4060fe7c4 100644 --- a/executor/src/witgen/range_constraints.rs +++ b/executor/src/witgen/range_constraints.rs @@ -84,14 +84,14 @@ impl RangeConstraint { /// Returns true if `v` is an allowed value for this range constraint. pub fn allows_value(&self, v: T) -> bool { - if self.min <= self.max { - if v < self.min || v > self.max { - return false; - } - } else if v < self.min && v > self.max { - return false; + let wrapping = self.min > self.max; + if (!wrapping && (v < self.min || v > self.max)) + || (wrapping && (v < self.min && v > self.max)) + { + false + } else { + v.to_integer() & self.mask == v.to_integer() } - v.to_integer() & self.mask == v.to_integer() } /// The range constraint of the sum of two expressions. From 479a96d5fe973a35e2bf82c415a0633681b88165 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 13:49:42 +0000 Subject: [PATCH 28/44] Store all range constraints with expressions. --- .../witgen/jit/affine_symbolic_expression.rs | 226 +++++++++--------- .../src/witgen/jit/symbolic_expression.rs | 4 +- 2 files changed, 118 insertions(+), 112 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index fd8c312583..ec775a6800 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -1,14 +1,14 @@ use std::{ collections::BTreeMap, fmt::{self, Display, Formatter}, - ops::{Add, Neg, Sub}, + ops::{Add, Mul, Neg, Sub}, }; use itertools::Itertools; use num_traits::Zero; use powdr_number::FieldElement; -use crate::witgen::{global_constraints::RangeConstraintSet, EvalError}; +use crate::witgen::EvalError; use super::{super::range_constraints::RangeConstraint, symbolic_expression::SymbolicExpression}; @@ -61,10 +61,12 @@ impl Assertion { /// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time /// (which can still include variables or symbols, which are only known at run-time), /// and the `x_i` are variables that are unknown at this point. +/// It also stores range constraints for all unknown variables. #[derive(Debug, Clone)] pub struct AffineSymbolicExpression { coefficients: BTreeMap>, offset: SymbolicExpression, + range_constraints: BTreeMap>, } /// Display for affine symbolic expressions, for informational purposes only. @@ -100,18 +102,20 @@ impl From> for AffineSymbolicExpres AffineSymbolicExpression { coefficients: Default::default(), offset: k, + range_constraints: Default::default(), } } } impl AffineSymbolicExpression { - pub fn from_known_variable(var: V) -> Self { - SymbolicExpression::from_var(var).into() + pub fn from_known_variable(var: V, rc: Option>) -> Self { + SymbolicExpression::from_var(var, rc).into() } - pub fn from_unknown_variable(var: V) -> Self { + pub fn from_unknown_variable(var: V, rc: Option>) -> Self { AffineSymbolicExpression { - coefficients: [(var, T::from(1).into())].into_iter().collect(), + coefficients: [(var.clone(), T::from(1).into())].into_iter().collect(), offset: SymbolicExpression::from(T::from(0)), + range_constraints: rc.into_iter().map(|rc| (var.clone(), rc)).collect(), } } pub fn from_number(n: T) -> Self { @@ -128,32 +132,20 @@ impl AffineSymbolicExpression { } /// Tries to multiply this expression with another one. - /// Returns `None` if the result would be quadratic. + /// Returns `None` if the result would be quadratic, i.e. + /// if both expressions contain unknown variables. pub fn try_mul(&self, other: &Self) -> Option { - if !self.coefficients.is_empty() && !other.coefficients.is_empty() { - return None; - } - let (multiplier, coefficients, offset) = if self.coefficients.is_empty() { - (&self.offset, &other.coefficients, &other.offset) + if let Some(multiplier) = other.try_to_known() { + Some(self.clone() * multiplier) + } else if let Some(multiplier) = self.try_to_known() { + Some(other.clone() * multiplier) } else { - (&other.offset, &self.coefficients, &self.offset) - }; - let coefficients = coefficients - .iter() - .map(|(var, coeff)| (var.clone(), coeff * multiplier)) - .collect(); - let offset = offset * multiplier; - Some(AffineSymbolicExpression { - coefficients, - offset, - }) + None + } } /// Solves the equation `self = 0` and returns how to compute the solution. - pub fn solve( - &self, - range_constraints: &impl RangeConstraintSet, - ) -> Result>, EvalError> { + pub fn solve(&self) -> Result>, EvalError> { Ok(match self.coefficients.len() { 0 => { if self.offset.is_known_zero() { @@ -178,45 +170,41 @@ impl AffineSymbolicExpression { } } _ => { - let r = self.solve_through_constraints(range_constraints); - if !r.is_empty() { - return Ok(r); - } - let negated = -self; - let r = negated.solve_through_constraints(range_constraints); + let r = self.solve_bit_decomposition(); if !r.is_empty() { r } else { - self.transfer_constraints(range_constraints) - .into_iter() - .chain(self.transfer_constraints(range_constraints)) - .collect() + let negated = -self; + let r = negated.solve_bit_decomposition(); + if !r.is_empty() { + r + } else { + self.transfer_constraints() + .into_iter() + .chain(negated.transfer_constraints()) + .collect() + } } } }) } /// Tries to solve a bit-decomposition equation. - fn solve_through_constraints( - &self, - range_constraints: &impl RangeConstraintSet, - ) -> Vec> { + fn solve_bit_decomposition(&self) -> Vec> { + // All the coefficients need to be known numbers and the + // variables need to be range-constrained. let constrained_coefficients = self .coefficients .iter() - .filter_map(|(var, coeff)| { - coeff.try_to_number().and_then(|c| { - range_constraints - .range_constraint(var.clone()) - .map(|rc| (var.clone(), c, rc)) - }) + .map(|(var, coeff)| { + let c = coeff.try_to_number()?; + let rc = self.range_constraints.get(var)?; + Some((var.clone(), c, rc)) }) - .collect_vec(); - - // All the coefficients need to have known range constraints. - if constrained_coefficients.len() != self.coefficients.len() { - return Vec::new(); - } + .collect::>>(); + let Some(constrained_coefficients) = constrained_coefficients else { + return vec![]; + }; // Check if they are mutually exclusive and compute assignments. let mut covered_bits: ::Integer = 0.into(); @@ -252,10 +240,7 @@ impl AffineSymbolicExpression { effects } - fn transfer_constraints( - &self, - range_constraints: &impl RangeConstraintSet, - ) -> Vec> { + fn transfer_constraints(&self) -> Vec> { // We are looking for X = a * Y + b * Z + ... or -X = a * Y + b * Z + ... // where X is least constrained. @@ -265,8 +250,8 @@ impl AffineSymbolicExpression { .filter(|(_var, coeff)| coeff.is_known_one() || coeff.is_known_minus_one()) .max_by_key(|(var, _c)| { // Sort so that we get the least constrained variable. - range_constraints - .range_constraint((*var).clone()) + self.range_constraints + .get(var) .map(|c| c.range_width()) .unwrap_or_else(|| T::modulus()) }) @@ -280,11 +265,9 @@ impl AffineSymbolicExpression { .iter() .filter(|(var, _)| *var != solve_for) .map(|(var, coeff)| { - coeff.try_to_number().and_then(|coeff| { - range_constraints - .range_constraint(var.clone()) - .map(|constraint| constraint.multiple(coeff)) - }) + let coeff = coeff.try_to_number()?; + let rc = self.range_constraints.get(var)?; + Some(rc.multiple(coeff)) }) .chain(std::iter::once(self.offset.range_constraint())) .collect::>>() @@ -308,17 +291,25 @@ impl Add for &AffineSymbolicExpression { fn add(self, rhs: Self) -> Self::Output { let mut coefficients = self.coefficients.clone(); + let mut range_constraints = self.range_constraints.clone(); for (var, coeff) in &rhs.coefficients { coefficients .entry(var.clone()) .and_modify(|f| *f = &*f + coeff) .or_insert_with(|| coeff.clone()); + if let Some(range_right) = rhs.range_constraints.get(var) { + range_constraints + .entry(var.clone()) + .and_modify(|rc| *rc = rc.conjunction(range_right)) + .or_insert_with(|| range_right.clone()); + } } coefficients.retain(|_, f| !f.is_known_zero()); let offset = &self.offset + &rhs.offset; AffineSymbolicExpression { coefficients, offset, + range_constraints, } } } @@ -358,6 +349,7 @@ impl Neg for &AffineSymbolicExpression { .map(|(var, coeff)| (var.clone(), -coeff)) .collect(), offset: -&self.offset, + range_constraints: self.range_constraints.clone(), } } } @@ -370,20 +362,27 @@ impl Neg for AffineSymbolicExpression { } } +/// Multiply by known symbolic expression. +impl Mul<&SymbolicExpression> + for AffineSymbolicExpression +{ + type Output = AffineSymbolicExpression; + + fn mul(mut self, rhs: &SymbolicExpression) -> Self::Output { + for coeff in self.coefficients.values_mut() { + *coeff = &*coeff * rhs; + } + self.offset = &self.offset * rhs; + self + } +} + #[cfg(test)] mod test { use powdr_number::GoldilocksField; use super::*; - #[derive(Default)] - struct SimpleRangeConstraintSet(BTreeMap<&'static str, RangeConstraint>); - impl RangeConstraintSet<&str, GoldilocksField> for SimpleRangeConstraintSet { - fn range_constraint(&self, var: &str) -> Option> { - self.0.get(var).cloned() - } - } - type Ase = AffineSymbolicExpression; fn from_number(x: i32) -> Ase { @@ -396,58 +395,73 @@ mod test { #[test] fn unsolvable() { - let r = from_number(10).solve(&SimpleRangeConstraintSet::default()); + let r = from_number(10).solve(); assert!(r.is_err()); } #[test] fn unsolvable_with_vars() { - let x = &Ase::from_known_variable("X"); - let y = &Ase::from_known_variable("Y"); + let x = &Ase::from_known_variable("X", None); + let y = &Ase::from_known_variable("Y", None); let constr = x + y - from_number(10); - let r = constr.solve(&SimpleRangeConstraintSet::default()); + let r = constr.solve(); assert!(r.is_err()); } #[test] fn solvable_without_vars() { let constr = &from_number(0); - assert!(constr - .solve(&SimpleRangeConstraintSet::default()) - .unwrap() - .is_empty()); + assert!(constr.solve().unwrap().is_empty()); } #[test] fn solve_simple_eq() { - let y = Ase::from_known_variable("y"); - let x = Ase::from_unknown_variable("X"); + let y = Ase::from_known_variable("y", None); + let x = Ase::from_unknown_variable("X", None); // 2 * X + 7 * y - 10 = 0 let two = from_number(2); let seven = from_number(7); let ten = from_number(10); let constr = mul(&two, &x) + mul(&seven, &y) - ten; - let effects = constr.solve(&SimpleRangeConstraintSet::default()).unwrap(); + let effects = constr.solve().unwrap(); assert_eq!(effects.len(), 1); let Effect::Assignment(var, expr) = &effects[0] else { panic!("Expected assignment"); }; assert_eq!(var.to_string(), "X"); - assert_eq!(expr.to_string(), "(((y * 7) + -10) / -2)"); + assert_eq!(expr.to_string(), "(((7 * y) + -10) / -2)"); + } + + #[test] + fn solve_div_by_range_constrained_var() { + let y = Ase::from_known_variable("y", None); + let z = Ase::from_known_variable("z", None); + let x = Ase::from_unknown_variable("X", None); + // z * X + 7 * y - 10 = 0 + let seven = from_number(7); + let ten = from_number(10); + let constr = mul(&z, &x) + mul(&seven, &y) - ten.clone(); + // If we do not range-constrain z, we cannot solve since we don't know if it might be zero. + let effects = constr.solve().unwrap(); + assert_eq!(effects.len(), 0); + let z = + Ase::from_known_variable("z", Some(RangeConstraint::from_range(10.into(), 20.into()))); + let constr = mul(&z, &x) + mul(&seven, &y) - ten; + let effects = constr.solve().unwrap(); + let Effect::Assignment(var, expr) = &effects[0] else { + panic!("Expected assignment"); + }; + assert_eq!(var.to_string(), "X"); + assert_eq!(expr.to_string(), "(((7 * y) + -10) / -z)"); } #[test] fn solve_bit_decomposition() { - let a = Ase::from_unknown_variable("a"); - let b = Ase::from_unknown_variable("b"); - let c = Ase::from_unknown_variable("c"); - let z = Ase::from_known_variable("Z"); - let range_constraints = SimpleRangeConstraintSet( - ["a", "b", "c"] - .into_iter() - .map(|var| (var, RangeConstraint::from_mask(0xffu32))) - .collect(), - ); + let rc = Some(RangeConstraint::from_mask(0xffu32)); + let a = Ase::from_unknown_variable("a", rc.clone()); + let b = Ase::from_unknown_variable("b", rc.clone()); + let c = Ase::from_unknown_variable("c", rc.clone()); + let z = Ase::from_known_variable("Z", None); // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 + Z = 0 let ten = from_number(10); let constr = mul(&a, &from_number(0x100)) @@ -456,13 +470,10 @@ mod test { + ten + z; // Without range constraints, this is not solvable. - assert!(constr - .solve(&SimpleRangeConstraintSet::default()) - .unwrap() - .is_empty()); + assert!(constr.solve().unwrap().is_empty()); // With range constraints, it should be solvable. let effects = constr - .solve(&range_constraints) + .solve() .unwrap() .into_iter() .map(|effect| match effect { @@ -493,16 +504,11 @@ assert (10 + Z) == ((10 + Z) | 4294967040); #[test] fn solve_constraint_transfer() { - let a = Ase::from_unknown_variable("a"); - let b = Ase::from_unknown_variable("b"); - let c = Ase::from_unknown_variable("c"); - let z = Ase::from_unknown_variable("Z"); - let range_constraints = SimpleRangeConstraintSet( - ["a", "b", "c"] - .into_iter() - .map(|var| (var, RangeConstraint::from_mask(0xffu32))) - .collect(), - ); + let rc = Some(RangeConstraint::from_mask(0xffu32)); + let a = Ase::from_unknown_variable("a", rc.clone()); + let b = Ase::from_unknown_variable("b", rc.clone()); + let c = Ase::from_unknown_variable("c", rc.clone()); + let z = Ase::from_unknown_variable("Z", None); // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 - Z = 0 let ten = from_number(10); let constr = mul(&a, &from_number(0x100)) @@ -511,7 +517,7 @@ assert (10 + Z) == ((10 + Z) | 4294967040); + ten - z; let effects = constr - .solve(&range_constraints) + .solve() .unwrap() .into_iter() .map(|effect| match effect { diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 08f2f8fa6d..db1bc634ca 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -45,8 +45,8 @@ pub enum UnaryOperator { } impl SymbolicExpression { - pub fn from_var(name: V) -> Self { - SymbolicExpression::Variable(name, None) + pub fn from_var(var: V, rc: Option>) -> Self { + SymbolicExpression::Variable(var, rc) } pub fn is_known_zero(&self) -> bool { From 1a77808c2b52e9f44a1184bb98d26ba52f3dd8ea Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 15:59:15 +0000 Subject: [PATCH 29/44] Improved allows_value. --- executor/src/witgen/range_constraints.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/executor/src/witgen/range_constraints.rs b/executor/src/witgen/range_constraints.rs index a4060fe7c4..cdd76a76a0 100644 --- a/executor/src/witgen/range_constraints.rs +++ b/executor/src/witgen/range_constraints.rs @@ -84,14 +84,13 @@ impl RangeConstraint { /// Returns true if `v` is an allowed value for this range constraint. pub fn allows_value(&self, v: T) -> bool { - let wrapping = self.min > self.max; - if (!wrapping && (v < self.min || v > self.max)) - || (wrapping && (v < self.min && v > self.max)) - { - false + let in_range = if self.min <= self.max { + self.min <= v && v <= self.max } else { - v.to_integer() & self.mask == v.to_integer() - } + v <= self.min || self.max <= v + }; + let in_mask = v.to_integer() & self.mask == v.to_integer(); + in_range && in_mask } /// The range constraint of the sum of two expressions. From 8c588f87548c51ed5f803484a26d798d7fa69839 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 16:15:01 +0000 Subject: [PATCH 30/44] Add another case for division. --- .../src/witgen/jit/affine_symbolic_expression.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index ec775a6800..ad47553802 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -156,16 +156,26 @@ impl AffineSymbolicExpression { } 1 => { let (var, coeff) = self.coefficients.iter().next().unwrap(); + // Solve "coeff * X + self.offset = 0" by division. assert!( !coeff.is_known_zero(), "Zero coefficient has not been removed." ); if coeff.is_known_nonzero() { + // In this case, we can always compute a solution. let value = self.offset.field_div(&-coeff); vec![Effect::Assignment(var.clone(), value)] + } else if self.offset.is_known_nonzero() { + // If the offset is not zero, then the coefficient must be non-zero, + // otherwise the constraint is violated. + let value = self.offset.field_div(&-coeff); + vec![ + Assertion::assert_is_nonzero(coeff.clone()), + Effect::Assignment(var.clone(), value), + ] } else { - // We can only solve this if we know that the coefficient cannot be zero. - // TODO We could do this by adding a runtime condition + // If this case, we could have an equation of the form + // 0 * X = 0, which is valid and generates no information about X. vec![] } } From 969e2ba977983df8ed53b67d5799115b286fe664 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 16:17:57 +0000 Subject: [PATCH 31/44] Fix tests. --- .../witgen/jit/affine_symbolic_expression.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index ad47553802..02e1d9c086 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -468,7 +468,8 @@ mod test { #[test] fn solve_bit_decomposition() { let rc = Some(RangeConstraint::from_mask(0xffu32)); - let a = Ase::from_unknown_variable("a", rc.clone()); + // First try without range constrain on a + let a = Ase::from_unknown_variable("a", None); let b = Ase::from_unknown_variable("b", rc.clone()); let c = Ase::from_unknown_variable("c", rc.clone()); let z = Ase::from_known_variable("Z", None); @@ -477,11 +478,17 @@ mod test { let constr = mul(&a, &from_number(0x100)) + mul(&b, &from_number(0x10000)) + mul(&c, &from_number(0x1000000)) - + ten - + z; + + ten.clone() + + z.clone(); // Without range constraints, this is not solvable. assert!(constr.solve().unwrap().is_empty()); - // With range constraints, it should be solvable. + // Now add the range constraint on a, it should be solvable. + let a = Ase::from_unknown_variable("a", rc.clone()); + let constr = mul(&a, &from_number(0x100)) + + mul(&b, &from_number(0x10000)) + + mul(&c, &from_number(0x1000000)) + + ten.clone() + + z; let effects = constr .solve() .unwrap() @@ -537,11 +544,11 @@ assert (10 + Z) == ((10 + Z) | 4294967040); .format("") .to_string(); // It appears twice because we solve the positive and the negated equation. - // Maybe it is enough to only solve one. + // Note that the negated version has a different bit mask. assert_eq!( effects, "Z: [10, 4294967050] & 0xffffff0a; -Z: [10, 4294967050] & 0xffffff0a; +Z: [10, 4294967050] & 0xffffffff; " ); } From 88cd2648659e6eaf0918cf0318ae90a089cbb61d Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 16:34:26 +0000 Subject: [PATCH 32/44] Rename Variable to Symbol and relax one check. --- .../witgen/jit/affine_symbolic_expression.rs | 50 +++++++++++-------- .../src/witgen/jit/symbolic_expression.rs | 19 +++---- 2 files changed, 39 insertions(+), 30 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 02e1d9c086..5da3713c5a 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -14,11 +14,11 @@ use super::{super::range_constraints::RangeConstraint, symbolic_expression::Symb /// The effect of solving a symbolic equation. pub enum Effect { - /// variable can be assigned a value. + /// Variable can be assigned a value. Assignment(V, SymbolicExpression), - /// we learnt a new range constraint on variable. + /// We learnt a new range constraint on variable. RangeConstraint(V, RangeConstraint), - /// a run-time assertion. If this fails, we have conflicting constraints. + /// A run-time assertion. If this fails, we have conflicting constraints. Assertion(Assertion), } @@ -107,9 +107,15 @@ impl From> for AffineSymbolicExpres } } +impl From for AffineSymbolicExpression { + fn from(k: T) -> Self { + SymbolicExpression::from(k).into() + } +} + impl AffineSymbolicExpression { - pub fn from_known_variable(var: V, rc: Option>) -> Self { - SymbolicExpression::from_var(var, rc).into() + pub fn from_known_symbol(symbol: V, rc: Option>) -> Self { + SymbolicExpression::from_symbol(symbol, rc).into() } pub fn from_unknown_variable(var: V, rc: Option>) -> Self { AffineSymbolicExpression { @@ -118,9 +124,6 @@ impl AffineSymbolicExpression { range_constraints: rc.into_iter().map(|rc| (var.clone(), rc)).collect(), } } - pub fn from_number(n: T) -> Self { - SymbolicExpression::from(n).into() - } /// If this expression does not contain unknown variables, returns the symbolic expression. pub fn try_to_known(&self) -> Option<&SymbolicExpression> { @@ -145,13 +148,16 @@ impl AffineSymbolicExpression { } /// Solves the equation `self = 0` and returns how to compute the solution. + /// The solution can contain assignments to multiple variables. + /// If no way to solve the equation has been found, returns the empty vector. + /// If the equation is known to be unsolvable, returns an error. pub fn solve(&self) -> Result>, EvalError> { Ok(match self.coefficients.len() { 0 => { - if self.offset.is_known_zero() { - vec![] - } else { + if self.offset.is_known_nonzero() { return Err(EvalError::ConstraintUnsatisfiable(self.to_string())); + } else { + vec![] } } 1 => { @@ -396,7 +402,7 @@ mod test { type Ase = AffineSymbolicExpression; fn from_number(x: i32) -> Ase { - Ase::from_number(GoldilocksField::from(x)) + GoldilocksField::from(x).into() } fn mul(a: &Ase, b: &Ase) -> Ase { @@ -411,11 +417,13 @@ mod test { #[test] fn unsolvable_with_vars() { - let x = &Ase::from_known_variable("X", None); - let y = &Ase::from_known_variable("Y", None); + let x = &Ase::from_known_symbol("X", None); + let y = &Ase::from_known_symbol("Y", None); let constr = x + y - from_number(10); - let r = constr.solve(); - assert!(r.is_err()); + // We cannot solve it but also cannot know it is unsolvable. + assert!(constr.solve().unwrap().is_empty()); + // But if we know the values, we can be sure there is a conflict. + assert!(from_number(10).solve().is_err()); } #[test] @@ -426,7 +434,7 @@ mod test { #[test] fn solve_simple_eq() { - let y = Ase::from_known_variable("y", None); + let y = Ase::from_known_symbol("y", None); let x = Ase::from_unknown_variable("X", None); // 2 * X + 7 * y - 10 = 0 let two = from_number(2); @@ -444,8 +452,8 @@ mod test { #[test] fn solve_div_by_range_constrained_var() { - let y = Ase::from_known_variable("y", None); - let z = Ase::from_known_variable("z", None); + let y = Ase::from_known_symbol("y", None); + let z = Ase::from_known_symbol("z", None); let x = Ase::from_unknown_variable("X", None); // z * X + 7 * y - 10 = 0 let seven = from_number(7); @@ -455,7 +463,7 @@ mod test { let effects = constr.solve().unwrap(); assert_eq!(effects.len(), 0); let z = - Ase::from_known_variable("z", Some(RangeConstraint::from_range(10.into(), 20.into()))); + Ase::from_known_symbol("z", Some(RangeConstraint::from_range(10.into(), 20.into()))); let constr = mul(&z, &x) + mul(&seven, &y) - ten; let effects = constr.solve().unwrap(); let Effect::Assignment(var, expr) = &effects[0] else { @@ -472,7 +480,7 @@ mod test { let a = Ase::from_unknown_variable("a", None); let b = Ase::from_unknown_variable("b", rc.clone()); let c = Ase::from_unknown_variable("c", rc.clone()); - let z = Ase::from_known_variable("Z", None); + let z = Ase::from_known_symbol("Z", None); // a * 0x100 + b * 0x10000 + c * 0x1000000 + 10 + Z = 0 let ten = from_number(10); let constr = mul(&a, &from_number(0x100)) diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index db1bc634ca..8f3ddd470f 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -12,11 +12,12 @@ use crate::witgen::range_constraints::RangeConstraint; /// involving known cells or variables and compile-time constants. /// Each of the sub-expressions can have its own range constraint. #[derive(Debug, Clone)] -pub enum SymbolicExpression { +pub enum SymbolicExpression { /// A concrete constant value known at compile time. Concrete(T), - /// A symbolic value known at run-time, referencing either a cell or a local variable. - Variable(V, Option>), + /// A symbolic value known at run-time, referencing a cell, + /// an input, a local variable or whatever it is used for. + Symbol(S, Option>), BinaryOperation( Rc, BinaryOperator, @@ -44,9 +45,9 @@ pub enum UnaryOperator { Neg, } -impl SymbolicExpression { - pub fn from_var(var: V, rc: Option>) -> Self { - SymbolicExpression::Variable(var, rc) +impl SymbolicExpression { + pub fn from_symbol(symbol: S, rc: Option>) -> Self { + SymbolicExpression::Symbol(symbol, rc) } pub fn is_known_zero(&self) -> bool { @@ -75,7 +76,7 @@ impl SymbolicExpression { pub fn range_constraint(&self) -> Option> { match self { SymbolicExpression::Concrete(v) => Some(RangeConstraint::from_value(*v)), - SymbolicExpression::Variable(.., rc) + SymbolicExpression::Symbol(.., rc) | SymbolicExpression::BinaryOperation(.., rc) | SymbolicExpression::UnaryOperation(.., rc) => rc.clone(), } @@ -84,7 +85,7 @@ impl SymbolicExpression { pub fn try_to_number(&self) -> Option { match self { SymbolicExpression::Concrete(n) => Some(*n), - SymbolicExpression::Variable(..) + SymbolicExpression::Symbol(..) | SymbolicExpression::BinaryOperation(..) | SymbolicExpression::UnaryOperation(..) => None, } @@ -102,7 +103,7 @@ impl Display for SymbolicExpression { write!(f, "-{}", -*n) } } - SymbolicExpression::Variable(name, _) => write!(f, "{name}"), + SymbolicExpression::Symbol(name, _) => write!(f, "{name}"), SymbolicExpression::BinaryOperation(lhs, op, rhs, _) => { write!(f, "({lhs} {op} {rhs})") } From 818362075d532b0acc2f0fbc7c3c422e890436a2 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 16:35:23 +0000 Subject: [PATCH 33/44] clippy --- executor/src/witgen/jit/affine_symbolic_expression.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 5da3713c5a..afdf683bd7 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -140,10 +140,9 @@ impl AffineSymbolicExpression { pub fn try_mul(&self, other: &Self) -> Option { if let Some(multiplier) = other.try_to_known() { Some(self.clone() * multiplier) - } else if let Some(multiplier) = self.try_to_known() { - Some(other.clone() * multiplier) } else { - None + self.try_to_known() + .map(|multiplier| other.clone() * multiplier) } } From be36d962aa355fcebf4e30c720d74591cc917636 Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 11 Dec 2024 16:59:06 +0000 Subject: [PATCH 34/44] fix test --- executor/src/witgen/jit/affine_symbolic_expression.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index adb7f0e44c..1c1e41a081 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -454,9 +454,9 @@ mod test { let x = &Ase::from_known_symbol("X", None); let y = &Ase::from_known_symbol("Y", None); let constr = x + y - from_number(10); - // We cannot solve it but also cannot know it is unsolvable. + // We cannot solve it, but we can also not learn anything new from it. let result = constr.solve().unwrap(); - assert!(!result.complete && result.effects.is_empty()); + assert!(result.complete && result.effects.is_empty()); // But if we know the values, we can be sure there is a conflict. assert!(from_number(10).solve().is_err()); } From bfb53719e5597293ff115141db3558fd6f4505d0 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 11:28:38 +0000 Subject: [PATCH 35/44] Access to fixed columns. --- executor/src/witgen/jit/witgen_inference.rs | 130 ++++++++++++++------ 1 file changed, 94 insertions(+), 36 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index e7ad5458be..eb13749bd7 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -3,8 +3,9 @@ use std::collections::{HashMap, HashSet}; use itertools::Itertools; use powdr_ast::analyzed::{ AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, - AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, LookupIdentity, - PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, SelectedExpressions, + AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Challenge, Identity, + LookupIdentity, PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, + SelectedExpressions, }; use powdr_number::FieldElement; @@ -20,20 +21,23 @@ use super::{ /// This component can generate code that solves identities. /// It needs a driver that tells it which identities to process on which rows. -pub struct WitgenInference<'a, T: FieldElement> { +pub struct WitgenInference<'a, T: FieldElement, RefEval: ReferenceEvaluator> { fixed_data: &'a FixedData<'a, T>, + reference_evaluator: RefEval, range_constraints: HashMap>, known_cells: HashSet, code: Vec>, } -impl<'a, T: FieldElement> WitgenInference<'a, T> { +impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, RefEval> { pub fn new( fixed_data: &'a FixedData<'a, T>, + reference_evaluator: RefEval, known_cells: impl IntoIterator, ) -> Self { Self { fixed_data, + reference_evaluator, range_constraints: Default::default(), known_cells: known_cells.into_iter().collect(), code: Default::default(), @@ -57,21 +61,11 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { self.process_polynomial_identity(expression, row_offset) } Identity::Lookup(LookupIdentity { - id, - source: _, - left, - right, + id, left, right, .. }) | Identity::PhantomLookup(PhantomLookupIdentity { - id, - source: _, - left, - right, - multiplicity: _, - }) => { - // TODO multiplicity? - self.process_lookup(*id, left, right, row_offset) - } + id, left, right, .. + }) => self.process_lookup(*id, left, right, row_offset), _ => { // TODO ProcessResult::empty() @@ -155,12 +149,12 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { for e in effects { match &e { Effect::Assignment(cell, assignment) => { + self.known_cells.insert(cell.clone()); if let Some(rc) = assignment.range_constraint() { // If the cell was determined to be a constant, we add this // as a range constraint, so we can use it in future evaluations. self.add_range_constraint(cell.clone(), rc); } - self.known_cells.insert(cell.clone()); self.code.push(e); } Effect::RangeConstraint(cell, rc) => { @@ -184,9 +178,14 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { let rc = self .range_constraint(cell.clone()) .map_or(rc.clone(), |existing_rc| existing_rc.conjunction(&rc)); - // TODO if the conjuntion results in a single value, make the cell known. - // TODO but we also need to generate code for the assignment! - self.range_constraints.insert(cell, rc); + if !self.known_cells.contains(&cell) { + if let Some(v) = rc.try_to_single_value() { + // Special case: Cell is fixed to a constant by range constraints only. + self.known_cells.insert(cell.clone()); + self.code.push(Effect::Assignment(cell.clone(), v.into())); + } + } + self.range_constraints.insert(cell.clone(), rc); } fn evaluate( @@ -197,16 +196,7 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { Some(match expr { Expression::Reference(r) => { if r.is_fixed() { - todo!() - // let mut row = self.latch_row as i64 + offset as i64; - // while row < 0 { - // row += self.block_size as i64; - // } - // // TODO at some point we should check that all of the fixed columns are periodic. - // // TODO We can only do this for block machines. - // // For dynamic machines, fixed columns are "known but symbolic" - // let v = self.fixed_data.fixed_cols[&r.poly_id].values_max_size()[row as usize]; - // EvalResult::from_number(v) + self.reference_evaluator.evaluate_fixed(r, offset)?.into() } else { let cell = Cell::from_reference(r, offset); // If a cell is known and has a compile-time constant value, @@ -221,8 +211,13 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { } } } - Expression::PublicReference(_) => return None, // TODO - Expression::Challenge(_) => return None, // TODO + Expression::PublicReference(public) => { + self.reference_evaluator.evaluate_public(public)?.into() + } + Expression::Challenge(challenge) => self + .reference_evaluator + .evaluate_challenge(challenge)? + .into(), Expression::Number(n) => (*n).into(), Expression::BinaryOperation(op) => self.evaulate_binary_operation(op, offset)?, Expression::UnaryOperation(op) => self.evaluate_unary_operation(op, offset)?, @@ -281,12 +276,29 @@ impl<'a, T: FieldElement> WitgenInference<'a, T> { } } +pub trait ReferenceEvaluator { + fn evaluate_fixed(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { + None + } + fn evaluate_challenge(&self, _challenge: &Challenge) -> Option { + None + } + fn evaluate_public(&self, _public: &String) -> Option { + None + } +} + #[cfg(test)] mod test { + use std::{fs, str::from_utf8}; + use powdr_ast::analyzed::Analyzed; use powdr_number::GoldilocksField; - use crate::witgen::{jit::affine_symbolic_expression::Assertion, FixedData}; + use crate::{ + constant_evaluator, + witgen::{jit::affine_symbolic_expression::Assertion, FixedData}, + }; use super::*; @@ -323,10 +335,25 @@ mod test { .join("\n") } + struct ReferenceEvaluatorForFixedData<'a>(&'a FixedData<'a, GoldilocksField>); + impl<'a> ReferenceEvaluator for ReferenceEvaluatorForFixedData<'a> { + fn evaluate_fixed( + &self, + var: &AlgebraicReference, + row_offset: i32, + ) -> Option { + assert!(var.is_fixed()); + let values = self.0.fixed_cols[&var.poly_id].values_max_size(); + let row = (row_offset as usize + var.next as usize) % values.len(); + Some(values[row]) + } + } + fn solve_on_rows(input: &str, rows: &[i32], known_cells: Vec<(&str, i32)>) -> String { let analyzed: Analyzed = powdr_pil_analyzer::analyze_string(input).unwrap(); - let fixed_data = FixedData::new(&analyzed, &[], &[], Default::default(), 0); + let fixed_col_vals = constant_evaluator::generate(&analyzed); + let fixed_data = FixedData::new(&analyzed, &fixed_col_vals, &[], Default::default(), 0); let known_cells = known_cells.iter().map(|(name, row_offset)| { let id = fixed_data.try_column_by_name(name).unwrap().id; Cell { @@ -335,7 +362,9 @@ mod test { row_offset: *row_offset, } }); - let mut witgen = WitgenInference::new(&fixed_data, known_cells); + + let ref_eval = ReferenceEvaluatorForFixedData(&fixed_data); + let mut witgen = WitgenInference::new(&fixed_data, ref_eval, known_cells); let mut complete = HashSet::new(); while complete.len() != analyzed.identities.len() * rows.len() { for row in rows { @@ -365,4 +394,33 @@ mod test { "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);\nX[2] = Y[1];\nY[2] = (X[1] + Y[1]);" ); } + + #[test] + fn fib_with_fixed() { + let input = " + namespace Fib(8); + col fixed FIRST = [1] + [0]*; + let x; + let y; + FIRST * (y - 1) = 0; + FIRST * (x - 1) = 0; + // This works in this test because we do not implement wrapping properly in this test. + x' - y = 0; + y' - (x + y) = 0; + "; + let code = solve_on_rows(&input, &[0, 1, 2, 3], vec![]); + assert_eq!( + code, + "Fib::y[0] = 1; +Fib::x[0] = 1; +Fib::x[1] = 1; +Fib::y[1] = 2; +Fib::x[2] = 2; +Fib::y[2] = 3; +Fib::x[3] = 3; +Fib::y[3] = 5; +Fib::x[4] = 5; +Fib::y[4] = 8;" + ); + } } From f7bc144beb6c8beb86be928d517fe2bdb7a78d39 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 11:37:33 +0000 Subject: [PATCH 36/44] Support permutations. --- .../witgen/jit/affine_symbolic_expression.rs | 6 +-- executor/src/witgen/jit/witgen_inference.rs | 38 +++++++++++-------- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 1c1e41a081..31f43d0b8f 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -20,8 +20,8 @@ pub enum Effect { RangeConstraint(V, RangeConstraint), /// A run-time assertion. If this fails, we have conflicting constraints. Assertion(Assertion), - /// a lookup / call to a different machine. - Lookup(u64, Vec>), + /// a call to a different machine. + MachineCall(u64, Vec>), } /// A run-time assertion. If this fails, we have conflicting constraints. @@ -59,7 +59,7 @@ impl Assertion { } } -pub enum LookupArgument { +pub enum MachineCallArgument { Known(SymbolicExpression), Unknown(AffineSymbolicExpression), } diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index eb13749bd7..ae7f9f01b5 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -4,13 +4,13 @@ use itertools::Itertools; use powdr_ast::analyzed::{ AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Challenge, Identity, - LookupIdentity, PhantomLookupIdentity, PolyID, PolynomialIdentity, PolynomialType, - SelectedExpressions, + LookupIdentity, PermutationIdentity, PhantomBusInteractionIdentity, PhantomLookupIdentity, + PhantomPermutationIdentity, PolyID, PolynomialIdentity, PolynomialType, SelectedExpressions, }; use powdr_number::FieldElement; use crate::witgen::{ - global_constraints::RangeConstraintSet, jit::affine_symbolic_expression::LookupArgument, + global_constraints::RangeConstraintSet, jit::affine_symbolic_expression::MachineCallArgument, }; use super::{ @@ -63,13 +63,21 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, Identity::Lookup(LookupIdentity { id, left, right, .. }) + | Identity::Permutation(PermutationIdentity { + id, left, right, .. + }) + | Identity::PhantomPermutation(PhantomPermutationIdentity { + id, left, right, .. + }) | Identity::PhantomLookup(PhantomLookupIdentity { id, left, right, .. }) => self.process_lookup(*id, left, right, row_offset), - _ => { - // TODO + Identity::PhantomBusInteraction(_) => { + // TODO Once we have a concept of "can_be_answered", bus interactions + // should be as easy as lookups. ProcessResult::empty() } + Identity::Connect(_) => ProcessResult::empty(), }; self.ingest_effects(result.effects); result.complete @@ -125,14 +133,14 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, .filter(|e| e.try_to_known().is_none()) .collect_vec(); if unknown.len() == 1 && unknown[0].single_unknown_variable().is_some() { - let effects = vec![Effect::Lookup( + let effects = vec![Effect::MachineCall( lookup_id, lhs.into_iter() .map(|e| { if let Some(val) = e.try_to_known() { - LookupArgument::Known(val.clone()) + MachineCallArgument::Known(val.clone()) } else { - LookupArgument::Unknown(e) + MachineCallArgument::Unknown(e) } }) .collect(), @@ -160,9 +168,9 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, Effect::RangeConstraint(cell, rc) => { self.add_range_constraint(cell.clone(), rc.clone()); } - Effect::Lookup(_, arguments) => { + Effect::MachineCall(_, arguments) => { for arg in arguments { - if let LookupArgument::Unknown(expr) = arg { + if let MachineCallArgument::Unknown(expr) = arg { let cell = expr.single_unknown_variable().unwrap(); self.known_cells.insert(cell.clone()); } @@ -219,12 +227,12 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, .evaluate_challenge(challenge)? .into(), Expression::Number(n) => (*n).into(), - Expression::BinaryOperation(op) => self.evaulate_binary_operation(op, offset)?, + Expression::BinaryOperation(op) => self.evaluate_binary_operation(op, offset)?, Expression::UnaryOperation(op) => self.evaluate_unary_operation(op, offset)?, }) } - fn evaulate_binary_operation( + fn evaluate_binary_operation( &self, op: &AlgebraicBinaryOperation, offset: i32, @@ -317,13 +325,13 @@ mod test { if *expected_equal { "==" } else { "!=" } ) } - Effect::Lookup(id, args) => { + Effect::MachineCall(id, args) => { format!( "lookup({id}, [{}]);", args.iter() .map(|arg| match arg { - LookupArgument::Known(k) => format!("Known({k})"), - LookupArgument::Unknown(u) => format!("Unknown({u})"), + MachineCallArgument::Known(k) => format!("Known({k})"), + MachineCallArgument::Unknown(u) => format!("Unknown({u})"), }) .join(", ") ) From 6ab422e5b561ef227db00d82d03053153822b3d5 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 14:22:47 +0000 Subject: [PATCH 37/44] Fix a bug and add another test. --- .../witgen/jit/affine_symbolic_expression.rs | 12 +-- executor/src/witgen/jit/witgen_inference.rs | 83 ++++++++++++++++++- 2 files changed, 87 insertions(+), 8 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 31f43d0b8f..3a1ac32274 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -286,13 +286,13 @@ impl AffineSymbolicExpression { return ProcessResult::empty(); } - // We need to assert that the masks cover the offset, + // We need to assert that the masks cover "-offset", // otherwise the equation is not solvable. - // We assert offset & !masks == 0 <=> offset == offset | masks. + // We assert -offset & !masks == 0 <=> -offset == -offset | masks. // We use the latter since we cannot properly bit-negate inside the field. effects.push(Assertion::assert_eq( - self.offset.clone(), - &self.offset | &T::from(covered_bits).into(), + -&self.offset, + -&self.offset | T::from(covered_bits).into(), )); ProcessResult::complete(effects) @@ -429,6 +429,8 @@ impl Mul<&SymbolicExpression> #[cfg(test)] mod test { + use pretty_assertions::assert_eq; + use powdr_number::GoldilocksField; use super::*; @@ -562,7 +564,7 @@ mod test { "a = ((-(10 + Z) & 65280) // 256); b = ((-(10 + Z) & 16711680) // 65536); c = ((-(10 + Z) & 4278190080) // 16777216); -assert (10 + Z) == ((10 + Z) | 4294967040); +assert -(10 + Z) == (-(10 + Z) | 4294967040); " ); } diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index ae7f9f01b5..20b427029e 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -300,12 +300,14 @@ pub trait ReferenceEvaluator { mod test { use std::{fs, str::from_utf8}; + use pretty_assertions::assert_eq; + use powdr_ast::analyzed::Analyzed; use powdr_number::GoldilocksField; use crate::{ constant_evaluator, - witgen::{jit::affine_symbolic_expression::Assertion, FixedData}, + witgen::{global_constraints, jit::affine_symbolic_expression::Assertion, FixedData}, }; use super::*; @@ -362,6 +364,8 @@ mod test { powdr_pil_analyzer::analyze_string(input).unwrap(); let fixed_col_vals = constant_evaluator::generate(&analyzed); let fixed_data = FixedData::new(&analyzed, &fixed_col_vals, &[], Default::default(), 0); + let (fixed_data, retained_identities) = + global_constraints::set_global_constraints(fixed_data, &analyzed.identities); let known_cells = known_cells.iter().map(|(name, row_offset)| { let id = fixed_data.try_column_by_name(name).unwrap().id; Cell { @@ -374,14 +378,17 @@ mod test { let ref_eval = ReferenceEvaluatorForFixedData(&fixed_data); let mut witgen = WitgenInference::new(&fixed_data, ref_eval, known_cells); let mut complete = HashSet::new(); - while complete.len() != analyzed.identities.len() * rows.len() { + let mut counter = 0; + while complete.len() != retained_identities.len() * rows.len() { + counter += 1; for row in rows { - for id in analyzed.identities.iter() { + for id in retained_identities.iter() { if !complete.contains(&(id.id(), *row)) && witgen.process_identity(id, *row) { complete.insert((id.id(), *row)); } } } + assert!(counter < 10000, "Solving took more than 10000 rounds."); } format_code(&witgen.code()) } @@ -431,4 +438,74 @@ Fib::x[4] = 5; Fib::y[4] = 8;" ); } + + #[test] + fn xor() { + let input = " +namespace Xor(256 * 256); + let latch: col = |i| { if (i % 4) == 3 { 1 } else { 0 } }; + let FACTOR: col = |i| { 1 << (((i + 1) % 4) * 8) }; + + let a: int -> int = |i| i % 256; + let b: int -> int = |i| (i / 256) % 256; + let P_A: col = a; + let P_B: col = b; + let P_C: col = |i| a(i) ^ b(i); + + let A_byte; + let B_byte; + let C_byte; + + [ A_byte, B_byte, C_byte ] in [ P_A, P_B, P_C ]; + + let A; + let B; + let C; + + A' = A * (1 - latch) + A_byte * FACTOR; + B' = B * (1 - latch) + B_byte * FACTOR; + C' = C * (1 - latch) + C_byte * FACTOR; +"; + let code = solve_on_rows( + &input, + // Use the second block to avoid wrap-around. + &[3, 4, 5, 6, 7], + vec![ + ("Xor::A", 7), + ("Xor::C", 7), // We solve it in reverse, just for fun. + ], + ); + assert_eq!( + code, + "\ +Xor::A_byte[6] = ((Xor::A[7] & 4278190080) // 16777216); +Xor::A[6] = (Xor::A[7] & 16777215); +assert Xor::A[7] == (Xor::A[7] | 4294967295); +Xor::C_byte[6] = ((Xor::C[7] & 4278190080) // 16777216); +Xor::C[6] = (Xor::C[7] & 16777215); +assert Xor::C[7] == (Xor::C[7] | 4294967295); +Xor::A_byte[5] = ((Xor::A[6] & 16711680) // 65536); +Xor::A[5] = (Xor::A[6] & 65535); +assert Xor::A[6] == (Xor::A[6] | 16777215); +Xor::C_byte[5] = ((Xor::C[6] & 16711680) // 65536); +Xor::C[5] = (Xor::C[6] & 65535); +assert Xor::C[6] == (Xor::C[6] | 16777215); +lookup(0, [Known(Xor::A_byte[6]), Unknown(Xor::B_byte[6]), Known(Xor::C_byte[6])]); +Xor::A_byte[4] = (65280 // 256); +Xor::A[4] = 255; +assert 65535 == 65535; +Xor::C_byte[4] = (65280 // 256); +Xor::C[4] = 255; +assert 65535 == 65535; +lookup(0, [Known(Xor::A_byte[5]), Unknown(Xor::B_byte[5]), Known(Xor::C_byte[5])]); +Xor::A_byte[3] = 255; +Xor::C_byte[3] = 255; +lookup(0, [Known(Xor::A_byte[4]), Unknown(Xor::B_byte[4]), Known(Xor::C_byte[4])]); +lookup(0, [Known(255), Unknown(Xor::B_byte[3]), Known(255)]); +Xor::B[4] = Xor::B_byte[3]; +Xor::B[5] = (Xor::B[4] + (Xor::B_byte[4] * 256)); +Xor::B[6] = (Xor::B[5] + (Xor::B_byte[5] * 65536)); +Xor::B[7] = (Xor::B[6] + (Xor::B_byte[6] * 16777216));" + ); + } } From 111902fb26f76c2b76dd59e649964a82f7cbe2c8 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 15:25:47 +0100 Subject: [PATCH 38/44] Update executor/src/witgen/jit/witgen_inference.rs Co-authored-by: Georg Wiese --- executor/src/witgen/jit/witgen_inference.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 20b427029e..4b0ce53368 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -73,7 +73,7 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, id, left, right, .. }) => self.process_lookup(*id, left, right, row_offset), Identity::PhantomBusInteraction(_) => { - // TODO Once we have a concept of "can_be_answered", bus interactions + // TODO(bus_interaction) Once we have a concept of "can_be_answered", bus interactions // should be as easy as lookups. ProcessResult::empty() } From 2b73e09ea20413302ccfb22c1ee0b6c78bd555d4 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 14:26:25 +0000 Subject: [PATCH 39/44] review --- executor/src/witgen/jit/affine_symbolic_expression.rs | 8 +++++--- executor/src/witgen/jit/witgen_inference.rs | 8 ++++---- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 3a1ac32274..d4b3c93571 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -185,8 +185,9 @@ impl AffineSymbolicExpression { /// Solves the equation `self = 0` and returns how to compute the solution. /// The solution can contain assignments to multiple variables. - /// If no way to solve the equation has been found, but it still contains - /// unknon variables, returns an empty, incomplete result. + /// If no way to solve the equation (and no way to derive new range + /// constraints) has been found, but it still contains + /// unknown variables, returns an empty, incomplete result. /// If the equation is known to be unsolvable, returns an error. pub fn solve(&self) -> Result, EvalError> { Ok(match self.coefficients.len() { @@ -466,7 +467,8 @@ mod test { #[test] fn solvable_without_vars() { let constr = &from_number(0); - assert!(constr.solve().unwrap().effects.is_empty()); + let result = constr.solve().unwrap(); + assert!(result.complete && result.effects.is_empty()); } #[test] diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 20b427029e..46d407ccbf 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -24,7 +24,7 @@ use super::{ pub struct WitgenInference<'a, T: FieldElement, RefEval: ReferenceEvaluator> { fixed_data: &'a FixedData<'a, T>, reference_evaluator: RefEval, - range_constraints: HashMap>, + derived_range_constraints: HashMap>, known_cells: HashSet, code: Vec>, } @@ -38,7 +38,7 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, Self { fixed_data, reference_evaluator, - range_constraints: Default::default(), + derived_range_constraints: Default::default(), known_cells: known_cells.into_iter().collect(), code: Default::default(), } @@ -193,7 +193,7 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, self.code.push(Effect::Assignment(cell.clone(), v.into())); } } - self.range_constraints.insert(cell.clone(), rc); + self.derived_range_constraints.insert(cell.clone(), rc); } fn evaluate( @@ -278,7 +278,7 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, next: false, }) .iter() - .chain(self.range_constraints.get(&cell)) + .chain(self.derived_range_constraints.get(&cell)) .cloned() .reduce(|gc, rc| gc.conjunction(&rc)) } From f57bdd77f920de8967a0d0c4883ac9bf5c9ed99c Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 14:37:24 +0000 Subject: [PATCH 40/44] Rename reference evaluator. --- executor/src/witgen/jit/witgen_inference.rs | 44 +++++++-------------- 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 10b98fc8fb..134b584c9f 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -21,23 +21,23 @@ use super::{ /// This component can generate code that solves identities. /// It needs a driver that tells it which identities to process on which rows. -pub struct WitgenInference<'a, T: FieldElement, RefEval: ReferenceEvaluator> { +pub struct WitgenInference<'a, T: FieldElement, FixedEval: FixedEvaluator> { fixed_data: &'a FixedData<'a, T>, - reference_evaluator: RefEval, + fixed_evaluator: FixedEval, derived_range_constraints: HashMap>, known_cells: HashSet, code: Vec>, } -impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, RefEval> { +impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, FixedEval> { pub fn new( fixed_data: &'a FixedData<'a, T>, - reference_evaluator: RefEval, + fixed_evaluator: FixedEval, known_cells: impl IntoIterator, ) -> Self { Self { fixed_data, - reference_evaluator, + fixed_evaluator, derived_range_constraints: Default::default(), known_cells: known_cells.into_iter().collect(), code: Default::default(), @@ -204,7 +204,7 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, Some(match expr { Expression::Reference(r) => { if r.is_fixed() { - self.reference_evaluator.evaluate_fixed(r, offset)?.into() + self.fixed_evaluator.evaluate(r, offset)?.into() } else { let cell = Cell::from_reference(r, offset); // If a cell is known and has a compile-time constant value, @@ -219,13 +219,10 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, } } } - Expression::PublicReference(public) => { - self.reference_evaluator.evaluate_public(public)?.into() + Expression::PublicReference(_) | Expression::Challenge(_) => { + // TODO we need to introduce a variable type for those. + return None; } - Expression::Challenge(challenge) => self - .reference_evaluator - .evaluate_challenge(challenge)? - .into(), Expression::Number(n) => (*n).into(), Expression::BinaryOperation(op) => self.evaluate_binary_operation(op, offset)?, Expression::UnaryOperation(op) => self.evaluate_unary_operation(op, offset)?, @@ -284,21 +281,14 @@ impl<'a, T: FieldElement, RefEval: ReferenceEvaluator> WitgenInference<'a, T, } } -pub trait ReferenceEvaluator { - fn evaluate_fixed(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { - None - } - fn evaluate_challenge(&self, _challenge: &Challenge) -> Option { - None - } - fn evaluate_public(&self, _public: &String) -> Option { +pub trait FixedEvaluator { + fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { None } } #[cfg(test)] mod test { - use std::{fs, str::from_utf8}; use pretty_assertions::assert_eq; @@ -345,13 +335,9 @@ mod test { .join("\n") } - struct ReferenceEvaluatorForFixedData<'a>(&'a FixedData<'a, GoldilocksField>); - impl<'a> ReferenceEvaluator for ReferenceEvaluatorForFixedData<'a> { - fn evaluate_fixed( - &self, - var: &AlgebraicReference, - row_offset: i32, - ) -> Option { + struct FixedEvaluatorForFixedData<'a>(&'a FixedData<'a, GoldilocksField>); + impl<'a> FixedEvaluator for FixedEvaluatorForFixedData<'a> { + fn evaluate(&self, var: &AlgebraicReference, row_offset: i32) -> Option { assert!(var.is_fixed()); let values = self.0.fixed_cols[&var.poly_id].values_max_size(); let row = (row_offset as usize + var.next as usize) % values.len(); @@ -375,7 +361,7 @@ mod test { } }); - let ref_eval = ReferenceEvaluatorForFixedData(&fixed_data); + let ref_eval = FixedEvaluatorForFixedData(&fixed_data); let mut witgen = WitgenInference::new(&fixed_data, ref_eval, known_cells); let mut complete = HashSet::new(); let mut counter = 0; From fea920e8f6751ed1b5994de05a95f98b53f7f13e Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 14:57:57 +0000 Subject: [PATCH 41/44] Fix range constraint for bitwise operations. --- executor/src/witgen/jit/symbolic_expression.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/symbolic_expression.rs b/executor/src/witgen/jit/symbolic_expression.rs index 8f3ddd470f..900135c7b3 100644 --- a/executor/src/witgen/jit/symbolic_expression.rs +++ b/executor/src/witgen/jit/symbolic_expression.rs @@ -286,7 +286,7 @@ impl BitAnd for &SymbolicExpression { Rc::new(rhs.clone()), self.range_constraint() .zip(rhs.range_constraint()) - .map(|(a, b)| a.conjunction(&b)), + .map(|(a, b)| RangeConstraint::from_mask(*a.mask() & *b.mask())), ) } } @@ -317,7 +317,9 @@ impl BitOr for &SymbolicExpression { Rc::new(self.clone()), BinaryOperator::BitOr, Rc::new(rhs.clone()), - None, + self.range_constraint() + .zip(rhs.range_constraint()) + .map(|(a, b)| RangeConstraint::from_mask(*a.mask() | *b.mask())), ) } } From fd363b2d15139b004e7051ee4ea8e92a343c5da1 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 15:11:20 +0000 Subject: [PATCH 42/44] Fix xor test. --- executor/src/witgen/jit/witgen_inference.rs | 35 ++++++++++++--------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 134b584c9f..35955bcb91 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -345,7 +345,12 @@ mod test { } } - fn solve_on_rows(input: &str, rows: &[i32], known_cells: Vec<(&str, i32)>) -> String { + fn solve_on_rows( + input: &str, + rows: &[i32], + known_cells: Vec<(&str, i32)>, + expected_complete: Option, + ) -> String { let analyzed: Analyzed = powdr_pil_analyzer::analyze_string(input).unwrap(); let fixed_col_vals = constant_evaluator::generate(&analyzed); @@ -365,7 +370,8 @@ mod test { let mut witgen = WitgenInference::new(&fixed_data, ref_eval, known_cells); let mut complete = HashSet::new(); let mut counter = 0; - while complete.len() != retained_identities.len() * rows.len() { + let expected_complete = expected_complete.unwrap_or(retained_identities.len() * rows.len()); + while complete.len() != expected_complete { counter += 1; for row in rows { for id in retained_identities.iter() { @@ -382,14 +388,14 @@ mod test { #[test] fn simple_polynomial_solving() { let input = "let X; let Y; let Z; X = 1; Y = X + 1; Z * Y = X + 10;"; - let code = solve_on_rows(input, &[0], vec![]); + let code = solve_on_rows(input, &[0], vec![], None); assert_eq!(code, "X[0] = 1;\nY[0] = 2;\nZ[0] = -9223372034707292155;"); } #[test] fn fib() { let input = "let X; let Y; X' = Y; Y' = X + Y;"; - let code = solve_on_rows(input, &[0, 1], vec![("X", 0), ("Y", 0)]); + let code = solve_on_rows(input, &[0, 1], vec![("X", 0), ("Y", 0)], None); assert_eq!( code, "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);\nX[2] = Y[1];\nY[2] = (X[1] + Y[1]);" @@ -409,7 +415,7 @@ mod test { x' - y = 0; y' - (x + y) = 0; "; - let code = solve_on_rows(&input, &[0, 1, 2, 3], vec![]); + let code = solve_on_rows(&input, &[0, 1, 2, 3], vec![], None); assert_eq!( code, "Fib::y[0] = 1; @@ -460,6 +466,7 @@ namespace Xor(256 * 256); ("Xor::A", 7), ("Xor::C", 7), // We solve it in reverse, just for fun. ], + Some(16), ); assert_eq!( code, @@ -477,17 +484,17 @@ Xor::C_byte[5] = ((Xor::C[6] & 16711680) // 65536); Xor::C[5] = (Xor::C[6] & 65535); assert Xor::C[6] == (Xor::C[6] | 16777215); lookup(0, [Known(Xor::A_byte[6]), Unknown(Xor::B_byte[6]), Known(Xor::C_byte[6])]); -Xor::A_byte[4] = (65280 // 256); -Xor::A[4] = 255; -assert 65535 == 65535; -Xor::C_byte[4] = (65280 // 256); -Xor::C[4] = 255; -assert 65535 == 65535; +Xor::A_byte[4] = ((Xor::A[5] & 65280) // 256); +Xor::A[4] = (Xor::A[5] & 255); +assert Xor::A[5] == (Xor::A[5] | 65535); +Xor::C_byte[4] = ((Xor::C[5] & 65280) // 256); +Xor::C[4] = (Xor::C[5] & 255); +assert Xor::C[5] == (Xor::C[5] | 65535); lookup(0, [Known(Xor::A_byte[5]), Unknown(Xor::B_byte[5]), Known(Xor::C_byte[5])]); -Xor::A_byte[3] = 255; -Xor::C_byte[3] = 255; +Xor::A_byte[3] = Xor::A[4]; +Xor::C_byte[3] = Xor::C[4]; lookup(0, [Known(Xor::A_byte[4]), Unknown(Xor::B_byte[4]), Known(Xor::C_byte[4])]); -lookup(0, [Known(255), Unknown(Xor::B_byte[3]), Known(255)]); +lookup(0, [Known(Xor::A_byte[3]), Unknown(Xor::B_byte[3]), Known(Xor::C_byte[3])]); Xor::B[4] = Xor::B_byte[3]; Xor::B[5] = (Xor::B[4] + (Xor::B_byte[4] * 256)); Xor::B[6] = (Xor::B[5] + (Xor::B_byte[5] * 65536)); From b715662dddcc895980b39a041b905adc9536ff18 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 15:13:46 +0000 Subject: [PATCH 43/44] Remove (for now) unused function. --- executor/src/witgen/jit/witgen_inference.rs | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 35955bcb91..9926ce53b9 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -3,9 +3,9 @@ use std::collections::{HashMap, HashSet}; use itertools::Itertools; use powdr_ast::analyzed::{ AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, - AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Challenge, Identity, - LookupIdentity, PermutationIdentity, PhantomBusInteractionIdentity, PhantomLookupIdentity, - PhantomPermutationIdentity, PolyID, PolynomialIdentity, PolynomialType, SelectedExpressions, + AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, LookupIdentity, + PermutationIdentity, PhantomLookupIdentity, PhantomPermutationIdentity, PolyID, + PolynomialIdentity, PolynomialType, SelectedExpressions, }; use powdr_number::FieldElement; @@ -44,10 +44,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F } } - pub fn known_cells(&self) -> &HashSet { - &self.known_cells - } - pub fn code(self) -> Vec> { self.code } @@ -415,7 +411,7 @@ mod test { x' - y = 0; y' - (x + y) = 0; "; - let code = solve_on_rows(&input, &[0, 1, 2, 3], vec![], None); + let code = solve_on_rows(input, &[0, 1, 2, 3], vec![], None); assert_eq!( code, "Fib::y[0] = 1; @@ -459,7 +455,7 @@ namespace Xor(256 * 256); C' = C * (1 - latch) + C_byte * FACTOR; "; let code = solve_on_rows( - &input, + input, // Use the second block to avoid wrap-around. &[3, 4, 5, 6, 7], vec![ From a905288b7ff3480a3b059872dbd8008a401a6592 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 12 Dec 2024 15:39:01 +0000 Subject: [PATCH 44/44] trick clippy. --- executor/src/witgen/jit/witgen_inference.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 9926ce53b9..f921bbcd31 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -1,3 +1,4 @@ +#![allow(unused)] use std::collections::{HashMap, HashSet}; use itertools::Itertools;