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,86 @@
+use tracing::{span, Level};
+
+use crate::constraint_framework::preprocessed_columns::gen_is_first;
+use crate::core::backend::simd::column::BaseColumn;
+use crate::core::backend::simd::SimdBackend;
+use crate::core::fields::m31::BaseField;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
+use crate::core::ColumnVec;
+
+// 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
+        )
+    }
+
+    // TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    pub const fn limb_bits(&self) -> u32 {
+        self.n_bits - self.n_expand_bits
+    }
+
+    // TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    pub const fn column_bits(&self) -> u32 {
+        2 * self.limb_bits()
+    }
+
+    /// Generates the Preprocessed trace for the xor table.
+    /// Returns the Preprocessed trace, the Preprocessed trace, and the claimed sum.
+    // TODO(Gali): Remove allow dead code.
+    #[allow(dead_code)]
+    #[allow(clippy::type_complexity)]
+    pub fn generate_constant_trace(
+        &self,
+    ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
+        let limb_bits = self.limb_bits();
+        let _span = span!(Level::INFO, "Xor Preprocessed trace").entered();
+
+        // Generate the constant columns. In reality, these should be generated before the
+        // proof even began.
+        let a_col: BaseColumn = (0..(1 << self.column_bits()))
+            .map(|i| BaseField::from_u32_unchecked((i >> limb_bits) as u32))
+            .collect();
+        let b_col: BaseColumn = (0..(1 << self.column_bits()))
+            .map(|i| BaseField::from_u32_unchecked((i & ((1 << limb_bits) - 1)) as u32))
+            .collect();
+        let c_col: BaseColumn = (0..(1 << self.column_bits()))
+            .map(|i| {
+                BaseField::from_u32_unchecked(
+                    ((i >> limb_bits) ^ (i & ((1 << limb_bits) - 1))) as u32,
+                )
+            })
+            .collect();
+
+        let mut constant_trace = [a_col, b_col, c_col]
+            .map(|x| {
+                CircleEvaluation::new(CanonicCoset::new(self.column_bits()).circle_domain(), x)
+            })
+            .to_vec();
+        constant_trace.push(gen_is_first(self.column_bits()));
+        constant_trace
+    }
+}

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

@@ -268,6 +268,23 @@ pub fn prove_fibonacci_plonk(
     (component, proof)
 }
 
+/// Preprocessed columns for describing a plonk circuit.
+/// Each plonk gate is described by input wires `a_wire` and `b_wire`, output wire `c_wire`, and
+/// operation `op` in P.  
+#[derive(Debug)]
+pub struct Plonk {
+    pub name: String,
+}
+impl Plonk {
+    pub const fn new(name: String) -> Self {
+        Self { name }
+    }
+
+    pub fn id(&self) -> String {
+        format!("preprocessed_plonk_{}", self.name)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::env;

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