Get lookup columns.

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

@@ -1,14 +1,15 @@
 use itertools::Itertools;
 use num_traits::Zero;
 
-use super::trace_gen::write_trace_row;
+use super::trace_gen::{write_lookup_column, write_trace_row};
 use crate::core::air::{Air, Component};
 use crate::core::backend::CPUBackend;
 use crate::core::fields::m31::BaseField;
 use crate::core::ColumnVec;
 
 pub const LOG_N_COLUMNS: usize = 8;
 pub const N_COLUMNS: usize = 1 << LOG_N_COLUMNS;
+pub const LOG_N_ROWS: usize = 4;
 
 /// Component that computes fibonacci numbers over [N_COLUMNS] columns.
 pub struct WideFibComponent {
@@ -36,6 +37,21 @@ impl WideFibComponent {
         dst
     }
 
+    pub fn lookup_columns(
+        &self,
+        trace: &[Vec<BaseField>],
+        // TODO(AlonH): Change alpha and z to SecureField.
+        alpha: BaseField,
+        z: BaseField,
+    ) -> ColumnVec<Vec<BaseField>> {
+        let n_rows = trace[0].len();
+        let zero_vec = vec![BaseField::zero(); n_rows];
+        let mut dst = vec![zero_vec; 2];
+        write_lookup_column(&mut dst[0], trace, 0, alpha, z);
+        write_lookup_column(&mut dst[1], trace, N_COLUMNS - 2, alpha, z);
+        dst
+    }
+
     pub fn log_column_size(&self) -> u32 {
         self.log_n_instances + self.log_fibonacci_size - LOG_N_COLUMNS as u32
     }

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

@@ -7,9 +7,9 @@ pub mod trace_gen;
 #[cfg(test)]
 mod tests {
     use itertools::Itertools;
-    use num_traits::Zero;
+    use num_traits::{One, Zero};
 
-    use super::component::{Input, WideFibAir, WideFibComponent, LOG_N_COLUMNS};
+    use super::component::{Input, WideFibAir, WideFibComponent, LOG_N_COLUMNS, LOG_N_ROWS};
     use crate::commitment_scheme::blake2_hash::Blake2sHasher;
     use crate::commitment_scheme::hasher::Hasher;
     use crate::core::air::accumulation::DomainEvaluationAccumulator;
@@ -34,8 +34,44 @@ mod tests {
         }
     }
 
+    pub fn assert_constraints_on_lookup_column(
+        columns: &[Vec<BaseField>],
+        input_trace: &[Vec<BaseField>],
+        alpha: BaseField,
+        z: BaseField,
+    ) {
+        let n_columns = input_trace.len();
+        let column_length = columns[0].len();
+        let mut column_0_prev_value = BaseField::one();
+        let mut column_1_prev_value = BaseField::one();
+        for i in 0..column_length {
+            assert_eq!(
+                (columns[0][i]
+                    - (input_trace[0][i] + alpha * input_trace[1][i] - z) * column_0_prev_value),
+                BaseField::zero()
+            );
+            assert_eq!(
+                (columns[1][i]
+                    - (input_trace[n_columns - 2][i] + alpha * input_trace[n_columns - 1][i] - z)
+                        * column_1_prev_value),
+                BaseField::zero()
+            );
+            column_0_prev_value = columns[0][i];
+            column_1_prev_value = columns[1][i];
+        }
+
+        assert_eq!(
+            (input_trace[0][0] + alpha * input_trace[1][0] - z) * columns[1][column_length - 1]
+                - (input_trace[n_columns - 2][column_length - 1]
+                    + alpha * input_trace[n_columns - 1][column_length - 1]
+                    - z)
+                    * columns[0][column_length - 1],
+            BaseField::zero()
+        );
+    }
+
     #[test]
-    fn test_wide_fib_trace() {
+    fn test_trace_row_constraints() {
         let wide_fib = WideFibComponent {
             log_fibonacci_size: LOG_N_COLUMNS as u32,
             log_n_instances: 1,
@@ -53,6 +89,25 @@ mod tests {
         assert_constraints_on_row(&trace1);
     }
 
+    #[test]
+    fn test_lookup_column_constraints() {
+        let wide_fib = WideFibComponent {
+            log_fibonacci_size: (LOG_N_ROWS + LOG_N_COLUMNS) as u32,
+            log_n_instances: 0,
+        };
+        let input = Input {
+            a: m31!(1),
+            b: m31!(1),
+        };
+
+        let alpha = m31!(1);
+        let z = m31!(2);
+        let trace = wide_fib.fill_initial_trace(vec![input]);
+        let lookup_trace = wide_fib.lookup_columns(&trace, alpha, z);
+
+        assert_constraints_on_lookup_column(&lookup_trace, &trace, alpha, z)
+    }
+
     #[test]
     fn test_composition_is_low_degree() {
         let wide_fib = WideFibComponent {

crates/prover/src/examples/wide_fibonacci/trace_gen.rs

@@ -1,3 +1,5 @@
+use num_traits::One;
+
 use super::component::{Input, N_COLUMNS};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::FieldExpOps;
@@ -18,3 +20,19 @@ pub fn write_trace_row(
 
     (dst[N_COLUMNS - 2][row_index], dst[N_COLUMNS - 1][row_index])
 }
+
+pub fn write_lookup_column(
+    dst: &mut [BaseField],
+    input_trace: &[Vec<BaseField>],
+    column_offset: usize,
+    alpha: BaseField,
+    z: BaseField,
+) {
+    let mut prev_value = BaseField::one();
+    for (i, cell) in dst.iter_mut().enumerate() {
+        let row_i_0 = input_trace[column_offset][i];
+        let row_i_1 = input_trace[column_offset + 1][i];
+        *cell = (row_i_0 + alpha * row_i_1 - z) * prev_value;
+        prev_value = *cell;
+    }
+}