Feat/airbuilder constraints

crypto/stark/src/air_builder.rs

@@ -0,0 +1,183 @@
+use crate::frame::Frame;
+use crate::trace::LDETraceTable;
+use math::field::element::FieldElement;
+use math::field::traits::{IsFFTField, IsField, IsSubFieldOf};
+
+/// Plonky3-style builder for fused constraint evaluation + alpha combination.
+///
+/// Constraints call `assert_zero(expr)` which internally accumulates
+/// alpha^i * expr into a running sum. No intermediate buffer, no vtable dispatch.
+///
+/// Parameterized by field type `F` (not an associated type) so that `dyn AirBuilder<F>`
+/// is object-safe. This allows the AIR trait to remain dyn-compatible while supporting
+/// the builder pattern: `eval_constraints(&self, builder: &mut dyn AirBuilder<E>)`.
+pub trait AirBuilder<F: IsField> {
+    /// Read main trace column at (row_offset, col). offset=0 is current row.
+    fn main(&self, offset: usize, col: usize) -> FieldElement<F>;
+
+    /// Read aux trace column at (row_offset, col).
+    fn aux(&self, offset: usize, col: usize) -> FieldElement<F>;
+
+    /// Assert expr == 0. Internally: accumulator += alpha^constraint_idx * expr.
+    fn assert_zero(&mut self, expr: FieldElement<F>);
+
+    /// RAP challenge by index.
+    fn challenge(&self, idx: usize) -> &FieldElement<F>;
+
+    /// Pre-computed LogUp alpha powers.
+    fn logup_alpha_power(&self, idx: usize) -> &FieldElement<F>;
+
+    /// LogUp table offset (L/N).
+    fn logup_table_offset(&self) -> &FieldElement<F>;
+}
+
+pub struct ProverBuilder<'a, F: IsSubFieldOf<E> + IsFFTField, E: IsField> {
+    lde_trace: &'a LDETraceTable<F, E>,
+    row: usize,
+    step_size: usize,
+    num_rows: usize,
+    accumulator: FieldElement<E>,
+    alpha: FieldElement<E>,
+    alpha_power: FieldElement<E>,
+    rap_challenges: &'a [FieldElement<E>],
+    logup_alpha_powers: &'a [FieldElement<E>],
+    logup_table_offset_val: &'a FieldElement<E>,
+}
+
+impl<'a, F, E> ProverBuilder<'a, F, E>
+where
+    F: IsSubFieldOf<E> + IsFFTField + Send + Sync,
+    E: IsField + Send + Sync,
+{
+    pub fn new(
+        lde_trace: &'a LDETraceTable<F, E>,
+        row: usize,
+        alpha: &FieldElement<E>,
+        rap_challenges: &'a [FieldElement<E>],
+        logup_alpha_powers: &'a [FieldElement<E>],
+        logup_table_offset: &'a FieldElement<E>,
+    ) -> Self {
+        Self {
+            lde_trace,
+            row,
+            step_size: lde_trace.lde_step_size,
+            num_rows: lde_trace.num_rows(),
+            accumulator: FieldElement::zero(),
+            alpha: alpha.clone(),
+            alpha_power: FieldElement::one(),
+            rap_challenges,
+            logup_alpha_powers,
+            logup_table_offset_val: logup_table_offset,
+        }
+    }
+
+    pub fn finish(self) -> FieldElement<E> {
+        self.accumulator
+    }
+}
+
+impl<'a, F, E> AirBuilder<E> for ProverBuilder<'a, F, E>
+where
+    F: IsSubFieldOf<E> + IsFFTField + Send + Sync,
+    E: IsField + Send + Sync,
+{
+    #[inline]
+    fn main(&self, offset: usize, col: usize) -> FieldElement<E> {
+        let lde_row = (self.row + offset * self.step_size) % self.num_rows;
+        self.lde_trace.get_main(lde_row, col).clone().to_extension()
+    }
+
+    #[inline]
+    fn aux(&self, offset: usize, col: usize) -> FieldElement<E> {
+        let lde_row = (self.row + offset * self.step_size) % self.num_rows;
+        self.lde_trace.get_aux(lde_row, col).clone()
+    }
+
+    #[inline]
+    fn assert_zero(&mut self, expr: FieldElement<E>) {
+        self.accumulator = &self.accumulator + &self.alpha_power * &expr;
+        self.alpha_power = &self.alpha_power * &self.alpha;
+    }
+
+    fn challenge(&self, idx: usize) -> &FieldElement<E> {
+        &self.rap_challenges[idx]
+    }
+
+    fn logup_alpha_power(&self, idx: usize) -> &FieldElement<E> {
+        &self.logup_alpha_powers[idx]
+    }
+
+    fn logup_table_offset(&self) -> &FieldElement<E> {
+        self.logup_table_offset_val
+    }
+}
+
+pub struct VerifierBuilder<'a, E: IsField> {
+    frame: &'a Frame<E, E>,
+    accumulator: FieldElement<E>,
+    alpha: FieldElement<E>,
+    alpha_power: FieldElement<E>,
+    rap_challenges: &'a [FieldElement<E>],
+    logup_alpha_powers: &'a [FieldElement<E>],
+    logup_table_offset_val: &'a FieldElement<E>,
+}
+
+impl<'a, E: IsField> VerifierBuilder<'a, E> {
+    pub fn new(
+        frame: &'a Frame<E, E>,
+        alpha: &FieldElement<E>,
+        rap_challenges: &'a [FieldElement<E>],
+        logup_alpha_powers: &'a [FieldElement<E>],
+        logup_table_offset: &'a FieldElement<E>,
+    ) -> Self {
+        Self {
+            frame,
+            accumulator: FieldElement::zero(),
+            alpha: alpha.clone(),
+            alpha_power: FieldElement::one(),
+            rap_challenges,
+            logup_alpha_powers,
+            logup_table_offset_val: logup_table_offset,
+        }
+    }
+
+    pub fn finish(self) -> FieldElement<E> {
+        self.accumulator
+    }
+}
+
+impl<'a, E: IsField> AirBuilder<E> for VerifierBuilder<'a, E> {
+    #[inline]
+    fn main(&self, offset: usize, col: usize) -> FieldElement<E> {
+        self.frame
+            .get_evaluation_step(offset)
+            .get_main_evaluation_element(0, col)
+            .clone()
+    }
+
+    #[inline]
+    fn aux(&self, offset: usize, col: usize) -> FieldElement<E> {
+        self.frame
+            .get_evaluation_step(offset)
+            .get_aux_evaluation_element(0, col)
+            .clone()
+    }
+
+    #[inline]
+    fn assert_zero(&mut self, expr: FieldElement<E>) {
+        self.accumulator = &self.accumulator + &self.alpha_power * &expr;
+        self.alpha_power = &self.alpha_power * &self.alpha;
+    }
+
+    fn challenge(&self, idx: usize) -> &FieldElement<E> {
+        &self.rap_challenges[idx]
+    }
+
+    fn logup_alpha_power(&self, idx: usize) -> &FieldElement<E> {
+        &self.logup_alpha_powers[idx]
+    }
+
+    fn logup_table_offset(&self) -> &FieldElement<E> {
+        self.logup_table_offset_val
+    }
+}

