Merge pull request #720 from powdr-labs/use_fixed_constraints

Use fixed constraints

compiler/tests/pil.rs

@@ -239,3 +239,8 @@ fn test_fixed_columns() {
 fn test_witness_via_let() {
     verify_pil("witness_via_let.pil", None);
 }
+
+#[test]
+fn conditional_fixed_constraints() {
+    verify_pil("conditional_fixed_constraints.pil", None);
+}

executor/src/witgen/global_constraints.rs

@@ -41,6 +41,18 @@ pub struct GlobalConstraints<T: FieldElement> {
     pub fixed_constraints: FixedColumnMap<Option<RangeConstraint<T>>>,
 }
 
+impl<T: FieldElement> RangeConstraintSet<&PolynomialReference, T> for GlobalConstraints<T> {
+    fn range_constraint(&self, id: &PolynomialReference) -> Option<RangeConstraint<T>> {
+        assert!(!id.next);
+        let poly_id = id.poly_id.unwrap();
+        match poly_id.ptype {
+            PolynomialType::Constant => self.fixed_constraints[&poly_id].clone(),
+            PolynomialType::Committed => self.witness_constraints[&poly_id].clone(),
+            PolynomialType::Intermediate => None,
+        }
+    }
+}
+
 /// Determines global constraints on witness and fixed columns.
 /// Removes identities that only serve to create range constraints from
 /// the identities vector and returns the remaining identities.

executor/src/witgen/identity_processor.rs

@@ -158,6 +158,7 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'b,
         // query the fixed lookup "machine"
         if let Some(result) = self.mutable_state.fixed_lookup.process_plookup(
             self.fixed_data,
+            rows,
             identity.kind,
             &left,
             &identity.right,

executor/src/witgen/machines/fixed_lookup_machine.rs

@@ -2,12 +2,15 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use std::mem;
 use std::num::NonZeroUsize;
 
-use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference};
+use ast::analyzed::{Expression, IdentityKind, PolyID, PolynomialReference, PolynomialType};
 use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 use number::FieldElement;
 
 use crate::witgen::affine_expression::AffineExpression;
+use crate::witgen::global_constraints::{GlobalConstraints, RangeConstraintSet};
+use crate::witgen::range_constraints::RangeConstraint;
+use crate::witgen::rows::RowPair;
 use crate::witgen::util::try_to_simple_poly_ref;
 use crate::witgen::{EvalError, EvalValue, IncompleteCause};
 use crate::witgen::{EvalResult, FixedData};
@@ -155,27 +158,23 @@ impl<T: FieldElement> IndexedColumns<T> {
 }
 
 /// Machine to perform a lookup in fixed columns only.
-#[derive(Default)]
-pub struct FixedLookup<T> {
+pub struct FixedLookup<T: FieldElement> {
+    global_constraints: GlobalConstraints<T>,
     indices: IndexedColumns<T>,
 }
 
 impl<T: FieldElement> FixedLookup<T> {
-    pub fn try_new(
-        _fixed_data: &FixedData<T>,
-        identities: &[&Identity<Expression<T>>],
-        witness_names: &HashSet<&str>,
-    ) -> Option<Self> {
-        if identities.is_empty() && witness_names.is_empty() {
-            Some(Default::default())
-        } else {
-            None
+    pub fn new(global_constraints: GlobalConstraints<T>) -> Self {
+        Self {
+            global_constraints,
+            indices: Default::default(),
         }
     }
 
     pub fn process_plookup<'b>(
         &mut self,
         fixed_data: &FixedData<T>,
+        rows: &RowPair<'_, '_, T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'b PolynomialReference, T>],
         right: &'b SelectedExpressions<Expression<T>>,
@@ -195,15 +194,24 @@ impl<T: FieldElement> FixedLookup<T> {
             .map(try_to_simple_poly_ref)
             .collect::<Option<Vec<_>>>()?;
 
-        Some(self.process_plookup_internal(fixed_data, left, right))
+        Some(self.process_plookup_internal(fixed_data, rows, left, right))
     }
 
     fn process_plookup_internal<'b>(
         &mut self,
         fixed_data: &FixedData<T>,
+        rows: &RowPair<'_, '_, T>,
         left: &[AffineExpression<&'b PolynomialReference, T>],
-        right: Vec<&PolynomialReference>,
+        right: Vec<&'b PolynomialReference>,
     ) -> EvalResult<'b, T> {
+        if left.len() == 1
+            && !left.first().unwrap().is_constant()
+            && right.first().unwrap().poly_id().ptype == PolynomialType::Constant
+        {
+            // Lookup of the form "c { X } in { B }". Might be a conditional range check.
+            return self.process_range_check(rows, left.first().unwrap(), right.first().unwrap());
+        }
+
         // split the fixed columns depending on whether their associated lookup variable is constant or not. Preserve the value of the constant arguments.
         // {1, 2, x} in {A, B, C} -> [[(A, 1), (B, 2)], [C, x]]
 
@@ -212,8 +220,7 @@ impl<T: FieldElement> FixedLookup<T> {
         let mut output_expressions = vec![];
 
         left.iter().zip(right).for_each(|(l, r)| {
-            let left_value = l.constant_value();
-            if let Some(value) = left_value {
+            if let Some(value) = l.constant_value() {
                 input_assignment.push((r, value));
             } else {
                 output_columns.push(r.poly_id());
@@ -274,4 +281,50 @@ impl<T: FieldElement> FixedLookup<T> {
 
         Ok(result)
     }
+
+    fn process_range_check<'b>(
+        &self,
+        rows: &RowPair<'_, '_, T>,
+        lhs: &AffineExpression<&'b PolynomialReference, T>,
+        rhs: &'b PolynomialReference,
+    ) -> EvalResult<'b, T> {
+        // Use AffineExpression::solve_with_range_constraints to transfer range constraints
+        // from the rhs to the lhs.
+        let equation = lhs.clone() - AffineExpression::from_variable_id(rhs);
+        let range_constraints = UnifiedRangeConstraints {
+            witness_constraints: rows,
+            global_constraints: &self.global_constraints,
+        };
+        let updates = equation.solve_with_range_constraints(&range_constraints)?;
+
+        // Filter out any updates to the fixed columns
+        Ok(EvalValue::incomplete_with_constraints(
+            updates
+                .constraints
+                .into_iter()
+                .filter(|(poly, _)| poly.poly_id().ptype == PolynomialType::Committed),
+            IncompleteCause::NotConcrete,
+        ))
+    }
+}
+
+/// Combines witness constraints on a concrete row with global range constraints
+/// (used for fixed columns).
+/// This is useful in order to transfer range constraints from fixed columns to
+/// witness columns (see [FixedLookup::process_range_check]).
+pub struct UnifiedRangeConstraints<'a, T: FieldElement> {
+    witness_constraints: &'a RowPair<'a, 'a, T>,
+    global_constraints: &'a GlobalConstraints<T>,
+}
+
+impl<T: FieldElement> RangeConstraintSet<&PolynomialReference, T>
+    for UnifiedRangeConstraints<'_, T>
+{
+    fn range_constraint(&self, poly: &PolynomialReference) -> Option<RangeConstraint<T>> {
+        match poly.poly_id().ptype {
+            PolynomialType::Committed => self.witness_constraints.range_constraint(poly),
+            PolynomialType::Constant => self.global_constraints.range_constraint(poly),
+            PolynomialType::Intermediate => unimplemented!(),
+        }
+    }
 }

executor/src/witgen/machines/machine_extractor.rs

@@ -29,7 +29,7 @@ pub fn split_out_machines<'a, T: FieldElement>(
     identities: Vec<&'a Identity<Expression<T>>>,
     global_range_constraints: &GlobalConstraints<T>,
 ) -> ExtractionOutput<'a, T> {
-    let fixed_lookup = FixedLookup::try_new(fixed, &[], &Default::default()).unwrap();
+    let fixed_lookup = FixedLookup::new(global_range_constraints.clone());
 
     let mut machines: Vec<KnownMachine<T>> = vec![];
 

executor/src/witgen/processor.rs

@@ -334,16 +334,16 @@ mod tests {
         let (constants, degree) = generate(&analyzed);
         let fixed_data = FixedData::new(&analyzed, degree, &constants, vec![]);
 
-        // No submachines
-        let mut fixed_lookup = FixedLookup::default();
-        let mut machines = vec![];
-
         // No global range constraints
         let global_range_constraints = GlobalConstraints {
             witness_constraints: fixed_data.witness_map_with(None),
             fixed_constraints: FixedColumnMap::new(None, fixed_data.fixed_cols.len()),
         };
 
+        // No submachines
+        let mut fixed_lookup = FixedLookup::new(global_range_constraints.clone());
+        let mut machines = vec![];
+
         let row_factory = RowFactory::new(&fixed_data, global_range_constraints);
         let data = (0..fixed_data.degree)
             .map(|i| row_factory.fresh_row(i))

test_data/pil/conditional_fixed_constraints.pil

@@ -0,0 +1,31 @@
+constant %N = 1024;
+
+namespace Main(%N);
+    let first = |i| match i { 0 => 1, _ => 0 };
+    let Y;
+    let Z;
+    let bytes = |i| i % 256;
+    let wrap_bit;
+    wrap_bit * (1 - wrap_bit) = 0;
+    let X_b4;
+    let X_b3;
+    let X_b2;
+    let X_b1;
+    { X_b4 } in { bytes };
+    { X_b3 } in { bytes };
+    { X_b2 } in { bytes };
+    { X_b1 } in { bytes };
+
+    let cond;
+
+    cond { Z } in { two_bits };
+    let two_bits = |i| i % 4;
+
+    0 = cond * (Y - (wrap_bit * 2**32 + X_b4 * 0x1000000 + X_b3 * 0x10000 + X_b2 * 0x100 + X_b1 * 4 + Z));
+    let six_bits = |i| i % 2**6;
+    cond { X_b1 } in { six_bits };
+
+    let odd = |i| i % 2;
+    cond = odd;
+
+    Y' = (1 - first') * (Y + 1);