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, ComponentVec, InteractionElements};
+use crate::core::{ColumnVec, ComponentVec, InteractionElements, LookupValues};
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -58,6 +58,7 @@ pub trait AirExt: Air {
         mask_values: &ComponentVec<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.0).for_each(|(component, mask)| {
@@ -66,6 +67,7 @@ pub trait AirExt: Air {
                 mask,
                 &mut evaluation_accumulator,
                 interaction_elements,
+                lookup_values,
             )
         });
         evaluation_accumulator.finalize()
@@ -130,6 +132,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()
@@ -146,10 +149,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;
@@ -74,6 +74,7 @@ pub trait Component {
         mask: &ColumnVec<Vec<SecureField>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
         interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     );
 }
 
@@ -93,7 +94,13 @@ 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 {
+        LookupValues::default()
+    }
 }
 
 /// A component trace is a set of polynomials for each column on that component.

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

@@ -406,6 +406,17 @@ macro_rules! impl_extension_field {
             }
         }
 
+        impl TryInto<M31> for $field_name {
+            type Error = ();
+
+            fn try_into(self) -> Result<M31, Self::Error> {
+                if self.1 != <$extended_field_name>::zero() {
+                    return Err(());
+                }
+                self.0.try_into().map_err(|_| ())
+            }
+        }
+
         impl AddAssign<M31> for $field_name {
             fn add_assign(&mut self, rhs: M31) {
                 *self = *self + rhs;

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};
@@ -39,6 +39,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,
 }
 
@@ -92,14 +93,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();
 
@@ -128,13 +139,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,
     })
 }
@@ -199,6 +212,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.
@@ -229,6 +250,7 @@ pub fn verify(
             &trace_oods_values,
             random_coeff,
             &interaction_elements,
+            &proof.lookup_values,
         )
     {
         return Err(VerificationError::OodsNotMatching);
@@ -353,7 +375,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>> {
@@ -431,6 +453,7 @@ mod tests {
             _mask: &crate::core::ColumnVec<Vec<SecureField>>,
             evaluation_accumulator: &mut PointEvaluationAccumulator,
             _interaction_elements: &InteractionElements,
+            _lookup_values: &LookupValues,
         ) {
             evaluation_accumulator.accumulate(qm31!(0, 0, 0, 1))
         }
@@ -452,6 +475,7 @@ mod tests {
             _trace: &ComponentTrace<'_, CpuBackend>,
             _evaluation_accumulator: &mut DomainEvaluationAccumulator<CpuBackend>,
             _interaction_elements: &InteractionElements,
+            _lookup_values: &LookupValues,
         ) {
             // Does nothing.
         }

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,
@@ -113,6 +113,7 @@ impl Component for FibonacciComponent {
         mask: &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][..].try_into().unwrap()),
@@ -142,6 +143,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);

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,
             &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: &ColumnVec<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();