Add lookup final boundary constraints.

crates/prover/src/core/air/air_ext.rs

@@ -10,7 +10,7 @@ use crate::core::pcs::{CommitmentTreeProver, TreeVec};
 use crate::core::poly::circle::SecureCirclePoly;
 use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
 use crate::core::vcs::ops::MerkleOps;
-use crate::core::{ColumnVec, InteractionElements};
+use crate::core::{ColumnVec, InteractionElements, LookupValues};
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -56,6 +56,7 @@ pub trait AirExt: Air {
         mask_values: &Vec<TreeVec<Vec<Vec<SecureField>>>>,
         random_coeff: SecureField,
         interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     ) -> SecureField {
         let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
         zip_eq(self.components(), mask_values).for_each(|(component, mask)| {
@@ -64,6 +65,7 @@ pub trait AirExt: Air {
                 mask,
                 &mut evaluation_accumulator,
                 interaction_elements,
+                lookup_values,
             )
         });
         evaluation_accumulator.finalize()
@@ -133,6 +135,7 @@ pub trait AirProverExt<B: Backend>: AirProver<B> {
         random_coeff: SecureField,
         component_traces: &[ComponentTrace<'_, B>],
         interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     ) -> SecureCirclePoly<B> {
         let total_constraints: usize = self
             .prover_components()
@@ -149,10 +152,18 @@ pub trait AirProverExt<B: Backend>: AirProver<B> {
                 trace,
                 &mut accumulator,
                 interaction_elements,
+                lookup_values,
             )
         });
         accumulator.finalize()
     }
+
+    fn lookup_values(&self, component_traces: &[ComponentTrace<'_, B>]) -> LookupValues {
+        let mut values = LookupValues::default();
+        zip_eq(self.prover_components(), component_traces)
+            .for_each(|(component, trace)| values.extend(component.lookup_values(trace)));
+        values
+    }
 }
 
 impl<B: Backend, A: AirProver<B>> AirProverExt<B> for A {}

crates/prover/src/core/air/mod.rs

@@ -7,7 +7,7 @@ use super::fields::qm31::SecureField;
 use super::pcs::TreeVec;
 use super::poly::circle::{CircleEvaluation, CirclePoly};
 use super::poly::BitReversedOrder;
-use super::{ColumnVec, InteractionElements};
+use super::{ColumnVec, InteractionElements, LookupValues};
 
 pub mod accumulation;
 mod air_ext;
@@ -67,13 +67,14 @@ pub trait Component {
     /// Returns the ids of the interaction elements used by the component.
     fn interaction_element_ids(&self) -> Vec<String>;
 
-    /// Evaluates the constraint quotients combination of the component, given the mask values.
+    /// Evaluates the constraint quotients combination of the component at a point.
     fn evaluate_constraint_quotients_at_point(
         &self,
         point: CirclePoint<SecureField>,
         mask: &TreeVec<ColumnVec<Vec<SecureField>>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
         interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     );
 }
 
@@ -93,7 +94,11 @@ pub trait ComponentProver<B: Backend>: Component {
         trace: &ComponentTrace<'_, B>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<B>,
         interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     );
+
+    /// Returns the values needed to evaluate the components lookup boundary constraints.
+    fn lookup_values(&self, _trace: &ComponentTrace<'_, B>) -> LookupValues;
 }
 
 /// A component trace is a set of polynomials for each column on that component.

crates/prover/src/core/fields/cm31.rs

@@ -52,6 +52,17 @@ impl Mul for CM31 {
     }
 }
 
