Add Structs using PreprocessedColumn Trait

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -3,6 +3,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};
@@ -15,6 +16,44 @@ const SIMD_ENUMERATION_0: PackedM31 = unsafe {
     ]))
 };
 
+/// A column with `1` at the first position, and `0` elsewhere.
+#[derive(Debug, Clone)]
+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)
+    }
+
+    pub fn id(&self) -> String {
+        format!("preprocessed_is_first_{}", self.log_size).to_string()
+    }
+}
+
 // TODO(ilya): Where should this enum be placed?
 // TODO(Gali): Consider making it a trait, 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,28 @@
+// TODO(Gali): Add documentation and remove allow dead code.
+#[allow(dead_code)]
+#[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,
+        }
+    }
+
+    // TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    pub fn id(&self) -> String {
+        format!(
+            "preprocessed_xor_table_{}_{}_{}",
+            self.n_bits, self.n_expand_bits, self.index_in_table
+        )
+    }
+}

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,14 @@
+// TODO(Gali): Add documentation.
+#[derive(Debug)]
+pub struct Plonk {
+    pub wire: usize,
+}
+impl Plonk {
+    pub const fn new(wire: usize) -> Self {
+        Self { wire }
+    }
+
+    pub fn id(&self) -> String {
+        format!("preprocessed_plonk_{}", self.wire)
+    }
+}

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,49 @@
+use num_traits::One;
+
+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)]
+    pub 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)
+    }
+
+    pub fn id(&self) -> String {
+        format!(
+            "preprocessed_is_step_with_offset_{}_{}_{}",
+            self.log_size, self.log_step, self.offset
+        )
+    }
+}