Add Structs using PreprocessedColumn Trait

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -5,6 +5,7 @@ use std::simd::Simd;
 use num_traits::{One, Zero};
 
 use crate::core::backend::simd::m31::{PackedM31, N_LANES};
+use crate::core::backend::simd::SimdBackend;
 use crate::core::backend::{Backend, Col, Column};
 use crate::core::fields::m31::{BaseField, M31};
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
@@ -38,6 +39,53 @@ impl Hash for dyn PreprocessedColumnTrait {
     }
 }
 
+/// 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
+    }
+}
+
 // TODO(ilya): Where should this enum be placed?
 // TODO(Gali): Add documentation for the rest of the variants.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]

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/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/state_machine/mod.rs

@@ -1,5 +1,6 @@
 use crate::constraint_framework::relation_tracker::RelationSummary;
 use crate::constraint_framework::Relation;
+use crate::core::poly::BitReversedOrder;
 pub mod components;
 pub mod gen;
 
@@ -12,19 +13,25 @@ use gen::{gen_interaction_trace, gen_trace};
 use itertools::{chain, Itertools};
 
 use crate::constraint_framework::preprocessed_columns::{
-    gen_preprocessed_columns, PreprocessedColumn,
+    gen_preprocessed_columns, IsFirst, PreprocessedColumn,
 };
 use crate::constraint_framework::TraceLocationAllocator;
 use crate::core::backend::simd::m31::LOG_N_LANES;
 use crate::core::backend::simd::SimdBackend;
 use crate::core::channel::Blake2sChannel;
-use crate::core::fields::m31::M31;
+use crate::core::fields::m31::{BaseField, M31};
 use crate::core::fields::qm31::QM31;
 use crate::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig, TreeVec};
-use crate::core::poly::circle::{CanonicCoset, PolyOps};
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
 use crate::core::prover::{prove, verify, VerificationError};
 use crate::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
 
+pub fn gen_is_first_columns(
+    columns: impl Iterator<Item = IsFirst>,
+) -> Vec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
+    columns.map(|col| col.gen_column_simd()).collect()
+}
+
 #[allow(unused)]
 pub fn prove_state_machine(
     log_n_rows: u32,

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
+    }
+}