+impl TryInto<M31> for CM31 {
+    type Error = ();
+
+    fn try_into(self) -> Result<M31, Self::Error> {
+        if self.1 != M31::zero() {
+            return Err(());
+        }
+        Ok(self.0)
+    }
+}
+
 impl FieldExpOps for CM31 {
     fn inverse(&self) -> Self {
         assert!(!self.is_zero(), "0 has no inverse");

crates/prover/src/core/fields/qm31.rs

@@ -66,6 +66,17 @@ impl Mul for QM31 {
     }
 }
 
+impl TryInto<M31> for QM31 {
+    type Error = ();
+
+    fn try_into(self) -> Result<M31, Self::Error> {
+        if self.1 != CM31::zero() {
+            return Err(());
+        }
+        self.0.try_into()
+    }
+}
+
 impl FieldExpOps for QM31 {
     fn inverse(&self) -> Self {
         assert!(!self.is_zero(), "0 has no inverse");

crates/prover/src/core/mod.rs

@@ -1,6 +1,8 @@
 use std::collections::BTreeMap;
 use std::ops::{Deref, DerefMut, Index};
 
+use fields::m31::BaseField;
+
 use self::fields::qm31::SecureField;
 
 pub mod air;
@@ -61,7 +63,7 @@ impl<T> DerefMut for ComponentVec<T> {
     }
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 pub struct InteractionElements(BTreeMap<String, SecureField>);
 
 impl InteractionElements {
@@ -82,3 +84,25 @@ impl Index<&str> for InteractionElements {
         &self.0[index]
     }
 }
+
+#[derive(Default, Debug)]
+pub struct LookupValues(BTreeMap<String, BaseField>);
+
+impl LookupValues {
+    pub fn new(values: BTreeMap<String, BaseField>) -> Self {
+        Self(values)
+    }
+
+    pub fn extend(&mut self, other: Self) {
+        self.0.extend(other.0);
+    }
+}
+
+impl Index<&str> for LookupValues {
+    type Output = BaseField;
+
+    fn index(&self, index: &str) -> &Self::Output {
+        // TODO(AlonH): Return an error if the key is not found.
+        &self.0[index]
+    }
+}

crates/prover/src/core/prover/mod.rs

@@ -10,7 +10,7 @@ use super::pcs::{CommitmentSchemeProof, TreeVec};
 use super::poly::circle::{CanonicCoset, SecureCirclePoly, MAX_CIRCLE_DOMAIN_LOG_SIZE};
 use super::poly::twiddles::TwiddleTree;
 use super::proof_of_work::ProofOfWorkVerificationError;
-use super::{ColumnVec, InteractionElements};
+use super::{ColumnVec, InteractionElements, LookupValues};
 use crate::core::air::{Air, AirExt, AirProverExt};
 use crate::core::backend::CpuBackend;
 use crate::core::channel::{Blake2sChannel, Channel as ChannelTrait};
@@ -38,6 +38,7 @@ pub const N_QUERIES: usize = 3;
 #[derive(Debug)]
 pub struct StarkProof {
     pub commitments: TreeVec<<ChannelHasher as Hasher>::Hash>,
+    pub lookup_values: LookupValues,
     pub commitment_scheme_proof: CommitmentSchemeProof,
 }
 
@@ -91,14 +92,24 @@ pub fn generate_proof<B: Backend + MerkleOps<MerkleHasher>>(
     twiddles: &TwiddleTree<B>,
     commitment_scheme: &mut CommitmentSchemeProver<B>,
 ) -> Result<StarkProof, ProvingError> {
+    let component_traces = air.component_traces(&commitment_scheme.trees);
+    let lookup_values = air.lookup_values(&component_traces);
+    channel.mix_felts(
+        &lookup_values
+            .0
+            .values()
+            .map(|v| SecureField::from(*v))
+            .collect_vec(),
+    );
+
     // Evaluate and commit on composition polynomial.
     let random_coeff = channel.draw_felt();
-
     let span = span!(Level::INFO, "Composition generation").entered();
     let composition_polynomial_poly = air.compute_composition_polynomial(
         random_coeff,
-        &air.component_traces(&commitment_scheme.trees),
+        &component_traces,
         interaction_elements,
+        &lookup_values,
     );
     span.exit();
 
@@ -127,13 +138,15 @@ pub fn generate_proof<B: Backend + MerkleOps<MerkleHasher>>(
             &trace_oods_values,
             random_coeff,
             interaction_elements,
+            &lookup_values,
         )
     {
         return Err(ProvingError::ConstraintsNotSatisfied);
     }
 
     Ok(StarkProof {
         commitments: commitment_scheme.roots(),
+        lookup_values,
         commitment_scheme_proof,
     })
 }
@@ -198,6 +211,14 @@ pub fn verify(
         commitment_scheme.commit(proof.commitments[1], &column_log_sizes[1], channel);
     }
 
+    channel.mix_felts(
+        &proof
+            .lookup_values
+            .0
+            .values()
+            .map(|v| SecureField::from(*v))
+            .collect_vec(),
+    );
     let random_coeff = channel.draw_felt();
 
     // Read composition polynomial commitment.
@@ -228,6 +249,7 @@ pub fn verify(
             &trace_oods_values,
             random_coeff,
             &interaction_elements,
+            &proof.lookup_values,
         )
     {
         return Err(VerificationError::OodsNotMatching);
@@ -332,7 +354,7 @@ mod tests {
     use crate::core::poly::BitReversedOrder;
     use crate::core::prover::{prove, ProvingError};
     use crate::core::test_utils::test_channel;
-    use crate::core::{ColumnVec, InteractionElements};
+    use crate::core::{ColumnVec, InteractionElements, LookupValues};
     use crate::qm31;
 
     struct TestAir<C: ComponentProver<CpuBackend>> {
@@ -410,6 +432,7 @@ mod tests {
             _mask: &TreeVec<Vec<Vec<SecureField>>>,
             evaluation_accumulator: &mut PointEvaluationAccumulator,
             _interaction_elements: &InteractionElements,
+            _lookup_values: &LookupValues,
         ) {
             evaluation_accumulator.accumulate(qm31!(0, 0, 0, 1))
         }
@@ -431,9 +454,14 @@ mod tests {
             _trace: &ComponentTrace<'_, CpuBackend>,
             _evaluation_accumulator: &mut DomainEvaluationAccumulator<CpuBackend>,
             _interaction_elements: &InteractionElements,
+            _lookup_values: &LookupValues,
         ) {
             // Does nothing.
         }
+
+        fn lookup_values(&self, _trace: &ComponentTrace<'_, CpuBackend>) -> LookupValues {
+            LookupValues::default()
+        }
     }
 
     // Ignored because it takes too long and too much memory (in the CI) to run.

crates/prover/src/examples/fibonacci/component.rs

@@ -15,7 +15,7 @@ use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use crate::core::poly::BitReversedOrder;
 use crate::core::utils::bit_reverse_index;
-use crate::core::{ColumnVec, InteractionElements};
+use crate::core::{ColumnVec, InteractionElements, LookupValues};
 
 pub struct FibonacciComponent {
     pub log_size: u32,
@@ -110,6 +110,7 @@ impl Component for FibonacciComponent {
         mask: &TreeVec<ColumnVec<Vec<SecureField>>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
         _interaction_elements: &InteractionElements,
+        _lookup_values: &LookupValues,
     ) {
         evaluation_accumulator.accumulate(
             self.step_constraint_eval_quotient_by_mask(point, &mask[0][0][..].try_into().unwrap()),
@@ -137,6 +138,7 @@ impl ComponentProver<CpuBackend> for FibonacciComponent {
         trace: &ComponentTrace<'_, CpuBackend>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<CpuBackend>,
         _interaction_elements: &InteractionElements,
+        _lookup_values: &LookupValues,
     ) {
         let poly = &trace.polys[0][0];
         let trace_domain = CanonicCoset::new(self.log_size);
@@ -166,4 +168,8 @@ impl ComponentProver<CpuBackend> for FibonacciComponent {
             }
         }
     }
+
+    fn lookup_values(&self, _trace: &ComponentTrace<'_, CpuBackend>) -> LookupValues {
+        LookupValues::default()
+    }
 }

crates/prover/src/examples/fibonacci/mod.rs

@@ -110,7 +110,6 @@ impl MultiFibonacci {
 #[cfg(test)]
 mod tests {
     use std::assert_matches::assert_matches;
-    use std::collections::BTreeMap;
     use std::iter::zip;
 
     use itertools::Itertools;
@@ -129,7 +128,7 @@ mod tests {
     use crate::core::prover::VerificationError;
     use crate::core::queries::Queries;
     use crate::core::utils::bit_reverse;
-    use crate::core::InteractionElements;
+    use crate::core::{InteractionElements, LookupValues};
     use crate::{m31, qm31};
 
     pub fn generate_test_queries(n_queries: usize, trace_length: usize) -> Vec<usize> {
@@ -159,7 +158,8 @@ mod tests {
         let composition_polynomial_poly = fib.air.compute_composition_polynomial(
             random_coeff,
             &component_traces,
-            &InteractionElements::new(BTreeMap::new()),
+            &InteractionElements::default(),
+            &LookupValues::default(),
         );
 
         // Evaluate this polynomial at another point out of the evaluation domain and compare to
@@ -181,7 +181,8 @@ mod tests {
             point,
             &TreeVec::new(vec![mask_values]),
             &mut evaluation_accumulator,
-            &InteractionElements::new(BTreeMap::new()),
+            &InteractionElements::default(),
+            &LookupValues::default(),
         );
         let oods_value = evaluation_accumulator.finalize();
 

crates/prover/src/examples/poseidon/mod.rs

@@ -26,7 +26,7 @@ use crate::core::fields::{FieldExpOps, FieldOps};
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
 use crate::core::poly::BitReversedOrder;
-use crate::core::{ColumnVec, InteractionElements};
+use crate::core::{ColumnVec, InteractionElements, LookupValues};
 
 const N_LOG_INSTANCES_PER_ROW: usize = 3;
 const N_INSTANCES_PER_ROW: usize = 1 << N_LOG_INSTANCES_PER_ROW;
@@ -124,6 +124,7 @@ impl Component for PoseidonComponent {
         mask: &TreeVec<Vec<Vec<SecureField>>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
         _interaction_elements: &InteractionElements,
+        _lookup_values: &LookupValues,
     ) {
         let constraint_zero_domain = CanonicCoset::new(self.log_column_size()).coset;
         let denom = coset_vanishing(constraint_zero_domain, point);
@@ -402,6 +403,7 @@ impl ComponentProver<SimdBackend> for PoseidonComponent {
         trace: &ComponentTrace<'_, SimdBackend>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<SimdBackend>,
         _interaction_elements: &InteractionElements,
+        _lookup_values: &LookupValues,
     ) {
         assert_eq!(trace.polys[0].len(), self.n_columns());
         let eval_domain = CanonicCoset::new(self.log_column_size() + LOG_EXPAND).circle_domain();
@@ -461,6 +463,10 @@ impl ComponentProver<SimdBackend> for PoseidonComponent {
             assert_eq!(evaluator.constraint_index, n_constraints);
         }
     }
+
+    fn lookup_values(&self, _trace: &ComponentTrace<'_, SimdBackend>) -> LookupValues {
+        LookupValues::default()
+    }
 }
 
 #[cfg(test)]