crypto/stark/src/constraints/evaluator.rs

@@ -49,8 +49,10 @@ where
         num_periodic: usize,
         offsets: &[usize],
         logup_table_offset: &FieldElement<FieldExtension>,
+        composition_alpha: &FieldElement<FieldExtension>,
     ) -> Vec<FieldElement<FieldExtension>> {
         let is_uniform = zerofier_data.is_uniform();
+        let use_builder = air.uses_builder();
 
         // Pre-compute LogUp alpha powers once for all LDE domain points.
         let logup_alpha_powers: Vec<FieldElement<FieldExtension>> =
@@ -96,6 +98,24 @@ where
                         )
                     },
                     |(transition_buf, periodic_buf, frame), (i, boundary)| {
+                        if use_builder {
+                            let mut builder = crate::air_builder::ProverBuilder::new(
+                                lde_trace,
+                                i,
+                                composition_alpha,
+                                rap_challenges,
+                                &logup_alpha_powers,
+                                logup_table_offset,
+                            );
+                            air.eval_constraints_with_builder(&mut builder);
+                            let acc = if is_uniform {
+                                zerofier_data.get_uniform(i) * &builder.finish()
+                            } else {
+                                panic!("AirBuilder requires uniform zerofiers")
+                            };
+                            return acc + boundary;
+                        }
+
                         frame.fill_from_lde(lde_trace, i, offsets);
 
                         for (j, col) in lde_periodic_columns.iter().enumerate() {
@@ -148,6 +168,24 @@ where
                 .into_iter()
                 .enumerate()
                 .map(|(i, boundary)| {
+                    if use_builder {
+                        let mut builder = crate::air_builder::ProverBuilder::new(
+                            lde_trace,
+                            i,
+                            composition_alpha,
+                            rap_challenges,
+                            &logup_alpha_powers,
+                            logup_table_offset,
+                        );
+                        air.eval_constraints_with_builder(&mut builder);
+                        let acc = if is_uniform {
+                            zerofier_data.get_uniform(i) * &builder.finish()
+                        } else {
+                            panic!("AirBuilder requires uniform zerofiers")
+                        };
+                        return acc + boundary;
+                    }
+
                     frame.fill_from_lde(lde_trace, i, offsets);
 
                     for (j, col) in lde_periodic_columns.iter().enumerate() {
@@ -315,6 +353,14 @@ where
         let num_periodic = lde_periodic_columns.len();
         let offsets = &air.context().transition_offsets;
 
+        // Extract the raw composition alpha (β) from the pre-computed powers [1, β, β², ...].
+        // The ProverBuilder generates its own powers internally via assert_zero().
+        let composition_alpha = if transition_coefficients.len() >= 2 {
+            transition_coefficients[1].clone()
+        } else {
+            FieldElement::one()
+        };
+
         Self::evaluate_transitions(
             air,
             lde_trace,
@@ -327,6 +373,7 @@ where
             num_periodic,
             offsets,
             &self.logup_table_offset,
+            &composition_alpha,
         )
     }
 }

crypto/stark/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod air_builder;
 #[cfg(feature = "debug-checks")]
 pub mod bus_debug;
 pub mod constraints;