Poseidon with logup

crates/prover/src/constraint_framework/assert.rs

@@ -7,6 +7,7 @@ use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::{CanonicCoset, CirclePoly};
+use crate::core::utils::circle_domain_order_to_coset_order;
 
 /// Evaluates expressions at a trace domain row, and asserts constraints. Mainly used for testing.
 pub struct AssertEvaluator<'a> {
@@ -66,10 +67,13 @@ pub fn assert_constraints<B: Backend>(
     let traces = trace_polys.as_ref().map(|tree| {
         tree.iter()
             .map(|poly| {
-                poly.evaluate(trace_domain.circle_domain())
-                    .bit_reverse()
-                    .values
-                    .to_cpu()
+                circle_domain_order_to_coset_order(
+                    &poly
+                        .evaluate(trace_domain.circle_domain())
+                        .bit_reverse()
+                        .values
+                        .to_cpu(),
+                )
             })
             .collect()
     });

crates/prover/src/constraint_framework/logup.rs

@@ -1,7 +1,8 @@
 use itertools::Itertools;
-use num_traits::Zero;
+use num_traits::{One, Zero};
 use tracing::{span, Level};
 
+use super::EvalAtRow;
 use crate::core::backend::simd::column::SecureFieldVec;
 use crate::core::backend::simd::m31::LOG_N_LANES;
 use crate::core::backend::simd::qm31::PackedSecureField;
@@ -14,10 +15,91 @@ use crate::core::fields::secure_column::SecureColumn;
 use crate::core::fields::FieldExpOps;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use crate::core::poly::BitReversedOrder;
-use crate::core::utils::bit_reverse_index;
+use crate::core::utils::{bit_reverse_index, shifted_secure_combination};
 use crate::core::ColumnVec;
 
-#[derive(Copy, Clone, Debug)]
+pub struct LogupAtRow<const BATCH_SIZE: usize, E: EvalAtRow> {
+    pub interaction: usize,
+    pub queue: [(E::EF, E::EF); BATCH_SIZE],
+    pub queue_size: usize,
+    pub claimed_sum: SecureField,
+    pub prev_mask: E::EF,
+    pub is_first: E::F,
+}
+impl<const BATCH_SIZE: usize, E: EvalAtRow> LogupAtRow<BATCH_SIZE, E> {
+    pub fn new(interaction: usize, claimed_sum: SecureField, is_first: E::F) -> Self {
+        Self {
+            interaction,
+            queue: [(E::EF::zero(), E::EF::zero()); BATCH_SIZE],
+            queue_size: 0,
+            claimed_sum,
+            prev_mask: E::EF::zero(),
+            is_first,
+        }
+    }
+    pub fn push_lookup(
+        &mut self,
+        eval: &mut E,
+        numerator: E::EF,
+        values: &[E::F],
+        lookup_elements: LookupElements,
+    ) {
+        let shifted_value = shifted_secure_combination(
+            values,
+            E::EF::zero() + lookup_elements.alpha,
+            E::EF::zero() + lookup_elements.z,
+        );
+        self.push_frac(eval, numerator, shifted_value);
+    }
+
+    pub fn push_frac(&mut self, eval: &mut E, p: E::EF, q: E::EF) {
+        if self.queue_size < BATCH_SIZE {
+            self.queue[self.queue_size] = (p, q);
+            self.queue_size += 1;
+            return;
+        }
+
+        // Compute sum_i pi/qi over batch, as a fraction, p/q.
+        let (num, denom) = self
+            .queue
+            .iter()
+            .copied()
+            .fold((E::EF::zero(), E::EF::one()), |(p0, q0), (pi, qi)| {
+                (p0 * qi + pi * q0, qi * q0)
+            });
+
+        self.queue[0] = (p, q);
+        self.queue_size = 1;
+
+        // Add a constraint that p / q = diff.
+        let cur = E::combine_ef(std::array::from_fn(|_| {
+            eval.next_interaction_mask(1, [0])[0]
+        }));
+        let diff = cur - self.prev_mask;
+        self.prev_mask = cur;
+        eval.add_constraint(diff * denom - num);
+    }
+
+    pub fn finalize(self, eval: &mut E) {
+        let (p, q) = self.queue[0..self.queue_size]
+            .iter()
+            .copied()
+            .fold((E::EF::zero(), E::EF::one()), |(p0, q0), (pi, qi)| {
+                (p0 * qi + pi * q0, qi * q0)
+            });
+
+        let cumulative_mask_values =
+            std::array::from_fn(|_| eval.next_interaction_mask(self.interaction, [0, -1]));
+        let cur = E::combine_ef(cumulative_mask_values.map(|[cur, _prev]| cur));
+        let up = E::combine_ef(cumulative_mask_values.map(|[_cur, prev]| prev));
+        let up = up - self.is_first * self.claimed_sum;
+        let diff = cur - up - self.prev_mask;
+
+        eval.add_constraint(diff * q - p);
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct LookupElements {
     pub z: SecureField,
     pub alpha: SecureField,

crates/prover/src/constraint_framework/mod.rs

@@ -7,7 +7,7 @@ mod point;
 mod simd_domain;
 
 use std::fmt::Debug;
-use std::ops::{Add, AddAssign, Mul, Sub};
+use std::ops::{Add, AddAssign, Mul, Neg, Sub};
 
 pub use assert::{assert_constraints, AssertEvaluator};
 pub use info::InfoEvaluator;
@@ -44,6 +44,7 @@ pub trait EvalAtRow {
         + Copy
         + Debug
         + Zero
+        + Neg<Output = Self::EF>
         + Add<SecureField, Output = Self::EF>
         + Sub<SecureField, Output = Self::EF>
         + Mul<SecureField, Output = Self::EF>