From dc1d681e5f1ce6ee0b79029dc6d49661a0fc0cdb Mon Sep 17 00:00:00 2001 From: chriseth Date: Wed, 29 Mar 2023 18:33:08 +0200 Subject: [PATCH] witness generation for fixed lookups --- src/witness_generator/bit_constraints.rs | 16 +-- .../machines/fixed_lookup_machine.rs | 103 +++++++++--------- src/witness_generator/util.rs | 20 +++- tests/pil.rs | 5 + tests/pil_data/pair_lookup.pil | 13 +++ 5 files changed, 88 insertions(+), 69 deletions(-) create mode 100644 tests/pil_data/pair_lookup.pil diff --git a/src/witness_generator/bit_constraints.rs b/src/witness_generator/bit_constraints.rs index c6e65012c..2a9d46e6e 100644 --- a/src/witness_generator/bit_constraints.rs +++ b/src/witness_generator/bit_constraints.rs @@ -1,12 +1,13 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt::{Display, Formatter}; -use crate::analyzer::{BinaryOperator, Expression, Identity, IdentityKind, PolynomialReference}; +use crate::analyzer::{BinaryOperator, Expression, Identity, IdentityKind}; use crate::number::{AbstractNumberType, GOLDILOCKS_MOD}; use crate::witness_generator::util::{contains_next_ref, WitnessColumnNamer}; use super::expression_evaluator::ExpressionEvaluator; use super::symbolic_evaluator::SymbolicEvaluator; +use super::util::is_simple_poly; use super::{Constraint, FixedData}; /// Constraint on the bit values of a variable X. @@ -320,19 +321,6 @@ fn try_transfer_constraints<'a>( }) } -fn is_simple_poly(expr: &Expression) -> Option<&str> { - if let Expression::PolynomialReference(PolynomialReference { - name, - index: None, - next: false, - }) = expr - { - Some(name) - } else { - None - } -} - fn smallest_period_candidate(fixed: &[AbstractNumberType]) -> Option { if fixed.first() != Some(&0.into()) { return None; diff --git a/src/witness_generator/machines/fixed_lookup_machine.rs b/src/witness_generator/machines/fixed_lookup_machine.rs index 81834f3bd..e35cdb948 100644 --- a/src/witness_generator/machines/fixed_lookup_machine.rs +++ b/src/witness_generator/machines/fixed_lookup_machine.rs @@ -1,14 +1,15 @@ use std::collections::{HashMap, HashSet}; -use super::Machine; -use crate::analyzer::{Expression, Identity, IdentityKind, SelectedExpressions}; -use crate::number::{AbstractNumberType, DegreeType}; +use num_bigint::BigInt; +use super::Machine; +use crate::analyzer::{Identity, IdentityKind, SelectedExpressions}; +use crate::number::AbstractNumberType; + +use crate::witness_generator::util::is_simple_poly; use crate::witness_generator::{ affine_expression::AffineExpression, eval_error::{self, EvalError}, - expression_evaluator::ExpressionEvaluator, - fixed_evaluator::FixedEvaluator, util::contains_witness_ref, EvalResult, FixedData, }; @@ -50,6 +51,15 @@ impl Machine for FixedLookup { return None; } + // get the values of the fixed columns + let right = right + .expressions + .iter() + .map(|right_key| { + is_simple_poly(right_key).and_then(|name| fixed_data.fixed_cols.get(name)) + }) + .collect::>()?; + // If we already know the LHS, skip it. if left .iter() @@ -74,62 +84,47 @@ impl FixedLookup { &mut self, fixed_data: &FixedData, left: &[Result], - right: &SelectedExpressions, + right: Vec<&&Vec>, ) -> EvalResult { - let left_key = match left[0].clone() { - Ok(v) => match v.constant_value() { - Some(v) => Ok(v), - None => Err(format!( - "First expression needs to be constant but is not: {}.", - v.format(fixed_data) - )), - }, - Err(err) => Err(format!("First expression on the LHS is unknown: {err}")), - }?; + let mut matches = (0..fixed_data.degree as usize) + // get all lookup rows which match the lhs + .filter_map(|row| { + left.iter() + .zip(right.iter()) + .map(|(left, right)| { + let right = &right[row]; + match left { + Ok(left) => left + .constant_value() + // if the lhs is constant, it's a match iff the values match + .map(|left| (left == *right).then_some(right)) + // if it's not constant, it's a match + .unwrap_or_else(|| Some(right)), + // if we do not know the lhs, it's a match + Err(_) => Some(right), + } + }) + .collect::>>() + }) + // deduplicate values + .collect::>(); - let right_key = right.expressions.first().unwrap(); - let rhs_row = if let Expression::PolynomialReference(poly) = right_key { - // TODO we really need a search index on this. - fixed_data.fixed_cols - .get(poly.name.as_str()) - .and_then(|values| values.iter().position(|v| *v == left_key)) - .ok_or_else(|| { - format!( - "Unable to find matching row on the RHS where the first element is {left_key} - only fixed columns supported there." - ) - }) - .map(|i| i as DegreeType) - } else { - Err("First item on the RHS must be a polynomial reference.".to_string()) - }?; - - // TODO we only support the following case: - // - The first component on the LHS has to be known - // - The first component on the RHS has to be a direct fixed column reference - // - The first match of those uniquely determines the rest of the RHS. - - // TODO in the other cases, we could at least return some improved bit constraints. - - let rhs_evaluator = - ExpressionEvaluator::new(FixedEvaluator::new(fixed_data, rhs_row as usize)); + let right_values = match matches.len() { + // no match, we error out + 0 => { + return Err(EvalError::Generic("Plookup is not satisfied".to_string())); + } + // a single match, we continue + 1 => matches.drain().next().unwrap(), + // multiple matches, we stop and learnt nothing + _ => return Ok(vec![]), + }; let mut reasons = vec![]; let mut result = vec![]; - for (l, r) in left.iter().zip(right.expressions.iter()).skip(1) { + for (l, r) in left.iter().zip(right_values) { match l { Ok(l) => { - // This needs to be a costant because symbolic variables - // would reference a different row! - let r = rhs_evaluator.evaluate(r).and_then(|r| { - r.constant_value().ok_or_else(|| { - format!("Constant value required: {}", r.format(fixed_data)).into() - }) - }); - if let Err(err) = r { - reasons.push(err); - continue; - } - let r = r.unwrap(); let evaluated = l.clone() - r.clone().into(); // TODO we could use bit constraints here match evaluated.solve() { diff --git a/src/witness_generator/util.rs b/src/witness_generator/util.rs index 4cb24a3c9..c2942f27d 100644 --- a/src/witness_generator/util.rs +++ b/src/witness_generator/util.rs @@ -1,4 +1,4 @@ -use crate::analyzer::Expression; +use crate::analyzer::{Expression, PolynomialReference}; use super::FixedData; @@ -35,6 +35,24 @@ pub fn contains_witness_ref(expr: &Expression, fixed_data: &FixedData) -> bool { }) } +/// Checks if an expression is +/// - a polynomial +/// - not part of a polynomial array +/// - not shifted with `'` +/// and return the polynomial's name if so +pub fn is_simple_poly(expr: &Expression) -> Option<&str> { + if let Expression::PolynomialReference(PolynomialReference { + name, + index: None, + next: false, + }) = expr + { + Some(name) + } else { + None + } +} + pub fn expr_any(expr: &Expression, f: &mut impl FnMut(&Expression) -> bool) -> bool { if f(expr) { true diff --git a/tests/pil.rs b/tests/pil.rs index ee5edc60e..4bd28955d 100644 --- a/tests/pil.rs +++ b/tests/pil.rs @@ -65,3 +65,8 @@ fn test_witness_lookup() { }), ); } + +#[test] +fn test_pair_lookup() { + verify_pil("pair_lookup.pil", None); +} diff --git a/tests/pil_data/pair_lookup.pil b/tests/pil_data/pair_lookup.pil new file mode 100644 index 000000000..11511455d --- /dev/null +++ b/tests/pil_data/pair_lookup.pil @@ -0,0 +1,13 @@ +constant %N = 256; + +namespace Main(%N); + col fixed A(i) { i % 16 }; + col fixed B(i) { (i >> 4) % 16 }; + col fixed C(i) { (A(i) + B(i)) & 0xf }; + + + col fixed a(i) { (i + 13) * 29 & 0xf }; + col fixed b(i) { (i * 17 + 3) & 0xf }; + col witness c; + + {a, b, c} in {Main.A, Main.B, Main.C}; \ No newline at end of file