Add Structs using PreprocessedColumn Trait

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

@@ -16,6 +16,7 @@ use crate::core::fields::m31::BaseField;
 use crate::core::fields::FieldExpOps;
 
 mod air;
+mod preprocessed_columns;
 mod round;
 mod scheduler;
 mod xor_table;

crates/prover/src/examples/blake/preprocessed_columns.rs

@@ -0,0 +1,36 @@
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumnTrait;
+
+// TODO(Gali): Add documentation.
+#[derive(Debug)]
+pub struct XorTable {
+    pub n_bits: u32,
+    pub n_expand_bits: u32,
+    pub index_in_table: usize,
+}
+impl XorTable {
+    //TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    pub const fn new(n_bits: u32, n_expand_bits: u32, index_in_table: usize) -> Self {
+        Self {
+            n_bits,
+            n_expand_bits,
+            index_in_table,
+        }
+    }
+}
+impl PreprocessedColumnTrait for XorTable {
+    fn name(&self) -> &'static str {
+        "preprocessed_xor_table"
+    }
+
+    fn id(&self) -> String {
+        format!(
+            "XorTable(n_bits: {}, n_expand_bits: {}, index_in_table: {})",
+            self.n_bits, self.n_expand_bits, self.index_in_table
+        )
+    }
+
+    fn log_size(&self) -> u32 {
+        2 * (self.n_bits - self.n_expand_bits)
+    }
+}

crates/prover/src/examples/mod.rs

@@ -1,6 +1,7 @@
 pub mod blake;
 pub mod plonk;
 pub mod poseidon;
+pub mod preprocessed_columns;
 pub mod state_machine;
 pub mod wide_fibonacci;
 pub mod xor;

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

@@ -1,3 +1,5 @@
+pub mod preprocessed_columns;
+
 use itertools::Itertools;
 use num_traits::One;
 use tracing::{span, Level};

crates/prover/src/examples/plonk/preprocessed_columns.rs

@@ -0,0 +1,25 @@
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumnTrait;
+
+// TODO(Gali): Add documentation.
+#[derive(Debug)]
+pub struct Plonk {
+    pub wire: usize,
+}
+impl Plonk {
+    pub const fn new(wire: usize) -> Self {
+        Self { wire }
+    }
+}
+impl PreprocessedColumnTrait for Plonk {
+    fn name(&self) -> &'static str {
+        "preprocessed_plonk"
+    }
+
+    fn id(&self) -> String {
+        format!("Plonk(wire: {})", self.wire)
+    }
+
+    fn log_size(&self) -> u32 {
+        todo!("Plonk::log_size")
+    }
+}

crates/prover/src/examples/preprocessed_columns.rs

@@ -0,0 +1,58 @@
+use std::simd::Simd;
+
+use num_traits::{One, Zero};
+
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumnTrait;
+use crate::core::backend::simd::m31::{PackedM31, N_LANES};
+use crate::core::backend::simd::SimdBackend;
+use crate::core::backend::{Col, Column};
+use crate::core::fields::m31::BaseField;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
+
+/// A column with `1` at the first position, and `0` elsewhere.
+#[derive(Debug)]
+pub struct IsFirst {
+    pub log_size: u32,
+}
+impl IsFirst {
+    pub const fn new(log_size: u32) -> Self {
+        Self { log_size }
+    }
+
+    pub fn packed_at(&self, vec_row: usize) -> PackedM31 {
+        assert!(vec_row < (1 << self.log_size) / N_LANES);
+        if vec_row == 0 {
+            unsafe {
+                PackedM31::from_simd_unchecked(Simd::from_array(std::array::from_fn(|i| {
+                    if i == 0 {
+                        1
+                    } else {
+                        0
+                    }
+                })))
+            }
+        } else {
+            PackedM31::zero()
+        }
+    }
+
+    pub fn gen_column_simd(&self) -> CircleEvaluation<SimdBackend, BaseField, BitReversedOrder> {
+        let mut col = Col::<SimdBackend, BaseField>::zeros(1 << self.log_size);
+        col.set(0, BaseField::one());
+        CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
+    }
+}
+impl PreprocessedColumnTrait for IsFirst {
+    fn name(&self) -> &'static str {
+        "preprocessed_is_first"
+    }
+
+    fn id(&self) -> String {
+        format!("IsFirst(log_size: {})", self.log_size)
+    }
+
+    fn log_size(&self) -> u32 {
+        self.log_size
+    }
+}

crates/prover/src/examples/xor/gkr_lookups/mod.rs

@@ -1,2 +1,3 @@
 pub mod accumulation;
 pub mod mle_eval;
+pub mod preprocessed_columns;

crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs

@@ -0,0 +1,59 @@
+use num_traits::One;
+
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumnTrait;
+use crate::core::backend::simd::SimdBackend;
+use crate::core::backend::{Col, Column};
+use crate::core::fields::m31::BaseField;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
+use crate::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index};
+
+/// A column with `1` at every `2^log_step` positions, `0` elsewhere, shifted by offset.
+#[derive(Debug)]
+pub struct IsStepWithOffset {
+    log_size: u32,
+    log_step: u32,
+    offset: usize,
+}
+impl IsStepWithOffset {
+    pub const fn new(log_size: u32, log_step: u32, offset: usize) -> Self {
+        Self {
+            log_size,
+            log_step,
+            offset,
+        }
+    }
+
+    // TODO(andrew): Consider optimizing. Is a quotients of two coset_vanishing (use succinct rep
+    // for verifier).
+    //TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    fn gen_column_simd(&self) -> CircleEvaluation<SimdBackend, BaseField, BitReversedOrder> {
+        let mut col = Col::<SimdBackend, BaseField>::zeros(1 << self.log_size);
+        let size = 1 << self.log_size;
+        let step = 1 << self.log_step;
+        let step_offset = self.offset % step;
+        for i in (step_offset..size).step_by(step) {
+            let circle_domain_index = coset_index_to_circle_domain_index(i, self.log_size);
+            let circle_domain_index_bit_rev = bit_reverse_index(circle_domain_index, self.log_size);
+            col.set(circle_domain_index_bit_rev, BaseField::one());
+        }
+        CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
+    }
+}
+impl PreprocessedColumnTrait for IsStepWithOffset {
+    fn name(&self) -> &'static str {
+        "preprocessed_is_step_with_offset"
+    }
+
+    fn id(&self) -> String {
+        format!(
+            "IsStepWithOffset(log_size: {}, log_step: {}, offset: {})",
+            self.log_size, self.log_step, self.offset
+        )
+    }
+
+    fn log_size(&self) -> u32 {
+        self.log_size
+    }
+}