witness generation for fixed lookups

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<u64> {
     if fixed.first() != Some(&0.into()) {
         return None;

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::<Option<_>>()?;
+
         // If we already know the LHS, skip it.
         if left
             .iter()
@@ -74,62 +84,47 @@ impl FixedLookup {
         &mut self,
         fixed_data: &FixedData,
         left: &[Result<AffineExpression, EvalError>],
-        right: &SelectedExpressions,
+        right: Vec<&&Vec<BigInt>>,
     ) -> 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::<Option<Vec<_>>>()
+            })
+            // deduplicate values
+            .collect::<HashSet<_>>();
 
-        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() {

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

tests/pil.rs

@@ -65,3 +65,8 @@ fn test_witness_lookup() {
         }),
     );
 }
+
+#[test]
+fn test_pair_lookup() {
+    verify_pil("pair_lookup.pil", None);
+}

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};