Add eq evals constraints

crates/prover/src/constraint_framework/constant_columns.rs

@@ -4,10 +4,34 @@ use crate::core::backend::{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};
 
 /// Generates a column with a single one at the first position, and zeros elsewhere.
 pub fn gen_is_first<B: Backend>(log_size: u32) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
     let mut col = Col::<B, BaseField>::zeros(1 << log_size);
     col.set(0, BaseField::one());
     CircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), col)
 }
+
+/// Generates a column with `1` at every `2^log_step` positions, `0` elsewhere, shifted by offset.
+// TODO(andrew): Consider optimizing. Is a quotients of two coset_vanishing (use succinct rep for
+// verifier).
+pub fn gen_is_step_with_offset<B: Backend>(
+    log_size: u32,
+    log_step: u32,
+    offset: usize,
+) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
+    let mut col = Col::<B, BaseField>::zeros(1 << log_size);
+
+    let size = 1 << log_size;
+    let step = 1 << log_step;
+    let step_offset = offset % step;
+
+    for i in (step_offset..size).step_by(step) {
+        let circle_domain_index = coset_index_to_circle_domain_index(i, log_size);
+        let circle_domain_index_bit_rev = bit_reverse_index(circle_domain_index, log_size);
+        col.set(circle_domain_index_bit_rev, BaseField::one());
+    }
+
+    CircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), col)
+}

crates/prover/src/core/backend/simd/column.rs

@@ -1,4 +1,5 @@
-use std::mem;
+use std::iter::zip;
+use std::{array, mem};
 
 use bytemuck::allocation::cast_vec;
 use bytemuck::{cast_slice, cast_slice_mut, Zeroable};
@@ -215,6 +216,29 @@ pub struct SecureColumn {
     pub length: usize,
 }
 
