Get lookup columns. (#572)

crates/prover/src/core/utils.rs

@@ -1,7 +1,8 @@
 use std::iter::Peekable;
 
-use num_traits::One;
+use num_traits::{One, Zero};
 
+use super::fields::m31::BaseField;
 use super::fields::qm31::SecureField;
 
 pub trait IteratorMutExt<'a, T: 'a>: Iterator<Item = &'a mut T> {
@@ -86,6 +87,19 @@ pub fn generate_secure_powers(felt: SecureField, n_powers: usize) -> Vec<SecureF
         .collect()
 }
 
+/// Securely combines the given values using the given random alpha and z.
+/// Alpha and z should be secure field elements for soundness.
+pub fn shifted_secure_combination(
+    values: &[BaseField],
+    alpha: BaseField,
+    z: BaseField,
+) -> BaseField {
+    let res = values
+        .iter()
+        .fold(BaseField::zero(), |acc, &value| acc * alpha + value);
+    res - z
+}
+
 #[cfg(test)]
 mod tests {
     use num_traits::One;

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

@@ -7,7 +7,7 @@ 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::constraint_eval::gen_trace;
@@ -22,8 +22,10 @@ mod tests {
     use crate::core::poly::circle::CanonicCoset;
     use crate::core::poly::BitReversedOrder;
     use crate::core::prover::{prove, verify};
+    use crate::core::utils::shifted_secure_combination;
     use crate::core::vcs::blake2_hash::Blake2sHasher;
     use crate::core::vcs::hasher::Hasher;
+    use crate::examples::wide_fibonacci::trace_gen::write_lookup_column;
     use crate::m31;
 
     pub fn assert_constraints_on_row(row: &[BaseField]) {
@@ -35,8 +37,49 @@ mod tests {
         }
     }
 
+    pub fn assert_constraints_on_lookup_column(
+        column: &[BaseField],
+        input_trace: &[Vec<BaseField>],
+        alpha: BaseField,
+        z: BaseField,
+    ) {
+        let n_columns = input_trace.len();
+        let column_length = column.len();
+        assert_eq!(column_length, input_trace[0].len());
+        let mut prev_value = BaseField::one();
+        for (i, cell) in column.iter().enumerate() {
+            assert_eq!(
+                *cell
+                    * shifted_secure_combination(
+                        &[input_trace[n_columns - 2][i], input_trace[n_columns - 1][i]],
+                        alpha,
+                        z,
+                    ),
+                shifted_secure_combination(&[input_trace[0][i], input_trace[1][i]], alpha, z)
+                    * prev_value
+            );
+            prev_value = *cell;
+        }
+
+        // Assert the last cell in the column is equal to the combination of the first two values
+        // divided by the combination of the last two values in the sequence (all other values
+        // should cancel out).
+        assert_eq!(
+            column[column_length - 1]
+                * shifted_secure_combination(
+                    &[
+                        input_trace[n_columns - 2][column_length - 1],
+                        input_trace[n_columns - 1][column_length - 1]
+                    ],
+                    alpha,
+                    z,
+                ),
+            (shifted_secure_combination(&[input_trace[0][0], input_trace[1][0]], alpha, z))
+        );
+    }
+
     #[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,
@@ -54,6 +97,25 @@ mod tests {
         assert_constraints_on_row(&row_1);
     }
 
+    #[test]
+    fn test_lookup_column_constraints() {
+        let wide_fib = WideFibComponent {
+            log_fibonacci_size: 4 + LOG_N_COLUMNS as u32,
+            log_n_instances: 0,
+        };
+        let input = Input {
+            a: m31!(1),
+            b: m31!(1),
+        };
+
+        let alpha = m31!(7);
+        let z = m31!(11);
+        let trace = gen_trace(&wide_fib, vec![input]);
+        let lookup_column = write_lookup_column(&trace, alpha, z);
+
+        assert_constraints_on_lookup_column(&lookup_column, &trace, alpha, z)
+    }
+
     #[test]
     fn test_composition_is_low_degree() {
         let wide_fib = WideFibComponent {

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

@@ -1,6 +1,9 @@
+use num_traits::One;
+
 use super::component::Input;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::FieldExpOps;
+use crate::core::utils::shifted_secure_combination;
 
 /// Given a private input, write the trace row for the wide Fibonacci example to dst. Returns the
 /// last two elements of the row in case the sequence is continued.
@@ -18,3 +21,29 @@ pub fn write_trace_row(
 
     (dst[n_columns - 2][row_index], dst[n_columns - 1][row_index])
 }
+
+pub fn write_lookup_column(
+    input_trace: &[Vec<BaseField>],
+    // TODO(AlonH): Change alpha and z to SecureField.
+    alpha: BaseField,
+    z: BaseField,
+) -> Vec<BaseField> {
+    let n_rows = input_trace[0].len();
+    let n_columns = input_trace.len();
+    let mut prev_value = BaseField::one();
+    (0..n_rows)
+        .map(|i| {
+            let numerator =
+                shifted_secure_combination(&[input_trace[0][i], input_trace[1][i]], alpha, z);
+            let denominator = shifted_secure_combination(
+                &[input_trace[n_columns - 2][i], input_trace[n_columns - 1][i]],
+                alpha,
+                z,
+            );
+            // TODO(AlonH): Use batch inversion.
+            let cell = (numerator / denominator) * prev_value;
+            prev_value = cell;
+            cell
+        })
+        .collect()
+}