+impl SecureColumn {
+    // Separates a single column of `PackedSecureField` elements into `SECURE_EXTENSION_DEGREE` many
+    // `PackedBaseField` coordinate columns.
+    pub fn into_secure_column_by_coords(self) -> SecureColumnByCoords<SimdBackend> {
+        if self.len() < N_LANES {
+            return self.to_cpu().into_iter().collect();
+        }
+
+        let length = self.length;
+        let packed_length = self.data.len();
+        let mut columns = array::from_fn(|_| Vec::with_capacity(packed_length));
+
+        for v in self.data {
+            let packed_coords = v.into_packed_m31s();
+            zip(&mut columns, packed_coords).for_each(|(col, packed_coord)| col.push(packed_coord));
+        }
+
+        SecureColumnByCoords {
+            columns: columns.map(|col| BaseColumn { data: col, length }),
+        }
+    }
+}
+
 impl Column<SecureField> for SecureColumn {
     fn zeros(length: usize) -> Self {
         Self {
@@ -276,9 +300,8 @@ impl FromIterator<SecureField> for SecureColumn {
 
 impl FromIterator<PackedSecureField> for SecureColumn {
     fn from_iter<I: IntoIterator<Item = PackedSecureField>>(iter: I) -> Self {
-        let data = (&mut iter.into_iter()).collect_vec();
+        let data = iter.into_iter().collect_vec();
         let length = data.len() * N_LANES;
-
         Self { data, length }
     }
 }

crates/prover/src/core/fields/secure_column.rs

@@ -1,3 +1,6 @@
+use std::array;
+use std::iter::zip;
+
 use super::m31::BaseField;
 use super::qm31::SecureField;
 use super::{ExtensionOf, FieldOps};
@@ -89,13 +92,15 @@ impl<'a> IntoIterator for &'a SecureColumnByCoords<CpuBackend> {
 }
 impl FromIterator<SecureField> for SecureColumnByCoords<CpuBackend> {
     fn from_iter<I: IntoIterator<Item = SecureField>>(iter: I) -> Self {
-        let mut columns = std::array::from_fn(|_| vec![]);
-        for value in iter.into_iter() {
-            let vals = value.to_m31_array();
-            for j in 0..SECURE_EXTENSION_DEGREE {
-                columns[j].push(vals[j]);
-            }
+        let values = iter.into_iter();
+        let (lower_bound, _) = values.size_hint();
+        let mut columns = array::from_fn(|_| Vec::with_capacity(lower_bound));
+
+        for value in values {
+            let coords = value.to_m31_array();
+            zip(&mut columns, coords).for_each(|(col, coord)| col.push(coord));
         }
+
         SecureColumnByCoords { columns }
     }
 }

crates/prover/src/core/utils.rs

@@ -121,11 +121,15 @@ pub(crate) fn coset_order_to_circle_domain_order<F: Field>(values: &[F]) -> Vec<
     circle_domain_order
 }
 
-pub fn coset_order_to_circle_domain_order_index(index: usize, log_size: u32) -> usize {
-    if index & 1 == 0 {
-        index / 2
+/// Converts an index within a [`Coset`] to the corresponding index in a [`CircleDomain`].
+///
+/// [`CircleDomain`]: crate::core::poly::circle::CircleDomain
+/// [`Coset`]: crate::core::circle::Coset
+pub fn coset_index_to_circle_domain_index(coset_index: usize, log_domain_size: u32) -> usize {
+    if coset_index % 2 == 0 {
+        coset_index / 2
     } else {
-        (1 << log_size) - (index + 1) / 2
+        ((2 << log_domain_size) - coset_index) / 2
     }
 }
 

crates/prover/src/examples/xor/eq_eval_constraints.rs

@@ -0,0 +1,258 @@
+use std::array;
+
+use educe::Educe;
+use num_traits::{One, Zero};
+
+use crate::constraint_framework::EvalAtRow;
+use crate::core::fields::qm31::SecureField;
+use crate::core::lookups::utils::eq;
+
+#[derive(Educe)]
+#[educe(Debug, Clone, Copy)]
+pub struct EqEvalsMask<E: EvalAtRow, const N_VARIABLES: usize> {
+    pub curr: E::EF,
+    pub next: E::EF,
+    pub next_next: E::EF,
+    pub is_first: E::F,
+    pub is_half_coset_step_by_log_size: [Option<CircleDomainStepMask<E>>; N_VARIABLES],
+}
+
+impl<E: EvalAtRow, const N_VARIABLES: usize> EqEvalsMask<E, N_VARIABLES> {
+    pub fn draw<const EQ_EVALS_TRACE: usize, const SELECTOR_TRACE: usize>(eval: &mut E) -> Self {
+        let [first] = eval.next_interaction_mask(SELECTOR_TRACE, [0]);
+        Self::draw_with_first::<EQ_EVALS_TRACE, SELECTOR_TRACE>(eval, first)
+    }
+
+    pub fn draw_with_first<const EQ_EVALS_TRACE: usize, const SELECTOR_TRACE: usize>(
+        eval: &mut E,
+        is_first: E::F,
+    ) -> Self {
+        let [curr, next, next_next] =
+            eval.next_extension_interaction_mask(EQ_EVALS_TRACE, [0, 1, 2]);
+
+        let mut is_half_coset_step_by_log_size = [None; N_VARIABLES];
+
+        for step in &mut is_half_coset_step_by_log_size[0..N_VARIABLES.saturating_sub(1)] {
+            *step = Some(CircleDomainStepMask::draw::<SELECTOR_TRACE>(eval));
+        }
+
+        Self {
+            curr,
+            next,
+            next_next,
+            is_first,
+            is_half_coset_step_by_log_size,
+        }
+    }
+
+    /// Evaluates the constraints on the mask to enforce correct [`eq`] evals.
+    ///
+    /// See <https://eprint.iacr.org/2023/1284.pdf> (Section 5.1).
+    ///
+    /// The [`eq`] trace evals appear ordered on a `CircleDomain` rather than a `Coset`. This
+    /// is why there are separate sets of constraints for each `CircleDomain` coset half.
+    pub fn eval(self, eval: &mut E, mle_eval_point: MleEvalPoint<N_VARIABLES>) {
+        let Self {
+            curr,
+            next,
+            next_next,
+            is_first,
+            is_half_coset_step_by_log_size,
+        } = self;
+
+        // Check initial value on half_coset0 and final value on half_coset1.
+        // TODO: Try use a single `selector.first` mult between both constraints.
+        eval.add_constraint((curr - mle_eval_point.eq_0_p) * is_first);
+        eval.add_constraint((next - mle_eval_point.eq_1_p) * is_first);
+
+        // Check all variables except the last (last variable is handled by the constraints above).
+        #[allow(clippy::needless_range_loop)]
+        for variable_i in 0..N_VARIABLES.saturating_sub(1) {
+            let half_coset0_next = next_next;
+            let half_coset1_prev = next_next;
+            let step = is_half_coset_step_by_log_size[variable_i].unwrap();
+            let carry_quotient = mle_eval_point.eq_carry_quotients[variable_i];
+            // Safe to combine these constraints as `is_step.half_coset0` and `is_step.half_coset1`
+            // are never non-zero at the same time on the trace.
+            let half_coset0_check = (curr - half_coset0_next * carry_quotient) * step.half_coset0;
+            let half_coset1_check = (curr * carry_quotient - half_coset1_prev) * step.half_coset1;
+            eval.add_constraint(half_coset0_check + half_coset1_check);
+        }
+    }
+}
+
+#[derive(Educe)]
+#[educe(Debug, Clone, Copy)]
+pub struct CircleDomainStepMask<E: EvalAtRow> {
+    half_coset0: E::F,
+    half_coset1: E::F,
+}
+
+impl<E: EvalAtRow> CircleDomainStepMask<E> {
+    pub fn draw<const TRACE: usize>(eval: &mut E) -> Self {
+        let [half_coset0, half_coset1] = eval.next_interaction_mask(TRACE, [0, -1]);
+        Self {
+            half_coset0,
+            half_coset1,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct MleEvalPoint<const N_VARIABLES: usize> {
+    // Equals `eq({0}^|p|, p)`.
+    eq_0_p: SecureField,
+    // Equals `eq({1}^|p|, p)`.
+    eq_1_p: SecureField,
+    // Index `i` stores `eq(({1}^|i|, 0), p[0..i+1]) / eq(({0}^|i|, 1), p[0..i+1])`.
+    eq_carry_quotients: [SecureField; N_VARIABLES],
+    // Point `p`.
+    _p: [SecureField; N_VARIABLES],
+}
+
+impl<const N_VARIABLES: usize> MleEvalPoint<N_VARIABLES> {
+    /// Creates new metadata from point `p`.
+    pub fn new(p: [SecureField; N_VARIABLES]) -> Self {
+        let zero = SecureField::zero();
+        let one = SecureField::one();
+
+        Self {
+            eq_0_p: eq(&[zero; N_VARIABLES], &p),
+            eq_1_p: eq(&[one; N_VARIABLES], &p),
+            eq_carry_quotients: array::from_fn(|i| {
+                let mut numer_assignment = vec![one; i + 1];
+                numer_assignment[i] = zero;
+                let mut denom_assignment = vec![zero; i + 1];
+                denom_assignment[i] = one;
+                eq(&numer_assignment, &p[..i + 1]) / eq(&denom_assignment, &p[..i + 1])
+            }),
+            _p: p,
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use std::array;
+
+    use num_traits::One;
+    use rand::rngs::SmallRng;
+    use rand::{Rng, SeedableRng};
+    use test_log::test;
+
+    use super::MleEvalPoint;
+    use crate::constraint_framework::assert_constraints;
+    use crate::constraint_framework::constant_columns::{gen_is_first, gen_is_step_with_offset};
+    use crate::core::backend::simd::SimdBackend;
+    use crate::core::fields::m31::BaseField;
+    use crate::core::fields::qm31::SecureField;
+    use crate::core::lookups::gkr_prover::GkrOps;
+    use crate::core::pcs::TreeVec;
+    use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+    use crate::core::poly::BitReversedOrder;
+    use crate::core::ColumnVec;
+    use crate::examples::xor::eq_eval_constraints::EqEvalsMask;
+
+    const EVALS_TRACE: usize = 0;
+    const CONST_TRACE: usize = 1;
+
+    #[test]
+    #[ignore = "SimdBackend `MIN_FFT_LOG_SIZE` is 5"]
+    fn eq_constraints_with_4_variables() {
+        const N_VARIABLES: usize = 4;
+        let mut rng = SmallRng::seed_from_u64(0);
+        let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen());
+        let base_trace = gen_evals_trace(&eval_point);
+        let constants_trace = gen_constants_trace(N_VARIABLES);
+        let traces = TreeVec::new(vec![base_trace, constants_trace]);
+        let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect());
+        let trace_domain = CanonicCoset::new(eval_point.len() as u32);
+        let mle_eval_point = MleEvalPoint::new(eval_point);
+
+        assert_constraints(&trace_polys, trace_domain, |mut eval| {
+            let mask = EqEvalsMask::draw::<EVALS_TRACE, CONST_TRACE>(&mut eval);
+            mask.eval(&mut eval, mle_eval_point);
+        });
+    }
+
+    #[test]
+    fn eq_constraints_with_5_variables() {
+        const N_VARIABLES: usize = 5;
+        let mut rng = SmallRng::seed_from_u64(0);
+        let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen());
+        let base_trace = gen_evals_trace(&eval_point);
+        let constants_trace = gen_constants_trace(N_VARIABLES);
+        let traces = TreeVec::new(vec![base_trace, constants_trace]);
+        let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect());
+        let trace_domain = CanonicCoset::new(eval_point.len() as u32);
+        let mle_eval_point = MleEvalPoint::new(eval_point);
+
+        assert_constraints(&trace_polys, trace_domain, |mut eval| {
+            let mask = EqEvalsMask::draw::<EVALS_TRACE, CONST_TRACE>(&mut eval);
+            mask.eval(&mut eval, mle_eval_point);
+        });
+    }
+
+    #[test]
+    fn eq_constraints_with_8_variables() {
+        const N_VARIABLES: usize = 8;
+        let mut rng = SmallRng::seed_from_u64(0);
+        let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen());
+        let base_trace = gen_evals_trace(&eval_point);
+        let constants_trace = gen_constants_trace(N_VARIABLES);
+        let traces = TreeVec::new(vec![base_trace, constants_trace]);
+        let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect());
+        let trace_domain = CanonicCoset::new(eval_point.len() as u32);
+        let mle_eval_point = MleEvalPoint::new(eval_point);
+
+        assert_constraints(&trace_polys, trace_domain, |mut eval| {
+            let mask = EqEvalsMask::draw::<EVALS_TRACE, CONST_TRACE>(&mut eval);
+            mask.eval(&mut eval, mle_eval_point);
+        });
+    }
+
+    /// Generates a trace.
+    ///
+    /// Trace structure:
+    ///
+    /// ```text
+    /// -----------------------------
+    /// |          eq evals         |
+    /// -----------------------------
+    /// |  c0  |  c1  |  c2  |  c3  |
+    /// -----------------------------
+    /// ```
+    pub fn gen_evals_trace(
+        eval_point: &[SecureField],
+    ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
+        // TODO(andrew): Consider storing eq evals as a SecureColumn.
+        let eq_evals = SimdBackend::gen_eq_evals(eval_point, SecureField::one()).into_evals();
+        let eq_evals_coordinate_columns = eq_evals.into_secure_column_by_coords().columns;
+
+        let n_variables = eval_point.len();
+        let domain = CanonicCoset::new(n_variables as u32).circle_domain();
+        eq_evals_coordinate_columns
+            .map(|col| CircleEvaluation::new(domain, col))
+            .into()
+    }
+
+    pub fn gen_constants_trace(
+        n_variables: usize,
+    ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
+        let log_size = n_variables as u32;
+        let mut constants_trace = Vec::new();
+        constants_trace.push(gen_is_first(log_size));
+
+        for variable_i in 1..n_variables as u32 {
+            let half_coset_log_step = variable_i;
+            let half_coset_offset = (1 << (half_coset_log_step - 1)) - 1;
+
+            let log_step = half_coset_log_step + 1;
+            let offset = half_coset_offset * 2;
+
+            constants_trace.push(gen_is_step_with_offset(log_size, log_step, offset))
+        }
+
+        constants_trace
+    }
+}

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

@@ -1 +1,2 @@
+pub mod eq_eval_constraints;
 pub mod prefix_sum_constraints;