From 1c6ef62076b303d31ad99c0750e00abd6560882b Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Tue, 28 Jul 2026 13:06:31 -0300 Subject: [PATCH 1/3] Replace the AIR builder stack with a compiled constraint IR Circuits are now data, not trait impls: frontend Expr/ExtExpr trees compile to a flat, hash-consed, base-only node graph (graph::ConstraintGraph) that the prover, verifier and witness generator all evaluate with a dense forward sweep. Extension constraints are coordinate-expanded at compile time (Karatsuba for D=2, schoolbook otherwise), so the extension degree never appears in the compiled artifact; constraint roots are canonicalized (const-zero dropped, const-nonzero rejected at setup, sorted and deduplicated). - expr.rs: frontend expression trees and the internal CircuitSpec - graph.rs: interning compiler producing ConstraintGraph - eval.rs: dense sweep evaluator shared by prover, verifier and the lookup witness (plus a recursive reference evaluator for tests) - lookup.rs: logUp stage-2 constraints synthesized as ordinary ExtExpr data and compiled like user constraints; Lookup keeps its push/pull constructors - system.rs: System over CircuitInputs (main width, preprocessed trace, constraints, lookups); stage-2 and public-input layout is derived at setup, not authored - p3_adapter.rs: Plonky3-style authoring preserved. An AirBuilder implementation records Air::eval into Expr trees; LookupAir pairs an AIR with its lookups and converts into CircuitInputs; System::new accepts anything Into. Existing AIR-based tests, examples and benches run unchanged except for import paths and the System type losing its AIR parameter. - builder/ (symbolic expressions, constraint folders, check builder) deleted; prover/verifier keep the same protocol and transcript --- benches/multi_stark.rs | 11 +- examples/lookup_proof.rs | 4 +- examples/preprocessed_proof.rs | 4 +- examples/simple_proof.rs | 2 +- src/builder/check.rs | 188 ---------- src/builder/folder.rs | 211 ----------- src/builder/mod.rs | 23 -- src/builder/symbolic.rs | 498 ------------------------- src/eval.rs | 212 +++++++++++ src/expr.rs | 305 ++++++++++++++++ src/graph.rs | 508 ++++++++++++++++++++++++++ src/lib.rs | 5 +- src/lookup.rs | 284 ++++++-------- src/p3_adapter.rs | 354 ++++++++++++++++++ src/prover.rs | 276 ++++++-------- src/system.rs | 346 ++++++++++++------ src/test_circuits/baby_bear_config.rs | 4 +- src/test_circuits/blake3.rs | 6 +- src/test_circuits/byte_operations.rs | 6 +- src/test_circuits/mod.rs | 2 +- src/test_circuits/u32_add.rs | 8 +- src/verifier.rs | 179 +++++---- 22 files changed, 1946 insertions(+), 1490 deletions(-) delete mode 100644 src/builder/check.rs delete mode 100644 src/builder/folder.rs delete mode 100644 src/builder/mod.rs delete mode 100644 src/builder/symbolic.rs create mode 100644 src/eval.rs create mode 100644 src/expr.rs create mode 100644 src/graph.rs create mode 100644 src/p3_adapter.rs diff --git a/benches/multi_stark.rs b/benches/multi_stark.rs index 5daae22..a910f9f 100644 --- a/benches/multi_stark.rs +++ b/benches/multi_stark.rs @@ -20,8 +20,8 @@ //! `TEXRAY_PREFIXES=` (empty) to also render Plonky3's internal spans. use criterion::{BenchmarkId, Criterion, criterion_group}; -use multi_stark::builder::symbolic::{SymbolicExpression, preprocessed_var, var}; -use multi_stark::lookup::{Lookup, LookupAir}; +use multi_stark::lookup::Lookup; +use multi_stark::p3_adapter::{LookupAir, SymbolicExpression, preprocessed_var, var}; use multi_stark::system::{ProverKey, System, SystemWitness}; use multi_stark::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; use multi_stark::{ @@ -168,10 +168,7 @@ impl U32CS { // Witness generation // --------------------------------------------------------------------------- -fn build_witness( - num_adds: usize, - system: &System, -) -> SystemWitness { +fn build_witness(num_adds: usize, system: &System) -> SystemWitness { let byte_width = 1; let add_width = 14; let add_height = num_adds.next_power_of_two(); @@ -261,7 +258,7 @@ fn bench_config() -> GoldilocksBlake3Config { } fn build_system() -> ( - System, + System, ProverKey, ) { let byte_table = LookupAir::new(U32CS::ByteTable, U32CS::ByteTable.lookups()); diff --git a/examples/lookup_proof.rs b/examples/lookup_proof.rs index 23bf945..0dc6c15 100644 --- a/examples/lookup_proof.rs +++ b/examples/lookup_proof.rs @@ -12,8 +12,8 @@ //! cargo run --example lookup_proof --release //! ``` -use multi_stark::builder::symbolic::{SymbolicExpression, var}; -use multi_stark::lookup::{Lookup, LookupAir}; +use multi_stark::lookup::Lookup; +use multi_stark::p3_adapter::{LookupAir, SymbolicExpression, var}; use multi_stark::system::{System, SystemWitness}; use multi_stark::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; use multi_stark::{ diff --git a/examples/preprocessed_proof.rs b/examples/preprocessed_proof.rs index cd5bc28..3da2efb 100644 --- a/examples/preprocessed_proof.rs +++ b/examples/preprocessed_proof.rs @@ -11,8 +11,8 @@ //! cargo run --example preprocessed_proof --release //! ``` -use multi_stark::builder::symbolic::{SymbolicExpression, preprocessed_var, var}; -use multi_stark::lookup::{Lookup, LookupAir}; +use multi_stark::lookup::Lookup; +use multi_stark::p3_adapter::{LookupAir, SymbolicExpression, preprocessed_var, var}; use multi_stark::system::{System, SystemWitness}; use multi_stark::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; use multi_stark::{ diff --git a/examples/simple_proof.rs b/examples/simple_proof.rs index 8cf3b61..3147130 100644 --- a/examples/simple_proof.rs +++ b/examples/simple_proof.rs @@ -9,7 +9,7 @@ //! cargo run --example simple_proof --release //! ``` -use multi_stark::lookup::LookupAir; +use multi_stark::p3_adapter::LookupAir; use multi_stark::system::{System, SystemWitness}; use multi_stark::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; use multi_stark::{ diff --git a/src/builder/check.rs b/src/builder/check.rs deleted file mode 100644 index b114978..0000000 --- a/src/builder/check.rs +++ /dev/null @@ -1,188 +0,0 @@ -use p3_air::{Air, AirBuilder, ExtensionBuilder, RowWindow}; -use p3_matrix::Matrix; -use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixView}; -use p3_matrix::stack::VerticalPair; - -use p3_field::{ExtensionField, Field}; - -use super::TwoStagedBuilder; - -pub fn check_constraints( - air: &A, - preprocessed: Option<&RowMajorMatrix>, - stage_1: &RowMajorMatrix, - stage_2: &RowMajorMatrix, - public_values: &[F], - stage_2_public_values: &[EF], -) where - F: Field, - EF: ExtensionField, - A: for<'a> Air>, -{ - let height = stage_1.height(); - - (0..height).for_each(|i| { - let i_next = (i + 1) % height; - - let stage_1_local = stage_1.row_slice(i).unwrap(); // i < height so unwrap should never fail. - let stage_1_next = stage_1.row_slice(i_next).unwrap(); // i_next < height so unwrap should never fail. - let stage_1 = VerticalPair::new( - RowMajorMatrixView::new_row(&*stage_1_local), - RowMajorMatrixView::new_row(&*stage_1_next), - ); - - let stage_2_local = stage_2.row_slice(i).unwrap(); // i < height so unwrap should never fail. - let stage_2_next = stage_2.row_slice(i_next).unwrap(); // i_next < height so unwrap should never fail. - let stage_2 = VerticalPair::new( - RowMajorMatrixView::new_row(&*stage_2_local), - RowMajorMatrixView::new_row(&*stage_2_next), - ); - let mut builder = DebugConstraintBuilder { - row_index: i, - preprocessed: RowWindow::from_two_rows(&[], &[]), - stage_1, - stage_2, - public_values, - stage_2_public_values, - is_first_row: F::from_bool(i == 0), - is_last_row: F::from_bool(i == height - 1), - is_transition: F::from_bool(i != height - 1), - }; - // We must call `eval` on the same block as the `preprocessed` matrix view, otherwise the borrow checker will complain. - // Mutation is used to remove code duplication. - if let Some(preprocessed) = preprocessed { - let preprocessed_local = preprocessed.row_slice(i).unwrap(); // i < height so unwrap should never fail. - let preprocessed_next = preprocessed.row_slice(i_next).unwrap(); // i_next < height so unwrap should never fail. - builder.preprocessed = - RowWindow::from_two_rows(&preprocessed_local, &preprocessed_next); - air.eval(&mut builder); - } else { - air.eval(&mut builder); - } - }); -} - -#[derive(Debug)] -pub struct DebugConstraintBuilder<'a, F: Field, EF: ExtensionField> { - /// The index of the row currently being evaluated. - row_index: usize, - /// A two-row window over the preprocessed trace (current and next row). - preprocessed: RowWindow<'a, F>, - stage_1: VerticalPair, RowMajorMatrixView<'a, F>>, - stage_2: VerticalPair, RowMajorMatrixView<'a, EF>>, - /// The public values provided for constraint validation (e.g. inputs or outputs). - public_values: &'a [F], - stage_2_public_values: &'a [EF], - /// A flag indicating whether this is the first row. - is_first_row: F, - /// A flag indicating whether this is the last row. - is_last_row: F, - /// A flag indicating whether this is a transition row (not the last row). - is_transition: F, -} - -impl<'a, F: Field, EF: ExtensionField> AirBuilder for DebugConstraintBuilder<'a, F, EF> { - type F = F; - type Expr = F; - type Var = F; - type PreprocessedWindow = RowWindow<'a, F>; - type MainWindow = RowWindow<'a, F>; - type PublicVar = F; - - fn main(&self) -> Self::MainWindow { - RowWindow::from_two_rows(self.stage_1.top.values, self.stage_1.bottom.values) - } - - fn preprocessed(&self) -> &Self::PreprocessedWindow { - &self.preprocessed - } - - fn is_first_row(&self) -> Self::Expr { - self.is_first_row - } - - fn is_last_row(&self) -> Self::Expr { - self.is_last_row - } - - /// # Panics - /// This function panics if `size` is not `2`. - fn is_transition_window(&self, size: usize) -> Self::Expr { - if size == 2 { - self.is_transition - } else { - panic!("only supports a window size of 2") - } - } - - fn assert_zero>(&mut self, x: I) { - assert_eq!( - x.into(), - F::ZERO, - "constraints had nonzero value on row {}", - self.row_index - ); - } - - fn assert_eq, I2: Into>(&mut self, x: I1, y: I2) { - let x = x.into(); - let y = y.into(); - assert_eq!( - x, y, - "values didn't match on row {}: {} != {}", - self.row_index, x, y - ); - } - - fn public_values(&self) -> &[Self::PublicVar] { - self.public_values - } -} - -impl> ExtensionBuilder for DebugConstraintBuilder<'_, F, EF> { - type EF = EF; - type ExprEF = EF; - type VarEF = EF; - - /// Assert that an extension field expression is zero. - fn assert_zero_ext(&mut self, x: I) - where - I: Into, - { - let x = x.into(); - assert_eq!( - x, - EF::ZERO, - "constraints had nonzero value on row {}", - self.row_index - ); - } - - fn assert_eq_ext(&mut self, x: I1, y: I2) - where - I1: Into, - I2: Into, - { - let x = x.into(); - let y = y.into(); - assert_eq!( - x, y, - "values didn't match on row {}: {} != {}", - self.row_index, x, y - ); - } -} - -impl<'a, F: Field, EF: ExtensionField> TwoStagedBuilder for DebugConstraintBuilder<'a, F, EF> { - type MP = VerticalPair, RowMajorMatrixView<'a, EF>>; - - type Stage2PublicVar = Self::EF; - - fn stage_2(&self) -> Self::MP { - self.stage_2 - } - - fn stage_2_public_values(&self) -> &[Self::Stage2PublicVar] { - self.stage_2_public_values - } -} diff --git a/src/builder/folder.rs b/src/builder/folder.rs deleted file mode 100644 index 3155157..0000000 --- a/src/builder/folder.rs +++ /dev/null @@ -1,211 +0,0 @@ -/// Constraint folders for the prover and verifier, adapted from Plonky3. -use p3_air::{AirBuilder, ExtensionBuilder, RowWindow}; -use p3_field::{Algebra, BasedVectorSpace, ExtensionField, Field}; -use p3_matrix::dense::RowMajorMatrixView; -use p3_matrix::stack::VerticalPair; - -use super::TwoStagedBuilder; - -/// Packed base-field values of `F`. -type PackedVal = ::Packing; -/// Packed extension-field values of `EF` over `F`. -type PackedExt = >::ExtensionPacking; - -pub struct ProverConstraintFolder<'a, F: Field, EF: ExtensionField> { - pub preprocessed: RowWindow<'a, PackedVal>, - pub stage_1: RowMajorMatrixView<'a, PackedVal>, - pub stage_2: RowMajorMatrixView<'a, PackedExt>, - pub stage_1_public_values: &'a [F], - pub stage_2_public_values: &'a [EF], - pub is_first_row: PackedVal, - pub is_last_row: PackedVal, - pub is_transition: PackedVal, - pub alpha_powers: &'a [EF], - pub decomposed_alpha_powers: &'a [Vec], - pub accumulator: PackedExt, - pub constraint_index: usize, -} - -type ViewPair<'a, T> = VerticalPair, RowMajorMatrixView<'a, T>>; - -pub struct VerifierConstraintFolder<'a, F: Field, EF: ExtensionField> { - pub preprocessed: RowWindow<'a, EF>, - pub stage_1: ViewPair<'a, EF>, - pub stage_2: ViewPair<'a, EF>, - pub stage_1_public_values: &'a [F], - pub stage_2_public_values: &'a [EF], - pub is_first_row: EF, - pub is_last_row: EF, - pub is_transition: EF, - pub alpha: EF, - pub accumulator: EF, -} - -impl<'a, F: Field, EF: ExtensionField> AirBuilder for ProverConstraintFolder<'a, F, EF> { - type F = F; - type Expr = PackedVal; - type Var = PackedVal; - type PreprocessedWindow = RowWindow<'a, PackedVal>; - type MainWindow = RowWindow<'a, PackedVal>; - type PublicVar = F; - - #[inline] - fn main(&self) -> Self::MainWindow { - RowWindow::from_view(&self.stage_1) - } - - fn preprocessed(&self) -> &Self::PreprocessedWindow { - &self.preprocessed - } - - #[inline] - fn is_first_row(&self) -> Self::Expr { - self.is_first_row - } - - #[inline] - fn is_last_row(&self) -> Self::Expr { - self.is_last_row - } - - /// # Panics - /// This function panics if `size` is not `2`. - #[inline] - fn is_transition_window(&self, size: usize) -> Self::Expr { - if size == 2 { - self.is_transition - } else { - panic!("multi-stark only supports a window size of 2") - } - } - - #[inline] - fn assert_zero>(&mut self, x: I) { - let x: PackedVal = x.into(); - let alpha_power = self.alpha_powers[self.constraint_index]; - self.accumulator += Into::>::into(alpha_power) * x; - self.constraint_index += 1; - } - - #[inline] - fn assert_zeros>(&mut self, array: [I; N]) { - let expr_array: [Self::Expr; N] = array.map(Into::into); - self.accumulator += PackedExt::::from_basis_coefficients_fn(|i| { - let alpha_powers = &self.decomposed_alpha_powers[i] - [self.constraint_index..(self.constraint_index + N)]; - PackedVal::::batched_linear_combination(&expr_array, alpha_powers) - }); - self.constraint_index += N; - } - - #[inline] - fn public_values(&self) -> &[Self::PublicVar] { - self.stage_1_public_values - } -} - -impl> ExtensionBuilder for ProverConstraintFolder<'_, F, EF> { - type EF = EF; - type ExprEF = PackedExt; - type VarEF = PackedExt; - - /// Assert that an extension field expression is zero. - fn assert_zero_ext(&mut self, x: I) - where - I: Into, - { - let x: PackedExt = x.into(); - let alpha_power = self.alpha_powers[self.constraint_index]; - self.accumulator += Into::>::into(alpha_power) * x; - self.constraint_index += 1; - } -} - -impl<'a, F: Field, EF: ExtensionField> TwoStagedBuilder for ProverConstraintFolder<'a, F, EF> { - type MP = RowMajorMatrixView<'a, PackedExt>; - - type Stage2PublicVar = Self::EF; - - fn stage_2(&self) -> Self::MP { - self.stage_2 - } - - fn stage_2_public_values(&self) -> &[Self::Stage2PublicVar] { - self.stage_2_public_values - } -} - -impl<'a, F: Field, EF: ExtensionField> AirBuilder for VerifierConstraintFolder<'a, F, EF> { - type F = F; - type Expr = EF; - type Var = EF; - type PreprocessedWindow = RowWindow<'a, EF>; - type MainWindow = RowWindow<'a, EF>; - type PublicVar = F; - - fn main(&self) -> Self::MainWindow { - RowWindow::from_two_rows(self.stage_1.top.values, self.stage_1.bottom.values) - } - - fn preprocessed(&self) -> &Self::PreprocessedWindow { - &self.preprocessed - } - - fn is_first_row(&self) -> Self::Expr { - self.is_first_row - } - - fn is_last_row(&self) -> Self::Expr { - self.is_last_row - } - - /// # Panics - /// This function panics if `size` is not `2`. - fn is_transition_window(&self, size: usize) -> Self::Expr { - if size == 2 { - self.is_transition - } else { - panic!("multi-stark only supports a window size of 2") - } - } - - fn assert_zero>(&mut self, x: I) { - let x: EF = x.into(); - self.accumulator *= self.alpha; - self.accumulator += x; - } - - fn public_values(&self) -> &[Self::PublicVar] { - self.stage_1_public_values - } -} - -impl> ExtensionBuilder for VerifierConstraintFolder<'_, F, EF> { - type EF = EF; - type ExprEF = EF; - type VarEF = EF; - - /// Assert that an extension field expression is zero. - fn assert_zero_ext(&mut self, x: I) - where - I: Into, - { - let x: EF = x.into(); - self.accumulator *= self.alpha; - self.accumulator += x; - } -} - -impl<'a, F: Field, EF: ExtensionField> TwoStagedBuilder for VerifierConstraintFolder<'a, F, EF> { - type MP = ViewPair<'a, EF>; - - type Stage2PublicVar = Self::EF; - - fn stage_2(&self) -> Self::MP { - self.stage_2 - } - - fn stage_2_public_values(&self) -> &[Self::Stage2PublicVar] { - self.stage_2_public_values - } -} diff --git a/src/builder/mod.rs b/src/builder/mod.rs deleted file mode 100644 index 7c41c5b..0000000 --- a/src/builder/mod.rs +++ /dev/null @@ -1,23 +0,0 @@ -use p3_air::ExtensionBuilder; -use p3_matrix::Matrix; - -pub mod check; -pub mod folder; -pub mod symbolic; - -/// Extension of [`ExtensionBuilder`] that provides access to a second-stage -/// trace and its associated public values (used by the lookup argument). -pub trait TwoStagedBuilder: ExtensionBuilder { - /// Matrix type for the stage 2 trace window. - type MP: Matrix; - - /// Variable type for stage 2 public values. - type Stage2PublicVar: Into + Copy; - - /// Returns the stage 2 trace window. - fn stage_2(&self) -> Self::MP; - - /// Returns the stage 2 public values (lookup and fingerprint challenges, - /// current accumulator, next accumulator). - fn stage_2_public_values(&self) -> &[Self::Stage2PublicVar]; -} diff --git a/src/builder/symbolic.rs b/src/builder/symbolic.rs deleted file mode 100644 index 557e45f..0000000 --- a/src/builder/symbolic.rs +++ /dev/null @@ -1,498 +0,0 @@ -/// Symbolic constraint builder and expressions, adapted from Plonky3. -use p3_air::{Air, AirBuilder, ExtensionBuilder}; -use p3_field::{Algebra, Dup, ExtensionField, Field, InjectiveMonomial, PrimeCharacteristicRing}; -use p3_matrix::dense::RowMajorMatrix; -use std::fmt::Debug; -use std::iter::{Product, Sum}; -use std::marker::PhantomData; -use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; - -use super::TwoStagedBuilder; - -#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] -pub enum Entry { - Preprocessed { offset: usize }, - Main { offset: usize }, - Stage2 { offset: usize }, - Public, - Stage2Public, - Challenge, -} - -/// A variable within the evaluation window, i.e. a column in either the local or next row. -#[derive(Copy, Clone, Debug)] -pub struct SymbolicVariable { - pub entry: Entry, - pub index: usize, - pub(crate) _phantom: PhantomData, -} - -impl SymbolicVariable { - pub const fn new(entry: Entry, index: usize) -> Self { - Self { - entry, - index, - _phantom: PhantomData, - } - } - - pub const fn degree_multiple(&self) -> usize { - match self.entry { - Entry::Preprocessed { .. } | Entry::Main { .. } | Entry::Stage2 { .. } => 1, - Entry::Public | Entry::Challenge | Entry::Stage2Public => 0, - } - } -} - -impl From> for SymbolicExpression { - fn from(value: SymbolicVariable) -> Self { - Self::Variable(value) - } -} - -impl Add for SymbolicVariable -where - T: Into>, -{ - type Output = SymbolicExpression; - - fn add(self, rhs: T) -> Self::Output { - SymbolicExpression::from(self) + rhs.into() - } -} - -impl Sub for SymbolicVariable -where - T: Into>, -{ - type Output = SymbolicExpression; - - fn sub(self, rhs: T) -> Self::Output { - SymbolicExpression::from(self) - rhs.into() - } -} - -impl Mul for SymbolicVariable -where - T: Into>, -{ - type Output = SymbolicExpression; - - fn mul(self, rhs: T) -> Self::Output { - SymbolicExpression::from(self) * rhs.into() - } -} - -/// An expression over `SymbolicVariable`s. -#[derive(Clone, Debug)] -pub enum SymbolicExpression { - Variable(SymbolicVariable), - IsFirstRow, - IsLastRow, - IsTransition, - Constant(F), - Add { - x: Box, - y: Box, - degree_multiple: usize, - }, - Sub { - x: Box, - y: Box, - degree_multiple: usize, - }, - Neg { - x: Box, - degree_multiple: usize, - }, - Mul { - x: Box, - y: Box, - degree_multiple: usize, - }, -} - -#[inline] -pub fn var(i: usize) -> SymbolicExpression { - SymbolicExpression::from(SymbolicVariable::new(Entry::Main { offset: 0 }, i)) -} - -#[inline] -pub fn preprocessed_var(i: usize) -> SymbolicExpression { - SymbolicExpression::from(SymbolicVariable::new(Entry::Preprocessed { offset: 0 }, i)) -} - -impl SymbolicExpression { - pub fn interpret, Var: Into + Clone>( - &self, - row: &[Var], - preprocessed: Option<&[Var]>, - ) -> Expr { - match self { - Self::Variable(var) => match &var.entry { - Entry::Main { offset: 0 } => row[var.index].clone().into(), - Entry::Preprocessed { offset: 0 } => { - preprocessed.unwrap()[var.index].clone().into() - } - _ => unimplemented!(), - }, - Self::Constant(c) => (*c).into(), - Self::Add { x, y, .. } => { - x.interpret(row, preprocessed) + y.interpret(row, preprocessed) - } - Self::Sub { x, y, .. } => { - x.interpret(row, preprocessed) - y.interpret(row, preprocessed) - } - Self::Neg { x, .. } => -x.interpret(row, preprocessed), - Self::Mul { x, y, .. } => { - x.interpret(row, preprocessed) * y.interpret(row, preprocessed) - } - _ => unimplemented!(), - } - } - - /// Returns the multiple of `n` (the trace length) in this expression's degree. - pub const fn degree_multiple(&self) -> usize { - match self { - Self::Variable(v) => v.degree_multiple(), - Self::IsFirstRow | Self::IsLastRow => 1, - Self::IsTransition | Self::Constant(_) => 0, - Self::Add { - degree_multiple, .. - } - | Self::Sub { - degree_multiple, .. - } - | Self::Neg { - degree_multiple, .. - } - | Self::Mul { - degree_multiple, .. - } => *degree_multiple, - } - } -} - -impl Dup for SymbolicExpression { - #[inline(always)] - fn dup(&self) -> Self { - self.clone() - } -} - -impl Default for SymbolicExpression { - fn default() -> Self { - Self::Constant(F::ZERO) - } -} - -/// Base-field values embed into expressions over any extension of the base -/// field (including the base field itself, via the reflexive -/// `ExtensionField for F`). This single impl replaces both the same-field -/// `From` conversion and the old concrete `Val` → `ExtVal` lifting. -impl> From for SymbolicExpression { - fn from(value: F) -> Self { - Self::Constant(value.into()) - } -} - -impl PrimeCharacteristicRing for SymbolicExpression { - type PrimeSubfield = F::PrimeSubfield; - - const ZERO: Self = Self::Constant(F::ZERO); - const ONE: Self = Self::Constant(F::ONE); - const TWO: Self = Self::Constant(F::TWO); - const NEG_ONE: Self = Self::Constant(F::NEG_ONE); - - #[inline] - fn from_prime_subfield(f: Self::PrimeSubfield) -> Self { - F::from_prime_subfield(f).into() - } -} - -impl> Algebra for SymbolicExpression {} - -impl Algebra> for SymbolicExpression {} - -impl, const N: u64> InjectiveMonomial for SymbolicExpression {} - -impl Add for SymbolicExpression -where - T: Into, -{ - type Output = Self; - - fn add(self, rhs: T) -> Self { - match (self, rhs.into()) { - (Self::Constant(lhs), rhs) if lhs == F::ZERO => rhs, - (lhs, Self::Constant(rhs)) if rhs == F::ZERO => lhs, - (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs + rhs), - (lhs, rhs) => Self::Add { - degree_multiple: lhs.degree_multiple().max(rhs.degree_multiple()), - x: Box::new(lhs), - y: Box::new(rhs), - }, - } - } -} - -impl AddAssign for SymbolicExpression -where - T: Into, -{ - fn add_assign(&mut self, rhs: T) { - *self = self.clone() + rhs.into(); - } -} - -impl Sum for SymbolicExpression -where - T: Into, -{ - fn sum>(iter: I) -> Self { - iter.map(Into::into) - .reduce(|x, y| x + y) - .unwrap_or(Self::ZERO) - } -} - -impl> Sub for SymbolicExpression { - type Output = Self; - - fn sub(self, rhs: T) -> Self { - match (self, rhs.into()) { - (Self::Constant(lhs), rhs) if lhs == F::ZERO => -rhs, - (lhs, Self::Constant(rhs)) if rhs == F::ZERO => lhs, - (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs - rhs), - (lhs, rhs) => Self::Sub { - degree_multiple: lhs.degree_multiple().max(rhs.degree_multiple()), - x: Box::new(lhs), - y: Box::new(rhs), - }, - } - } -} - -impl SubAssign for SymbolicExpression -where - T: Into, -{ - fn sub_assign(&mut self, rhs: T) { - *self = self.clone() - rhs.into(); - } -} - -impl Neg for SymbolicExpression { - type Output = Self; - - fn neg(self) -> Self { - match self { - Self::Constant(c) => Self::Constant(-c), - expr => Self::Neg { - degree_multiple: expr.degree_multiple(), - x: Box::new(expr), - }, - } - } -} - -impl> Mul for SymbolicExpression { - type Output = Self; - - fn mul(self, rhs: T) -> Self { - match (self, rhs.into()) { - (Self::Constant(lhs), rhs) if lhs == F::ONE => rhs, - (lhs, Self::Constant(rhs)) if rhs == F::ONE => lhs, - (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs * rhs), - (lhs, rhs) => Self::Mul { - degree_multiple: lhs.degree_multiple() + rhs.degree_multiple(), - x: Box::new(lhs), - y: Box::new(rhs), - }, - } - } -} - -impl MulAssign for SymbolicExpression -where - T: Into, -{ - fn mul_assign(&mut self, rhs: T) { - *self = self.clone() * rhs.into(); - } -} - -impl> Product for SymbolicExpression { - fn product>(iter: I) -> Self { - iter.map(Into::into) - .reduce(|x, y| x * y) - .unwrap_or(Self::ONE) - } -} - -pub fn get_symbolic_constraints( - air: &A, - preprocessed_width: usize, - stage_1_width: usize, - stage_2_width: usize, - num_public_values: usize, - num_stage_2_public_values: usize, -) -> Vec> -where - F: Field, - EF: ExtensionField, - A: Air>, -{ - let mut builder = SymbolicAirBuilder::new( - preprocessed_width, - stage_1_width, - stage_2_width, - num_public_values, - num_stage_2_public_values, - ); - air.eval(&mut builder); - builder.constraints -} - -pub fn get_max_constraint_degree(constraints: &[SymbolicExpression]) -> usize { - constraints - .iter() - .map(|c| c.degree_multiple()) - .max() - .unwrap_or(0) -} - -/// An `AirBuilder` for evaluating constraints symbolically, and recording them for later use. -/// -/// All variables and expressions are tagged with the extension field `EF`, -/// even those referring to base-field trace columns. The tag only affects the -/// type of embedded constants — using a single tag lets `Expr` and `ExprEF` -/// be the same type, which sidesteps the coherence problems of converting -/// between `SymbolicExpression` and `SymbolicExpression`. -#[derive(Debug)] -pub struct SymbolicAirBuilder> { - preprocessed: RowMajorMatrix>, - stage_1: RowMajorMatrix>, - stage_2: RowMajorMatrix>, - public_values: Vec>, - stage_2_public_values: Vec>, - constraints: Vec>, - _phantom: PhantomData, -} - -impl> SymbolicAirBuilder { - pub(crate) fn new( - preprocessed_width: usize, - stage_1_width: usize, - stage_2_width: usize, - num_public_values: usize, - num_stage_2_public_values: usize, - ) -> Self { - let prep_values = [0, 1] - .into_iter() - .flat_map(|offset| { - (0..preprocessed_width) - .map(move |index| SymbolicVariable::new(Entry::Preprocessed { offset }, index)) - }) - .collect(); - let stage_1_values = [0, 1] - .into_iter() - .flat_map(|offset| { - (0..stage_1_width) - .map(move |index| SymbolicVariable::new(Entry::Main { offset }, index)) - }) - .collect(); - let stage_2_values = [0, 1] - .into_iter() - .flat_map(|offset| { - (0..stage_2_width) - .map(move |index| SymbolicVariable::new(Entry::Stage2 { offset }, index)) - }) - .collect(); - let public_values = (0..num_public_values) - .map(move |index| SymbolicVariable::new(Entry::Public, index)) - .collect(); - let stage_2_public_values = (0..num_stage_2_public_values) - .map(move |index| SymbolicVariable::new(Entry::Stage2Public, index)) - .collect(); - Self { - preprocessed: RowMajorMatrix::new(prep_values, preprocessed_width), - stage_1: RowMajorMatrix::new(stage_1_values, stage_1_width), - stage_2: RowMajorMatrix::new(stage_2_values, stage_2_width), - public_values, - stage_2_public_values, - constraints: vec![], - _phantom: PhantomData, - } - } -} - -impl> AirBuilder for SymbolicAirBuilder { - type F = F; - type Expr = SymbolicExpression; - type Var = SymbolicVariable; - type PreprocessedWindow = RowMajorMatrix; - type MainWindow = RowMajorMatrix; - type PublicVar = SymbolicVariable; - - fn main(&self) -> Self::MainWindow { - self.stage_1.clone() - } - - fn preprocessed(&self) -> &Self::PreprocessedWindow { - &self.preprocessed - } - - fn is_first_row(&self) -> Self::Expr { - SymbolicExpression::IsFirstRow - } - - fn is_last_row(&self) -> Self::Expr { - SymbolicExpression::IsLastRow - } - - /// # Panics - /// This function panics if `size` is not `2`. - fn is_transition_window(&self, size: usize) -> Self::Expr { - if size == 2 { - SymbolicExpression::IsTransition - } else { - panic!("multi-stark only supports a window size of 2") - } - } - - fn assert_zero>(&mut self, x: I) { - self.constraints.push(x.into()); - } - - fn public_values(&self) -> &[Self::PublicVar] { - &self.public_values - } -} - -impl> ExtensionBuilder for SymbolicAirBuilder { - type EF = EF; - type ExprEF = SymbolicExpression; - type VarEF = SymbolicVariable; - - fn assert_zero_ext(&mut self, x: I) - where - I: Into, - { - self.constraints.push(x.into()); - } -} - -impl> TwoStagedBuilder for SymbolicAirBuilder { - type MP = RowMajorMatrix; - - type Stage2PublicVar = Self::VarEF; - - fn stage_2(&self) -> Self::MP { - self.stage_2.clone() - } - - fn stage_2_public_values(&self) -> &[Self::Stage2PublicVar] { - &self.stage_2_public_values - } -} diff --git a/src/eval.rs b/src/eval.rs new file mode 100644 index 0000000..87faec8 --- /dev/null +++ b/src/eval.rs @@ -0,0 +1,212 @@ +//! Evaluation. +//! +//! Compiled circuits are evaluated by a single dense forward sweep: one +//! buffer slot per node, children always at smaller indices. The frontend +//! trees also get a direct recursive evaluator, used as the reference in +//! tests (with genuine extension-field arithmetic for `ExtExpr`). + +use p3_field::{Algebra, Field}; + +use crate::expr::{ColRef, Expr, ExtExpr, RowOffset, Source}; +use crate::graph::{ConstraintGraph, ExtensionParams, Node, NodeId}; +use crate::lookup::Lookup; + +/// Concrete leaf values for a sweep, in the working type `W`: the two-row +/// window of each trace matrix, the public inputs and the selector values. +/// +/// `W` is the type the sweep computes in — `F` for the debug/witness paths, +/// `PackedVal` on the quotient domain (prover), `EF` at ζ (verifier). The +/// compiled node constants (base field `F`) embed into `W` via `Algebra`. +#[derive(Clone, Debug)] +pub struct VarValues<'a, W> { + /// `[current_row, next_row]` of the preprocessed trace. + pub preprocessed: [&'a [W]; 2], + /// `[current_row, next_row]` of the main trace. + pub main: [&'a [W]; 2], + /// `[current_row, next_row]` of the stage-2 trace (flattened base columns). + pub stage2: [&'a [W]; 2], + /// Public inputs, as base-field coordinates embedded into `W`. + pub publics: &'a [W], + pub is_first_row: W, + pub is_last_row: W, + pub is_transition: W, +} + +impl VarValues<'_, W> { + fn var(&self, col: &ColRef) -> W { + let rows = match col.source { + Source::Preprocessed => &self.preprocessed, + Source::Main => &self.main, + Source::Stage2 => &self.stage2, + }; + let row = match col.offset { + RowOffset::Current => rows[0], + RowOffset::Next => rows[1], + }; + row[col.index as usize] + } +} + +impl ConstraintGraph { + /// Dense forward sweep over the whole node vector, in working type `W`; + /// fills `buf` with one value per node. + pub fn sweep + Copy>(&self, values: &VarValues<'_, W>, buf: &mut Vec) { + self.sweep_range(values, buf, self.nodes.len()); + } + + /// Sweeps only the lookup prefix (partial evaluation for the lookup + /// witness). + pub fn sweep_lookup_prefix + Copy>( + &self, + values: &VarValues<'_, W>, + buf: &mut Vec, + ) { + self.sweep_range(values, buf, self.lookup_prefix_len); + } + + fn sweep_range + Copy>( + &self, + values: &VarValues<'_, W>, + buf: &mut Vec, + len: usize, + ) { + buf.clear(); + buf.reserve(len); + // Raw pointer into the reserved storage; captured by the `child` + // reader (a `*mut` is `Copy`, so it borrows nothing) and used for the + // writes, sidestepping the borrow checker in the hot loop. + let ptr = buf.spare_capacity_mut().as_mut_ptr(); + // SAFETY throughout: the node vector is topologically ordered — every + // child's index is strictly less than its parent's — so when node `i` + // is processed, every index it reads (`< i`) has already been written + // and initialized; the storage has room for `len` values (reserved, + // no realloc); and `W: Copy` means the reads take copies with no drop + // obligations. This is the constraint-evaluation hot loop (quotient + // domain, per packet), so eliding per-node bounds/capacity checks and + // the shared-subexpression recomputation matters. + let child = |id: NodeId| unsafe { (*ptr.add(id.index())).assume_init_read() }; + for i in 0..len { + let node = unsafe { self.nodes.get_unchecked(i) }; + let value = match node { + Node::Const(c) => (*c).into(), + Node::Var(col) => values.var(col), + Node::Public(idx) => values.publics[*idx as usize], + Node::IsFirstRow => values.is_first_row, + Node::IsLastRow => values.is_last_row, + Node::IsTransition => values.is_transition, + Node::Add(a, b) => child(*a) + child(*b), + Node::Sub(a, b) => child(*a) - child(*b), + Node::Mul(a, b) => child(*a) * child(*b), + Node::Neg(a) => -child(*a), + }; + unsafe { (*ptr.add(i)).write(value) }; + } + // SAFETY: the loop initialized exactly `len` values. + unsafe { buf.set_len(len) }; + } + + /// The constraint values, read off a full sweep buffer. + pub fn constraint_values(&self, buf: &[W]) -> Vec { + self.zeros.iter().map(|z| buf[z.index()]).collect() + } + + /// The concrete lookup values, read off a (prefix or full) sweep buffer. + pub fn lookup_values(&self, buf: &[W]) -> Vec> { + self.lookups + .iter() + .map(|lookup| Lookup { + multiplicity: buf[lookup.multiplicity.index()], + args: lookup.args.iter().map(|a| buf[a.index()]).collect(), + }) + .collect() + } + + /// Convenience: sweep and return the constraint values. + pub fn evaluate_constraints + Copy>(&self, values: &VarValues<'_, W>) -> Vec { + let mut buf = Vec::new(); + self.sweep(values, &mut buf); + self.constraint_values(&buf) + } +} + +/// Reference evaluation of a frontend base expression (recursive). +pub fn eval_expr(expr: &Expr, values: &VarValues<'_, F>) -> F { + match expr { + Expr::Const(c) => *c, + Expr::Var(col) => values.var(col), + Expr::Public(i) => values.publics[*i as usize], + Expr::IsFirstRow => values.is_first_row, + Expr::IsLastRow => values.is_last_row, + Expr::IsTransition => values.is_transition, + Expr::Add(a, b) => eval_expr(a, values) + eval_expr(b, values), + Expr::Sub(a, b) => eval_expr(a, values) - eval_expr(b, values), + Expr::Mul(a, b) => eval_expr(a, values) * eval_expr(b, values), + Expr::Neg(a) => -eval_expr(a, values), + } +} + +/// Reference evaluation of a frontend extension expression, in genuine +/// extension-field arithmetic: coordinates as `Vec` of length D, +/// products reduced mod `X^D − W` (schoolbook — deliberately independent +/// of the compiled Karatsuba path). +pub fn eval_ext_expr( + expr: &ExtExpr, + values: &VarValues<'_, F>, + params: &ExtensionParams, +) -> Vec { + let d = params.degree; + match expr { + ExtExpr::Coords(coords) => { + assert_eq!(coords.len(), d, "reference eval: bad Coords length"); + coords.iter().map(|c| eval_expr(c, values)).collect() + } + ExtExpr::Base(base) => { + let mut out = vec![F::ZERO; d]; + out[0] = eval_expr(base, values); + out + } + ExtExpr::Add(a, b) => { + let a = eval_ext_expr(a, values, params); + let b = eval_ext_expr(b, values, params); + a.into_iter().zip(b).map(|(x, y)| x + y).collect() + } + ExtExpr::Sub(a, b) => { + let a = eval_ext_expr(a, values, params); + let b = eval_ext_expr(b, values, params); + a.into_iter().zip(b).map(|(x, y)| x - y).collect() + } + ExtExpr::Neg(a) => { + let a = eval_ext_expr(a, values, params); + a.into_iter().map(|x| -x).collect() + } + ExtExpr::Mul(a, b) => { + let a = eval_ext_expr(a, values, params); + let b = eval_ext_expr(b, values, params); + let mut out = vec![F::ZERO; d]; + for i in 0..d { + for j in 0..d { + let term = a[i] * b[j]; + if i + j < d { + out[i + j] += term; + } else { + out[i + j - d] += params.w * term; + } + } + } + out + } + } +} + +/// Checks the topological invariant: every node's children have strictly +/// smaller indices. Test helper, but useful as a deserialization check too. +pub fn check_topological_order(circuit: &ConstraintGraph) -> bool { + circuit.nodes.iter().enumerate().all(|(i, node)| { + let ok = |child: NodeId| child.index() < i; + match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => ok(a) && ok(b), + Node::Neg(a) => ok(a), + _ => true, + } + }) +} diff --git a/src/expr.rs b/src/expr.rs new file mode 100644 index 0000000..2cb9744 --- /dev/null +++ b/src/expr.rs @@ -0,0 +1,305 @@ +//! Frontend expression trees. +//! +//! These exist only while a circuit is being described; compilation (see +//! [`crate::circuit`]) flattens them into the interned format everything +//! else consumes. Operators fold constants as they go. + +use std::ops::{Add, Mul, Neg, Sub}; + +use p3_field::Field; + +use crate::lookup::Lookup; + +/// Which committed matrix a column variable refers to. +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +pub enum Source { + Preprocessed, + Main, + /// The stage-2 trace in its flattened base-column layout. + Stage2, +} + +/// The evaluation window: only the current and next row are addressable. +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +pub enum RowOffset { + Current, + Next, +} + +/// A column of a committed matrix, at the current or next row. +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +pub struct ColRef { + pub source: Source, + pub offset: RowOffset, + pub index: u32, +} + +/// A base-field expression. +#[derive(Clone, Debug)] +pub enum Expr { + Const(F), + Var(ColRef), + /// A public input, as a base-field coordinate. + Public(u32), + IsFirstRow, + IsLastRow, + IsTransition, + Add(Box, Box), + Sub(Box, Box), + Mul(Box, Box), + Neg(Box), +} + +/// An extension-field expression. The primitive is `Coords`: an array of +/// D base-field coordinates representing `Σ_j coord_j · b_j`. +#[derive(Clone, Debug)] +pub enum ExtExpr { + /// The D-tuple primitive; length is checked at compile time. + Coords(Vec>), + /// Implicit embedding of a base expression, so mixed arithmetic like + /// `m * inv` needs no explicit lift. + Base(Box>), + Add(Box, Box), + Sub(Box, Box), + Mul(Box, Box), + Neg(Box), +} + +/// What the constraint compiler consumes, per circuit: the authored +/// [`crate::system::CircuitInputs`] with the derived widths filled in and +/// the synthesized lookup constraints appended. Built internally by +/// `System::new`; not part of the public API. +#[derive(Clone, Debug)] +pub(crate) struct CircuitSpec { + pub main_width: usize, + pub preprocessed_width: usize, + /// Width of the stage-2 trace in flattened base columns. + pub stage2_width: usize, + /// Number of public inputs (base-field coordinates). + pub num_publics: usize, + /// Base-field constraints: each must evaluate to 0 on every row. + pub constraints: Vec>, + /// Extension-field constraints, same vanishing semantics. + pub ext_constraints: Vec>, + /// Lookup specifications. + pub lookups: Vec>>, +} + +impl Default for CircuitSpec { + fn default() -> Self { + Self { + main_width: 0, + preprocessed_width: 0, + stage2_width: 0, + num_publics: 0, + constraints: vec![], + ext_constraints: vec![], + lookups: vec![], + } + } +} + +impl Expr { + pub fn constant(value: F) -> Self { + Self::Const(value) + } + + pub fn var(source: Source, offset: RowOffset, index: u32) -> Self { + Self::Var(ColRef { + source, + offset, + index, + }) + } + + pub fn main(index: u32) -> Self { + Self::var(Source::Main, RowOffset::Current, index) + } + + pub fn main_next(index: u32) -> Self { + Self::var(Source::Main, RowOffset::Next, index) + } + + pub fn preprocessed(index: u32) -> Self { + Self::var(Source::Preprocessed, RowOffset::Current, index) + } + + pub fn preprocessed_next(index: u32) -> Self { + Self::var(Source::Preprocessed, RowOffset::Next, index) + } + + pub fn stage2(index: u32) -> Self { + Self::var(Source::Stage2, RowOffset::Current, index) + } + + pub fn stage2_next(index: u32) -> Self { + Self::var(Source::Stage2, RowOffset::Next, index) + } + + pub fn public(index: u32) -> Self { + Self::Public(index) + } +} + +impl ExtExpr { + /// An extension constant from its D basis coordinates. + pub fn constant(coords: Vec) -> Self { + Self::Coords(coords.into_iter().map(Expr::Const).collect()) + } + + /// The extension value stored in stage-2 slot `slot`: D consecutive + /// flattened base columns starting at `slot * d`. + pub fn stage2(slot: u32, d: u32, offset: RowOffset) -> Self { + Self::Coords( + (0..d) + .map(|j| Expr::var(Source::Stage2, offset, slot * d + j)) + .collect(), + ) + } + + /// The extension public value `k`: D consecutive base coordinates + /// starting at `k * d`. + pub fn public(k: u32, d: u32) -> Self { + Self::Coords((0..d).map(|j| Expr::Public(k * d + j)).collect()) + } +} + +impl From for Expr { + fn from(value: F) -> Self { + Self::Const(value) + } +} + +impl From> for ExtExpr { + fn from(value: Expr) -> Self { + Self::Base(Box::new(value)) + } +} + +impl Add for Expr { + type Output = Self; + + fn add(self, rhs: Self) -> Self { + match (self, rhs) { + (Self::Const(a), Self::Const(b)) => Self::Const(a + b), + (Self::Const(z), x) | (x, Self::Const(z)) if z == F::ZERO => x, + (a, b) => Self::Add(Box::new(a), Box::new(b)), + } + } +} + +impl Sub for Expr { + type Output = Self; + + fn sub(self, rhs: Self) -> Self { + match (self, rhs) { + (Self::Const(a), Self::Const(b)) => Self::Const(a - b), + (x, Self::Const(z)) if z == F::ZERO => x, + (Self::Const(z), x) if z == F::ZERO => -x, + (a, b) => Self::Sub(Box::new(a), Box::new(b)), + } + } +} + +impl Mul for Expr { + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + match (self, rhs) { + (Self::Const(a), Self::Const(b)) => Self::Const(a * b), + (Self::Const(z), _) | (_, Self::Const(z)) if z == F::ZERO => Self::Const(F::ZERO), + (Self::Const(o), x) | (x, Self::Const(o)) if o == F::ONE => x, + (a, b) => Self::Mul(Box::new(a), Box::new(b)), + } + } +} + +impl Neg for Expr { + type Output = Self; + + fn neg(self) -> Self { + match self { + Self::Const(c) => Self::Const(-c), + Self::Neg(inner) => *inner, + e => Self::Neg(Box::new(e)), + } + } +} + +impl Add for ExtExpr { + type Output = Self; + + fn add(self, rhs: Self) -> Self { + Self::Add(Box::new(self), Box::new(rhs)) + } +} + +impl Sub for ExtExpr { + type Output = Self; + + fn sub(self, rhs: Self) -> Self { + Self::Sub(Box::new(self), Box::new(rhs)) + } +} + +impl Mul for ExtExpr { + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + Self::Mul(Box::new(self), Box::new(rhs)) + } +} + +impl Neg for ExtExpr { + type Output = Self; + + fn neg(self) -> Self { + Self::Neg(Box::new(self)) + } +} + +/// Mixed base/extension operators, routed through `ExtExpr::Base` so a +/// lone base value combines with extension values without a visible lift. +macro_rules! mixed_ops { + ($($trait:ident :: $method:ident),*) => { + $( + impl $trait> for Expr { + type Output = ExtExpr; + + fn $method(self, rhs: ExtExpr) -> ExtExpr { + ExtExpr::from(self).$method(rhs) + } + } + + impl $trait> for ExtExpr { + type Output = ExtExpr; + + fn $method(self, rhs: Expr) -> ExtExpr { + self.$method(ExtExpr::from(rhs)) + } + } + )* + }; +} + +mixed_ops!(Add::add, Sub::sub, Mul::mul); + +impl ExtExpr { + /// True if the expression is built solely from base embeddings — its + /// expansion would be one real constraint plus D−1 trivial zeros, so + /// it should have been stated as a base constraint. + pub fn is_purely_base(&self) -> bool { + match self { + Self::Coords(_) => false, + Self::Base(_) => true, + Self::Add(a, b) | Self::Sub(a, b) | Self::Mul(a, b) => { + a.is_purely_base() && b.is_purely_base() + } + Self::Neg(a) => a.is_purely_base(), + } + } +} + +// NOTE: deep frontend trees drop recursively; the production version needs +// an iterative `Drop` (which in turn requires the operators to avoid +// by-value destructuring, E0509). Out of scope for this lab. diff --git a/src/graph.rs b/src/graph.rs new file mode 100644 index 0000000..25826ac --- /dev/null +++ b/src/graph.rs @@ -0,0 +1,508 @@ +//! Compilation: frontend trees → flat, interned, base-only node vector. +//! +//! Properties established here (see the design doc): +//! - children by index, topological order (children before parents); +//! - hash-consing with commutative normalization: index equality ⇔ +//! structural equality; +//! - constant folding, so the zero coordinates of base embeddings vanish; +//! - coordinate expansion of extension constraints (scalar detection, +//! Karatsuba for D=2, schoolbook otherwise); +//! - lookup expressions interned first, so they occupy a prefix of the +//! vector (partial evaluation for the lookup witness); +//! - per-node degree multiples. + +use std::collections::HashMap; + +use p3_field::Field; + +use crate::expr::{CircuitSpec, ColRef, Expr, ExtExpr, Source}; +use crate::lookup::Lookup; + +/// Index of a node in the compiled vector. +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)] +pub struct NodeId(pub u32); + +impl NodeId { + #[inline] + pub fn index(self) -> usize { + self.0 as usize + } +} + +/// A compiled expression node. Base-field only: extension constraints +/// have been coordinate-expanded away by the time nodes exist. +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +pub enum Node { + Const(F), + Var(ColRef), + Public(u32), + IsFirstRow, + IsLastRow, + IsTransition, + Add(NodeId, NodeId), + Sub(NodeId, NodeId), + Mul(NodeId, NodeId), + Neg(NodeId), +} + +/// Binomial extension parameters: `X^degree = w`. +#[derive(Copy, Clone, Debug)] +pub struct ExtensionParams { + pub degree: usize, + pub w: F, + /// Emit the 3-multiplication Karatsuba form for D=2 products + /// (matching native extension-multiplication cost). Schoolbook + /// otherwise; both are exposed so tests can cross-check them. + pub karatsuba: bool, +} + +/// A compiled constraint graph: the hash-consed DAG of expression nodes +/// shared by all constraints and lookups of one circuit. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ConstraintGraph { + /// The unique expression vector; all constraints and lookups share it. + pub nodes: Vec>, + /// Degree multiple of each node. + pub degrees: Vec, + /// Constraint roots, in canonical order: base constraints, then each + /// extension constraint's D coordinate roots. + pub zeros: Vec, + /// Compiled lookups; all their nodes live in the prefix. + pub lookups: Vec>, + /// `nodes[..lookup_prefix_len]` covers everything reachable from the + /// lookup expressions. + pub lookup_prefix_len: usize, + pub max_constraint_degree: u32, +} + +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum CompileError { + CoordsLength { + constraint: usize, + expected: usize, + got: usize, + }, + ColumnOutOfRange { + source: Source, + index: u32, + width: usize, + }, + PublicOutOfRange { + index: u32, + count: usize, + }, + /// Stage-2 columns are extension coordinates; base constraints and + /// lookup expressions must not touch them. + Stage2InBaseContext, + /// A fully base-field extension constraint: state it as a base + /// constraint instead. + PurelyBaseExtConstraint { + constraint: usize, + }, + /// A constraint (or extension-constraint coordinate) that compiles to + /// a nonzero constant: it can never vanish, so the system would be + /// unprovable. `coordinate` is `None` for a base constraint, `Some(k)` + /// for coordinate `k` of extension constraint `constraint`. + UnsatisfiableConstant { + constraint: usize, + coordinate: Option, + }, +} + +/// Compiles a circuit spec. Interning order: lookups first (they form the +/// prefix), then base constraints, then expanded extension constraints. +/// +/// Constraint roots are canonicalized: a root that compiles to a constant +/// is dropped if zero (vacuous) or rejected if nonzero (unsatisfiable), +/// and the survivors are sorted by node id and deduplicated. The compiled +/// constraint set is therefore independent of the order and multiplicity +/// in which constraints were authored. +pub(crate) fn compile( + spec: &CircuitSpec, + params: &ExtensionParams, +) -> Result, CompileError> { + let mut interner = Interner::new(); + + let mut lookups = Vec::with_capacity(spec.lookups.len()); + for lookup in &spec.lookups { + let multiplicity = interner.compile_expr(&lookup.multiplicity, spec, false)?; + let args = lookup + .args + .iter() + .map(|arg| interner.compile_expr(arg, spec, false)) + .collect::, _>>()?; + lookups.push(Lookup { multiplicity, args }); + } + let lookup_prefix_len = interner.nodes.len(); + + let mut zeros = Vec::new(); + for (i, constraint) in spec.constraints.iter().enumerate() { + let root = interner.compile_expr(constraint, spec, false)?; + record_zero(&interner, &mut zeros, root, i, None)?; + } + for (i, constraint) in spec.ext_constraints.iter().enumerate() { + if constraint.is_purely_base() { + return Err(CompileError::PurelyBaseExtConstraint { constraint: i }); + } + let coords = interner.expand_ext(constraint, spec, params, i)?; + for (coord, root) in coords.into_iter().enumerate() { + record_zero(&interner, &mut zeros, root, i, Some(coord))?; + } + } + // Canonicalize: sort by node id and drop duplicate roots. Index + // equality is structural equality, so this deduplicates constraints + // that became equal under interning (e.g. `a + b` and `b + a`). + zeros.sort_unstable(); + zeros.dedup(); + + // Bottom-up interning guarantees children precede parents. The dense + // evaluator's unchecked sweep (`ConstraintGraph::sweep_range`) relies on + // this, so pin it here. + debug_assert!( + interner + .nodes + .iter() + .enumerate() + .all(|(i, node)| match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => + a.index() < i && b.index() < i, + Node::Neg(a) => a.index() < i, + _ => true, + }), + "compiled nodes must be topologically ordered" + ); + + let max_constraint_degree = zeros + .iter() + .map(|z| interner.degrees[z.index()]) + .max() + .unwrap_or(0); + Ok(ConstraintGraph { + nodes: interner.nodes, + degrees: interner.degrees, + zeros, + lookups, + lookup_prefix_len, + max_constraint_degree, + }) +} + +/// Records a constraint root, handling the constant cases: a nonzero +/// constant is rejected (unsatisfiable), a zero constant is dropped +/// (vacuous), anything else is appended. +fn record_zero( + interner: &Interner, + zeros: &mut Vec, + root: NodeId, + constraint: usize, + coordinate: Option, +) -> Result<(), CompileError> { + match interner.as_const(root) { + Some(c) if c == F::ZERO => {} + Some(_) => { + return Err(CompileError::UnsatisfiableConstant { + constraint, + coordinate, + }); + } + None => zeros.push(root), + } + Ok(()) +} + +/// Bottom-up interner. Children are canonical ids before a parent is +/// interned, so the map lookup is O(1) and sharing is complete. +struct Interner { + nodes: Vec>, + degrees: Vec, + map: HashMap, NodeId>, +} + +impl Interner { + fn new() -> Self { + Self { + nodes: vec![], + degrees: vec![], + map: HashMap::new(), + } + } + + fn intern(&mut self, node: Node) -> NodeId { + if let Some(&id) = self.map.get(&node) { + return id; + } + let id = NodeId(u32::try_from(self.nodes.len()).expect("node count exceeds u32")); + let degree = self.degree_of(&node); + self.nodes.push(node); + self.degrees.push(degree); + self.map.insert(node, id); + id + } + + fn degree_of(&self, node: &Node) -> u32 { + match *node { + Node::Const(_) | Node::Public(_) | Node::IsTransition => 0, + Node::Var(_) | Node::IsFirstRow | Node::IsLastRow => 1, + Node::Add(a, b) | Node::Sub(a, b) => { + self.degrees[a.index()].max(self.degrees[b.index()]) + } + Node::Mul(a, b) => self.degrees[a.index()] + self.degrees[b.index()], + Node::Neg(a) => self.degrees[a.index()], + } + } + + fn as_const(&self, id: NodeId) -> Option { + match self.nodes[id.index()] { + Node::Const(c) => Some(c), + _ => None, + } + } + + fn constant(&mut self, value: F) -> NodeId { + self.intern(Node::Const(value)) + } + + fn add(&mut self, a: NodeId, b: NodeId) -> NodeId { + match (self.as_const(a), self.as_const(b)) { + (Some(x), Some(y)) => return self.constant(x + y), + (Some(x), None) if x == F::ZERO => return b, + (None, Some(y)) if y == F::ZERO => return a, + _ => {} + } + // Commutative normalization: sorted operand ids. + let (a, b) = if a <= b { (a, b) } else { (b, a) }; + self.intern(Node::Add(a, b)) + } + + fn sub(&mut self, a: NodeId, b: NodeId) -> NodeId { + if a == b { + return self.constant(F::ZERO); + } + match (self.as_const(a), self.as_const(b)) { + (Some(x), Some(y)) => return self.constant(x - y), + (None, Some(y)) if y == F::ZERO => return a, + (Some(x), None) if x == F::ZERO => return self.neg(b), + _ => {} + } + self.intern(Node::Sub(a, b)) + } + + fn mul(&mut self, a: NodeId, b: NodeId) -> NodeId { + match (self.as_const(a), self.as_const(b)) { + (Some(x), Some(y)) => return self.constant(x * y), + (Some(x), None) => { + if x == F::ZERO { + return a; + } + if x == F::ONE { + return b; + } + } + (None, Some(y)) => { + if y == F::ZERO { + return b; + } + if y == F::ONE { + return a; + } + } + _ => {} + } + // Commutative normalization: sorted operand ids. + let (a, b) = if a <= b { (a, b) } else { (b, a) }; + self.intern(Node::Mul(a, b)) + } + + fn neg(&mut self, a: NodeId) -> NodeId { + if let Some(x) = self.as_const(a) { + return self.constant(-x); + } + if let Node::Neg(inner) = self.nodes[a.index()] { + return inner; + } + self.intern(Node::Neg(a)) + } + + fn compile_expr( + &mut self, + expr: &Expr, + spec: &CircuitSpec, + allow_stage2: bool, + ) -> Result { + Ok(match expr { + Expr::Const(c) => self.constant(*c), + Expr::Var(col) => { + let width = match col.source { + Source::Preprocessed => spec.preprocessed_width, + Source::Main => spec.main_width, + Source::Stage2 => { + if !allow_stage2 { + return Err(CompileError::Stage2InBaseContext); + } + spec.stage2_width + } + }; + if col.index as usize >= width { + return Err(CompileError::ColumnOutOfRange { + source: col.source, + index: col.index, + width, + }); + } + self.intern(Node::Var(*col)) + } + Expr::Public(index) => { + if *index as usize >= spec.num_publics { + return Err(CompileError::PublicOutOfRange { + index: *index, + count: spec.num_publics, + }); + } + self.intern(Node::Public(*index)) + } + Expr::IsFirstRow => self.intern(Node::IsFirstRow), + Expr::IsLastRow => self.intern(Node::IsLastRow), + Expr::IsTransition => self.intern(Node::IsTransition), + Expr::Add(a, b) => { + let a = self.compile_expr(a, spec, allow_stage2)?; + let b = self.compile_expr(b, spec, allow_stage2)?; + self.add(a, b) + } + Expr::Sub(a, b) => { + let a = self.compile_expr(a, spec, allow_stage2)?; + let b = self.compile_expr(b, spec, allow_stage2)?; + self.sub(a, b) + } + Expr::Mul(a, b) => { + let a = self.compile_expr(a, spec, allow_stage2)?; + let b = self.compile_expr(b, spec, allow_stage2)?; + self.mul(a, b) + } + Expr::Neg(a) => { + let a = self.compile_expr(a, spec, allow_stage2)?; + self.neg(a) + } + }) + } + + /// Coordinate expansion: returns the D coordinates of an extension + /// expression as interned base nodes. + fn expand_ext( + &mut self, + expr: &ExtExpr, + spec: &CircuitSpec, + params: &ExtensionParams, + constraint: usize, + ) -> Result, CompileError> { + let d = params.degree; + Ok(match expr { + ExtExpr::Coords(coords) => { + if coords.len() != d { + return Err(CompileError::CoordsLength { + constraint, + expected: d, + got: coords.len(), + }); + } + coords + .iter() + .map(|c| self.compile_expr(c, spec, true)) + .collect::, _>>()? + } + ExtExpr::Base(base) => { + let zero = self.constant(F::ZERO); + let mut coords = vec![zero; d]; + coords[0] = self.compile_expr(base, spec, true)?; + coords + } + ExtExpr::Add(a, b) => { + let a = self.expand_ext(a, spec, params, constraint)?; + let b = self.expand_ext(b, spec, params, constraint)?; + (0..d).map(|k| self.add(a[k], b[k])).collect() + } + ExtExpr::Sub(a, b) => { + let a = self.expand_ext(a, spec, params, constraint)?; + let b = self.expand_ext(b, spec, params, constraint)?; + (0..d).map(|k| self.sub(a[k], b[k])).collect() + } + ExtExpr::Neg(a) => { + let a = self.expand_ext(a, spec, params, constraint)?; + a.into_iter().map(|c| self.neg(c)).collect() + } + ExtExpr::Mul(a, b) => { + let a = self.expand_ext(a, spec, params, constraint)?; + let b = self.expand_ext(b, spec, params, constraint)?; + self.ext_mul(&a, &b, params) + } + }) + } + + /// True if all coordinates but the first are the interned zero — the + /// shape of an embedded base value. + fn is_scalar(&self, coords: &[NodeId]) -> bool { + coords[1..] + .iter() + .all(|&c| self.as_const(c) == Some(F::ZERO)) + } + + fn ext_mul(&mut self, a: &[NodeId], b: &[NodeId], params: &ExtensionParams) -> Vec { + let d = params.degree; + // Scalar structure is preserved by construction: an embedded base + // value multiplies coordinate-wise, D base multiplications. + if self.is_scalar(a) { + return b.iter().map(|&bk| self.mul(a[0], bk)).collect(); + } + if self.is_scalar(b) { + return a.iter().map(|&ak| self.mul(b[0], ak)).collect(); + } + if d == 2 && params.karatsuba { + // (a0 + a1 X)(b0 + b1 X) mod (X² − W), 3 multiplications: + // c0 = a0 b0 + W · a1 b1 + // c1 = (a0 + a1)(b0 + b1) − a0 b0 − a1 b1 + let p0 = self.mul(a[0], b[0]); + let p1 = self.mul(a[1], b[1]); + let sa = self.add(a[0], a[1]); + let sb = self.add(b[0], b[1]); + let s = self.mul(sa, sb); + let w = self.constant(params.w); + let wp1 = self.mul(w, p1); + let c0 = self.add(p0, wp1); + let t = self.sub(s, p0); + let c1 = self.sub(t, p1); + return vec![c0, c1]; + } + // Schoolbook: c_k = Σ_{i+j=k} a_i b_j + W · Σ_{i+j=k+D} a_i b_j. + let w = self.constant(params.w); + (0..d) + .map(|k| { + let mut low: Option = None; + let mut high: Option = None; + for (i, &ai) in a.iter().enumerate() { + for (j, &bj) in b.iter().enumerate() { + if i + j == k { + let term = self.mul(ai, bj); + low = Some(match low { + Some(acc) => self.add(acc, term), + None => term, + }); + } else if i + j == k + d { + let term = self.mul(ai, bj); + high = Some(match high { + Some(acc) => self.add(acc, term), + None => term, + }); + } + } + } + let low = low.expect("k < d always has low terms"); + match high { + Some(high) => { + let high = self.mul(w, high); + self.add(low, high) + } + None => low, + } + }) + .collect() + } +} diff --git a/src/lib.rs b/src/lib.rs index 1262998..0df70f6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,9 @@ -pub mod builder; pub mod config; +pub mod eval; +pub mod expr; +pub mod graph; pub mod lookup; +pub mod p3_adapter; pub mod prover; pub mod system; #[cfg(test)] diff --git a/src/lookup.rs b/src/lookup.rs index 7fb7fa3..d27974f 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -1,45 +1,62 @@ -use p3_air::{Air, BaseAir, ExtensionBuilder, WindowAccess}; +//! Lookups: the [`Lookup`] type, the logUp stage-2 constraint synthesis, and +//! the concrete lookup witness ([`LookupValues`]) plus its stage-2 trace +//! construction. +//! +//! Synthesis is a frontend-to-frontend function: given the lookups, it +//! produces the extension-field constraints for the running-accumulator +//! argument as ordinary [`ExtExpr`] data. [`crate::system::System::new`] +//! appends these to a circuit's `ext_constraints` before compilation, so they +//! are interned, coordinate-expanded and canonicalized like any other +//! constraint. +//! +//! # Layout conventions (owned here) +//! +//! - **Publics** (each an extension value, stored as `d` base coordinates): +//! slot 0 = β (lookup challenge), 1 = γ (fingerprint challenge), +//! 2 = current accumulator, 3 = next accumulator. So `num_publics = 4·d`. +//! - **Stage-2 columns** (extension values, flattened to base): slot 0 = the +//! running accumulator, slots `1..=L` = the message inverse of each lookup, +//! in lookup order. So `stage2_width = (1 + L)·d` base columns. + use p3_field::{ExtensionField, Field, PrimeCharacteristicRing, batch_multiplicative_inverse}; -use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use p3_matrix::dense::RowMajorMatrix; use p3_maybe_rayon::prelude::*; -use crate::builder::{TwoStagedBuilder, symbolic::SymbolicExpression}; - -/// Each circuit is required to have 4 arguments for the second stage. Namely, -/// the lookup challenge, fingerprint challenge, current accumulator and next -/// accumulator. -pub const LOOKUP_PUBLIC_SIZE: usize = 4; +use crate::expr::{Expr, ExtExpr, RowOffset}; -#[derive(Clone)] -pub struct Lookup { - pub multiplicity: Expr, - pub args: Vec, +/// A lookup: a multiplicity and a vector of arguments. `E` is a frontend +/// expression in a [`crate::system::CircuitInputs`] and a node id once +/// compiled. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct Lookup { + pub multiplicity: E, + pub args: Vec, } -impl Lookup { +impl Lookup { /// Returns a [`Lookup`] with multiplicity zero and no arguments. #[inline] pub fn empty() -> Self where - Expr: PrimeCharacteristicRing, + E: PrimeCharacteristicRing, { Self { - multiplicity: Expr::ZERO, + multiplicity: E::ZERO, args: vec![], } } /// "Pushing" has the semantics of adding a claim to the claim set. #[inline] - pub fn push(multiplicity: Expr, args: Vec) -> Self { + pub fn push(multiplicity: E, args: Vec) -> Self { Self { multiplicity, args } } /// "Pulling" has the semantics of removing a claim from the claim set. #[inline] - pub fn pull(multiplicity: Expr, args: Vec) -> Self + pub fn pull(multiplicity: E, args: Vec) -> Self where - Expr: std::ops::Neg, + E: std::ops::Neg, { Self { multiplicity: -multiplicity, @@ -48,27 +65,68 @@ impl Lookup { } } -pub struct LookupAir { - pub inner_air: A, - pub lookups: Vec>>, - pub preprocessed: Option>, +/// Number of extension-valued public inputs the lookup argument uses: +/// β, γ, current accumulator, next accumulator. +pub const LOOKUP_PUBLIC_SIZE: usize = 4; + +/// Public-input width (base coordinates) for a circuit whose extension +/// degree is `d`. +pub fn num_publics(d: usize) -> usize { + LOOKUP_PUBLIC_SIZE * d } -impl, F: Field> LookupAir { - pub fn new(inner_air: A, lookups: Vec>>) -> Self { - let preprocessed = inner_air.preprocessed_trace(); - Self { - inner_air, - lookups, - preprocessed, +/// Stage-2 trace width (flattened base columns) for `num_lookups` lookups +/// at extension degree `d`: one accumulator column plus one inverse column +/// per lookup. +pub fn stage2_width(num_lookups: usize, d: usize) -> usize { + (1 + num_lookups) * d +} + +/// Builds the logUp stage-2 constraints for `lookups` at extension degree +/// `d`. Order: the per-lookup message/inverse constraints in lookup order, +/// then the first-row, transition, and last-row accumulator constraints. +pub fn synthesize_lookups(lookups: &[Lookup>], d: usize) -> Vec> { + let d = u32::try_from(d).expect("extension degree exceeds u32"); + let beta = ExtExpr::public(0, d); + let gamma = ExtExpr::public(1, d); + let acc = ExtExpr::public(2, d); + let next_acc = ExtExpr::public(3, d); + let acc_col = ExtExpr::stage2(0, d, RowOffset::Current); + let next_acc_col = ExtExpr::stage2(0, d, RowOffset::Next); + + let mut constraints = Vec::with_capacity(lookups.len() + 3); + // acc_expr = acc_col + Σ_j multiplicity_j · inv_j, built once and reused + // by the transition and last-row constraints (the interner shares it). + let mut acc_expr = acc_col.clone(); + for (j, lookup) in lookups.iter().enumerate() { + let slot = 1 + u32::try_from(j).expect("lookup count exceeds u32"); + let inv = ExtExpr::stage2(slot, d, RowOffset::Current); + + // fingerprint = Σ_i args[i] · γ^i, via Horner over the reversed args. + let mut coeffs = lookup.args.iter().rev(); + let mut fingerprint = match coeffs.next() { + Some(arg) => ExtExpr::from(arg.clone()), + None => ExtExpr::from(Expr::constant(F::ZERO)), + }; + for arg in coeffs { + fingerprint = fingerprint * gamma.clone() + arg.clone(); } - } - /// One column for the accumulator and one column for the inverse of the - /// message associated with each lookup. - pub fn stage_2_width(&self) -> usize { - 1 + self.lookups.len() + // message = β + fingerprint; constraint: message · inv − 1 = 0. + let message = beta.clone() + fingerprint; + constraints.push(message * inv.clone() - Expr::constant(F::ONE)); + + acc_expr = acc_expr + lookup.multiplicity.clone() * inv; } + + // First row: acc_col = acc. + constraints.push(Expr::IsFirstRow * (acc_col - acc)); + // Transition: acc_expr = next row's accumulator column. + constraints.push(Expr::IsTransition * (acc_expr.clone() - next_acc_col)); + // Last row: acc_expr = next_acc. + constraints.push(Expr::IsLastRow * (acc_expr - next_acc)); + + constraints } /// Computes a fingerprint of the coefficients using Horner's method. @@ -106,54 +164,13 @@ pub struct LookupValues { } impl LookupValues { - /// Evaluates the symbolic lookups of a circuit on every row of its trace. - pub fn compute( - lookups: &[Lookup>], - trace: &RowMajorMatrix, - preprocessed: Option<&RowMajorMatrix>, - ) -> Self { - let height = trace.height(); - let num_lookups = lookups.len(); - let mut arg_offsets = Vec::with_capacity(num_lookups + 1); - arg_offsets.push(0); - for lookup in lookups { - arg_offsets.push(arg_offsets.last().unwrap() + lookup.args.len()); - } - let args_width = *arg_offsets.last().unwrap(); - let mut multiplicities = Vec::with_capacity(height * num_lookups); - let mut args = Vec::with_capacity(height * args_width); - let mut eval_row = |row: &[F], preprocessed_row: Option<&[F]>| { - for lookup in lookups { - multiplicities.push(lookup.multiplicity.interpret(row, preprocessed_row)); - for arg in &lookup.args { - args.push(arg.interpret(row, preprocessed_row)); - } - } - }; - match preprocessed { - Some(preprocessed) => trace - .row_slices() - .zip(preprocessed.row_slices()) - .for_each(|(row, preprocessed_row)| eval_row(row, Some(preprocessed_row))), - None => trace.row_slices().for_each(|row| eval_row(row, None)), - } - Self { - height, - num_lookups, - multiplicities, - arg_offsets, - args, - } - } - /// Builds flat lookup values from per-row, per-lookup concrete lookups. /// /// Every row must have the same number of lookups. Argument counts may - /// vary across rows within a lookup slot (e.g. rows holding - /// [`Lookup::empty`]); shorter argument lists are zero-padded to the - /// slot's maximum width. Padding is transparent to the protocol: the - /// message fingerprint is a Horner evaluation, so trailing zero - /// coefficients do not change its value. + /// vary across rows within a lookup slot; shorter argument lists are + /// zero-padded to the slot's maximum width. Padding is transparent to the + /// protocol: the message fingerprint is a Horner evaluation, so trailing + /// zero coefficients do not change its value. /// /// # Panics /// Panics if rows have differing numbers of lookups. @@ -291,10 +308,10 @@ impl LookupValues { /// Incremental, allocation-free constructor for [`LookupValues`]. /// -/// Rows start zeroed — multiplicity zero and zero arguments in every slot, -/// i.e. [`Lookup::empty`] — and are filled in place through the parallel row -/// writers of [`Self::par_rows_mut`]. Values are written directly into the -/// final flat storage, so no per-lookup allocation ever happens. +/// Rows start zeroed — multiplicity zero and zero arguments in every slot — +/// and are filled in place through the parallel row writers of +/// [`Self::rows_mut`]. Values are written directly into the final flat +/// storage, so no per-lookup allocation ever happens. pub struct LookupValuesBuilder { height: usize, num_lookups: usize, @@ -408,106 +425,13 @@ impl LookupRowMut<'_, F> { } } -impl BaseAir for LookupAir -where - A: BaseAir, - F: Field, -{ - fn width(&self) -> usize { - self.inner_air.width() - } - - fn preprocessed_trace(&self) -> Option> { - self.preprocessed.clone() - } -} - -impl Air for LookupAir -where - A: Air, - F: Field, - AB: TwoStagedBuilder, -{ - fn eval(&self, builder: &mut AB) { - if self.preprocessed.is_some() { - let preprocessed = builder.preprocessed().clone(); - let preprocessed_row = preprocessed.current_slice(); - self.eval_with_preprocessed_row(builder, Some(preprocessed_row)) - } else { - self.eval_with_preprocessed_row(builder, None) - } - } -} - -impl LookupAir { - fn eval_with_preprocessed_row(&self, builder: &mut AB, preprocessed_row: Option<&[AB::Var]>) - where - A: Air, - AB: TwoStagedBuilder, - { - // Call `eval` for regular stage 1 constraints. - self.inner_air.eval(builder); - - // Extract challenges and accumulators from stage 2 public values. - let stage_2_public_values = builder.stage_2_public_values(); - debug_assert_eq!(stage_2_public_values.len(), LOOKUP_PUBLIC_SIZE); - let lookup_challenge = stage_2_public_values[0].into(); - let fingerprint_challenge = stage_2_public_values[1].into(); - let acc = stage_2_public_values[2]; - let next_acc = stage_2_public_values[3]; - - // Bind relevant variables to construct the stage 2 constraints. - let stage_2 = builder.stage_2(); - let stage_2_row = stage_2.row_slice(0).unwrap(); - let stage_2_next_row = stage_2.row_slice(1).unwrap(); - let acc_col = stage_2_row[0]; - let next_acc_col = stage_2_next_row[0]; - let messages_inverses = &stage_2_row[1..]; - let lookups = &self.lookups; - debug_assert_eq!(messages_inverses.len(), lookups.len()); - - // Compute the final accumulator for the current row with the inverses - // of the messages from the stage 2 trace while asserting that these - // inverses are indeed the inverses of the messages computed on the main - // trace. - let main = builder.main(); - let row = main.current_slice(); - let mut acc_expr = acc_col.into(); - for (lookup, &message_inverse) in lookups.iter().zip(messages_inverses) { - let multiplicity: AB::ExprEF = - lookup.multiplicity.interpret(row, preprocessed_row).into(); - let args = lookup - .args - .iter() - .map(|arg| arg.interpret(row, preprocessed_row)); - let fingerprint = fingerprint(&fingerprint_challenge, args); - let message: AB::ExprEF = lookup_challenge.clone() + fingerprint; - let message_inverse = message_inverse.into(); - builder.assert_one_ext(message * message_inverse.clone()); - acc_expr += multiplicity * message_inverse; - } - - // The initial accumulator value must be set correctly. - builder.when_first_row().assert_eq_ext(acc_col, acc); - - // The accumulator computed on the main trace for the current row must - // equal the accumulator of the next row from the stage 2 trace. - builder - .when_transition() - .assert_eq_ext(acc_expr.clone(), next_acc_col); - - // The final accumulator must match the expected value. - builder.when_last_row().assert_eq_ext(acc_expr, next_acc); - } -} - #[cfg(test)] mod tests { - use p3_air::{AirBuilder, WindowAccess}; + use p3_air::{Air, AirBuilder, BaseAir, WindowAccess}; use p3_field::Field; use crate::{ - builder::symbolic::var, + p3_adapter::{LookupAir, SymbolicExpression, var}, system::{ProverKey, System, SystemWitness}, types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}, }; @@ -611,7 +535,7 @@ mod tests { }; fn system() -> ( - System, + System, ProverKey, ) { let config = GoldilocksBlake3Config::new(COMMITMENT_PARAMETERS, FRI_PARAMETERS); @@ -654,7 +578,7 @@ mod tests { traces } - fn witness(system: &System) -> SystemWitness { + fn witness(system: &System) -> SystemWitness { SystemWitness::from_stage_1(witness_traces(), system) } diff --git a/src/p3_adapter.rs b/src/p3_adapter.rs new file mode 100644 index 0000000..f29c348 --- /dev/null +++ b/src/p3_adapter.rs @@ -0,0 +1,354 @@ +//! Plonky3 adapter: build [`CircuitInputs`] from a Plonky3-style AIR. +//! +//! [`circuit_inputs_from_air`] runs [`Air::eval`] against a recording +//! builder whose expression type wraps the frontend [`Expr`], so an AIR +//! written against the `p3_air` traits (`builder.main()`, `assert_zero`, +//! `when_transition`, ...) compiles down to the same declarative +//! [`CircuitInputs`] that hand-authored circuits use. The main-trace width +//! and the optional preprocessed trace are taken from the AIR's +//! [`p3_air::BaseAir`] impl. +//! +//! Scope: +//! - Base constraints only. Extension constraints and public inputs are +//! owned by the lookup argument in this system (`System::new` derives +//! and synthesizes them), so `ExtensionBuilder` is not implemented and +//! `public_values()` returns the empty default. +//! - Lookups have no `p3_air` vocabulary; push them onto the returned +//! `CircuitInputs::lookups` afterwards. +//! - Periodic columns are not supported. + +use std::iter::{Product, Sum}; +use std::marker::PhantomData; +use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; + +use p3_air::{Air, AirBuilder, WindowAccess}; +use p3_field::{Algebra, Dup, Field, PrimeCharacteristicRing}; + +use crate::expr::{Expr, RowOffset, Source}; +use crate::lookup::Lookup; +use crate::system::CircuitInputs; + +/// Compatibility alias: the recording expression type, under the name the +/// symbolic builder used before the constraint-IR restructure. +pub type SymbolicExpression = P3Expr; + +/// Main-trace variable (current row) as an expression. +pub fn var(index: usize) -> P3Expr { + let index = u32::try_from(index).expect("column index exceeds u32"); + P3Expr(Expr::main(index)) +} + +/// Preprocessed-trace variable (current row) as an expression. +pub fn preprocessed_var(index: usize) -> P3Expr { + let index = u32::try_from(index).expect("column index exceeds u32"); + P3Expr(Expr::preprocessed(index)) +} + +/// A trace variable: a column reference in the current or next row of the +/// main or preprocessed trace. +#[derive(Debug)] +pub struct P3Var { + source: Source, + offset: RowOffset, + index: u32, + _phantom: PhantomData, +} + +// Manual impls: `derive` would put a bound on `F`, but the phantom +// parameter is only there to tie the variable to its expression type. +impl Copy for P3Var {} +impl Clone for P3Var { + fn clone(&self) -> Self { + *self + } +} + +impl P3Var { + const fn new(source: Source, offset: RowOffset, index: u32) -> Self { + Self { + source, + offset, + index, + _phantom: PhantomData, + } + } +} + +/// Expression wrapper implementing the `p3_air`/`p3_field` operator and +/// ring traits over the frontend [`Expr`] tree. +#[derive(Clone, Debug)] +pub struct P3Expr(pub Expr); + +impl From for P3Expr { + fn from(value: F) -> Self { + Self(Expr::Const(value)) + } +} + +impl From> for P3Expr { + fn from(value: P3Var) -> Self { + Self(Expr::var(value.source, value.offset, value.index)) + } +} + +impl Dup for P3Expr { + fn dup(&self) -> Self { + self.clone() + } +} + +impl Default for P3Expr { + fn default() -> Self { + Self(Expr::Const(F::ZERO)) + } +} + +impl> Add for P3Expr { + type Output = Self; + + fn add(self, rhs: T) -> Self { + Self(self.0 + rhs.into().0) + } +} + +impl> AddAssign for P3Expr { + fn add_assign(&mut self, rhs: T) { + *self = self.clone() + rhs.into(); + } +} + +impl> Sum for P3Expr { + fn sum>(iter: I) -> Self { + iter.map(Into::into) + .reduce(|x, y| x + y) + .unwrap_or(Self::ZERO) + } +} + +impl> Sub for P3Expr { + type Output = Self; + + fn sub(self, rhs: T) -> Self { + Self(self.0 - rhs.into().0) + } +} + +impl> SubAssign for P3Expr { + fn sub_assign(&mut self, rhs: T) { + *self = self.clone() - rhs.into(); + } +} + +impl Neg for P3Expr { + type Output = Self; + + fn neg(self) -> Self { + Self(-self.0) + } +} + +impl> Mul for P3Expr { + type Output = Self; + + fn mul(self, rhs: T) -> Self { + Self(self.0 * rhs.into().0) + } +} + +impl> MulAssign for P3Expr { + fn mul_assign(&mut self, rhs: T) { + *self = self.clone() * rhs.into(); + } +} + +impl> Product for P3Expr { + fn product>(iter: I) -> Self { + iter.map(Into::into) + .reduce(|x, y| x * y) + .unwrap_or(Self::ONE) + } +} + +impl PrimeCharacteristicRing for P3Expr { + type PrimeSubfield = F::PrimeSubfield; + + const ZERO: Self = Self(Expr::Const(F::ZERO)); + const ONE: Self = Self(Expr::Const(F::ONE)); + const TWO: Self = Self(Expr::Const(F::TWO)); + const NEG_ONE: Self = Self(Expr::Const(F::NEG_ONE)); + + fn from_prime_subfield(f: Self::PrimeSubfield) -> Self { + F::from_prime_subfield(f).into() + } +} + +impl Algebra for P3Expr {} + +impl Algebra> for P3Expr {} + +impl>> Add for P3Var { + type Output = P3Expr; + + fn add(self, rhs: T) -> P3Expr { + P3Expr::from(self) + rhs.into() + } +} + +impl>> Sub for P3Var { + type Output = P3Expr; + + fn sub(self, rhs: T) -> P3Expr { + P3Expr::from(self) - rhs.into() + } +} + +impl>> Mul for P3Var { + type Output = P3Expr; + + fn mul(self, rhs: T) -> P3Expr { + P3Expr::from(self) * rhs.into() + } +} + +/// Two-row window of symbolic variables over one trace. +#[derive(Clone, Debug)] +pub struct P3Window { + current: Vec>, + next: Vec>, +} + +impl P3Window { + fn new(source: Source, width: usize) -> Self { + let width = u32::try_from(width).expect("trace width exceeds u32"); + let vars = |offset| { + (0..width) + .map(|index| P3Var::new(source, offset, index)) + .collect() + }; + Self { + current: vars(RowOffset::Current), + next: vars(RowOffset::Next), + } + } +} + +impl WindowAccess> for P3Window { + fn current_slice(&self) -> &[P3Var] { + &self.current + } + + fn next_slice(&self) -> &[P3Var] { + &self.next + } +} + +/// An `AirBuilder` that records the asserted constraints as [`Expr`] trees. +pub struct P3AirBuilder { + main: P3Window, + preprocessed: P3Window, + constraints: Vec>, +} + +impl AirBuilder for P3AirBuilder { + type F = F; + type Expr = P3Expr; + type Var = P3Var; + type PreprocessedWindow = P3Window; + type MainWindow = P3Window; + // Public inputs are reserved for the lookup argument in this system, so + // the builder exposes none (the default `public_values()` is empty). + type PublicVar = P3Var; + + fn main(&self) -> Self::MainWindow { + self.main.clone() + } + + fn preprocessed(&self) -> &Self::PreprocessedWindow { + &self.preprocessed + } + + fn is_first_row(&self) -> Self::Expr { + P3Expr(Expr::IsFirstRow) + } + + fn is_last_row(&self) -> Self::Expr { + P3Expr(Expr::IsLastRow) + } + + /// # Panics + /// Panics if `size` is not `2`; only two-row windows are supported. + fn is_transition_window(&self, size: usize) -> Self::Expr { + assert_eq!(size, 2, "multi-stark only supports a window size of 2"); + P3Expr(Expr::IsTransition) + } + + fn assert_zero>(&mut self, x: I) { + self.constraints.push(x.into().0); + } +} + +/// A Plonky3-style AIR paired with its lookups. +/// +/// This is the unit `System::new` accepts for AIR-authored circuits: it +/// converts into [`CircuitInputs`] by evaluating the AIR (see +/// [`circuit_inputs_from_air`]) and attaching the lookups. +pub struct LookupAir { + pub inner_air: A, + pub lookups: Vec>>, +} + +impl LookupAir { + pub fn new(inner_air: A, lookups: Vec>>) -> Self { + Self { inner_air, lookups } + } +} + +impl>> From> for CircuitInputs { + fn from(air: LookupAir) -> Self { + let mut inputs = circuit_inputs_from_air(&air.inner_air); + inputs.lookups = air + .lookups + .into_iter() + .map(|lookup| Lookup { + multiplicity: lookup.multiplicity.0, + args: lookup.args.into_iter().map(|arg| arg.0).collect(), + }) + .collect(); + inputs + } +} + +/// Evaluates a Plonky3-style AIR into [`CircuitInputs`]: the main width and +/// preprocessed trace come from [`p3_air::BaseAir`], the constraints from +/// [`Air::eval`]. Lookups can be pushed onto the result afterwards. +/// +/// # Panics +/// Panics if the AIR declares periodic columns or public values, which +/// this system does not support. +pub fn circuit_inputs_from_air>>(air: &A) -> CircuitInputs { + assert_eq!( + air.num_periodic_columns(), + 0, + "periodic columns are not supported" + ); + assert_eq!( + air.num_public_values(), + 0, + "public inputs are reserved for the lookup argument" + ); + let main_width = air.width(); + let preprocessed = air.preprocessed_trace(); + let preprocessed_width = preprocessed.as_ref().map_or(0, |m| m.width); + let mut builder = P3AirBuilder { + main: P3Window::new(Source::Main, main_width), + preprocessed: P3Window::new(Source::Preprocessed, preprocessed_width), + constraints: vec![], + }; + air.eval(&mut builder); + CircuitInputs { + main_width, + preprocessed, + constraints: builder.constraints, + ..Default::default() + } +} diff --git a/src/prover.rs b/src/prover.rs index 2022e11..9c1a3fe 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -1,4 +1,10 @@ -//! Multi-circuit STARK prover. +//! Multi-circuit STARK prover, constraint-IR version. +//! +//! The composition polynomial is built by running the compiled circuit's +//! dense sweep in `PackedVal` on the quotient domain and folding the +//! constraint values with the constraint challenge. Stage-2 columns are +//! treated as ordinary base columns (they were committed flattened), so no +//! extension packing is needed here. //! //! The proving protocol proceeds in several stages sharing one Fiat-Shamir //! transcript. The challenger starts from a seed binding a domain tag and all @@ -124,7 +130,7 @@ //! ``` //! //! Here eval_cost(k_i) denotes the per-row cost of evaluating all k_i folded -//! constraints, which depends on the specific AIR. +//! constraints, which depends on the compiled constraint sweep. //! //! ## FRI opening //! @@ -164,27 +170,22 @@ //! with a logarithmic factor from the FFT. Doubling n_i approximately doubles //! the prover's work. -use std::ops::Deref; - -use crate::{ - builder::folder::ProverConstraintFolder, - config::{ - Com, Domain, EvaluationsOnDomain, PackedChallenge, PackedVal, PcsProof, StarkGenericConfig, - Val, - }, - lookup::{LookupValues, fingerprint}, - system::{ProverKey, System, SystemWitness}, +use crate::config::{ + Com, Domain, EvaluationsOnDomain, PackedChallenge, PackedVal, PcsProof, StarkGenericConfig, Val, }; -use bincode::{ - config::{Configuration, Fixint, LittleEndian, standard}, - error::{DecodeError, EncodeError}, - serde::{decode_from_slice, encode_to_vec}, -}; -use p3_air::{Air, BaseAir, RowWindow}; +use crate::eval::VarValues; +use crate::lookup::{LookupValues, fingerprint}; +use crate::system::{ProverKey, System, SystemWitness}; + +use bincode::config::{Configuration, Fixint, LittleEndian, standard}; +use bincode::error::{DecodeError, EncodeError}; +use bincode::serde::{decode_from_slice, encode_to_vec}; use p3_challenger::{CanObserve, FieldChallenger}; use p3_commit::{LagrangeSelectors, OpenedValuesForRound, Pcs, PolynomialSpace}; use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; -use p3_field::{BasedVectorSpace, Field, PackedValue, PrimeCharacteristicRing, TwoAdicField}; +use p3_field::{ + Algebra, BasedVectorSpace, Field, PackedValue, PrimeCharacteristicRing, TwoAdicField, +}; use p3_matrix::{Matrix, dense::RowMajorMatrix}; use p3_maybe_rayon::prelude::*; use p3_util::log2_strict_usize; @@ -246,13 +247,12 @@ impl Proof { } } -impl System +impl System where SC: StarkGenericConfig, // Two-adicity is needed to re-base the quotient's coefficient slices // onto the trace domain; every FRI-based config is two-adic anyway. Val: TwoAdicField + Ord, - A: BaseAir> + for<'a> Air, SC::Challenge>>, { /// Generates a STARK proof for the system with a single claim. /// @@ -285,7 +285,6 @@ where claims: &[&[Val]], witness: SystemWitness>, ) -> Proof { - // initialize pcs and challenger let pcs = self.config.pcs(); let mut challenger = self.config.initialise_challenger(); @@ -364,13 +363,13 @@ where challenger.observe_slice(claim); } - // generate lookup challenges + // Lookup challenges. let lookup_argument_challenge: SC::Challenge = challenger.sample_algebra_element(); challenger.observe_algebra_element(lookup_argument_challenge); let fingerprint_challenge: SC::Challenge = challenger.sample_algebra_element(); challenger.observe_algebra_element(fingerprint_challenge); - // construct the accumulator from the claims + // Initial accumulator from the claims. let mut acc = SC::Challenge::ZERO; for claim in claims { let message = lookup_argument_challenge @@ -419,7 +418,7 @@ where challenger.observe_algebra_element(*acc); } - // generate constraint challenge + // Constraint challenge. let constraint_challenge: SC::Challenge = challenger.sample_algebra_element(); // Cost: "Quotient computation and commit" — constraint evaluation on the @@ -436,7 +435,6 @@ where .enumerate() .map(|(pos, ((&ci, log_degree), next_acc))| { let circuit = &self.circuits[ci]; - let air = &circuit.air; let quotient_degree = circuit.quotient_degree(); let log_quotient_degree = log2_strict_usize(quotient_degree); let trace_domain = pcs.natural_domain_for_degree(1 << log_degree); @@ -457,25 +455,30 @@ where pcs.get_evaluations_on_domain(&stage_1_trace_data, pos, quotient_domain); let stage_2_trace_on_quotient_domain = pcs.get_evaluations_on_domain(&stage_2_trace_data, pos, quotient_domain); - // compute the quotient values which are elements of the extension field and flatten it to the base field - let public_values: [Val; 0] = []; - let stage_2_public_values = [ + + // The four lookup publics (β, γ, acc, next_acc) as flat base + // coordinates, in the layout the synthesized constraints read. + let mut lookup_publics: Vec> = Vec::new(); + for ef in [ lookup_argument_challenge, fingerprint_challenge, acc, *next_acc, - ]; - let quotient_values = quotient_values::( - air, - &public_values, - &stage_2_public_values, + ] { + lookup_publics.extend_from_slice(ef.as_basis_coefficients_slice()); + } + + // compute the quotient values which are elements of the extension field and flatten it to the base field + let quotient_values = quotient_values::( + circuit, + &lookup_publics, trace_domain, quotient_domain, &preprocessed_trace_on_quotient_domain, &stage_1_trace_on_quotient_domain, &stage_2_trace_on_quotient_domain, constraint_challenge, - circuit.constraint_count, + circuit.constraint_count(), ); let quotient_flat = RowMajorMatrix::new_col(quotient_values).flatten_to_base::>(); @@ -514,7 +517,6 @@ where challenger.observe(quotient_commit.clone()); drop(_g); - // save the commitments let commitments = Commitments { stage_1_trace: stage_1_trace_commit, stage_2_trace: stage_2_trace_commit, @@ -591,11 +593,12 @@ where } } +/// Evaluates the folded constraints on the quotient domain and divides by +/// the vanishing polynomial, producing the quotient values. #[allow(clippy::too_many_arguments)] -fn quotient_values( - air: &A, - public_values: &[Val], - stage_2_public_values: &[SC::Challenge], +fn quotient_values( + circuit: &crate::system::Circuit>, + lookup_publics: &[Val], trace_domain: Domain, quotient_domain: Domain, preprocessed_on_quotient_domain: &Option>, @@ -606,14 +609,12 @@ fn quotient_values( ) -> Vec where SC: StarkGenericConfig, - A: for<'a> Air, SC::Challenge>>, + Val: TwoAdicField + Ord, { let quotient_size = quotient_domain.size(); - let stage_1_width = stage_1_on_quotient_domain.width(); - let stage_2_width = stage_2_on_quotient_domain.width(); - let preprocessed_width = preprocessed_on_quotient_domain - .as_ref() - .map_or(0, |mat| mat.width()); + let main_width = circuit.main_width; + let stage_2_width = circuit.stage_2_width; + let preprocessed_width = circuit.preprocessed_width; let mut sels = trace_domain.selectors_on_coset(quotient_domain); let qdb = log2_strict_usize(quotient_domain.size()) - log2_strict_usize(trace_domain.size()); @@ -626,148 +627,125 @@ where sels.inv_vanishing.push(Val::::default()); } + // α powers in reverse (constraint i of k weighted by α^{k-1-i}), + // decomposed per basis coordinate for the batched base-field fold. let mut alpha_powers = alpha.powers().collect_n(constraint_count); alpha_powers.reverse(); + let decomposed_alpha_powers: Vec>> = + (0..>>::DIMENSION) + .map(|i| { + alpha_powers + .iter() + .map(|x| x.as_basis_coefficients_slice()[i]) + .collect() + }) + .collect(); - let decomposed_alpha_powers: Vec<_> = (0 - ..>>::DIMENSION) - .map(|i| { - alpha_powers - .iter() - .map(|x| x.as_basis_coefficients_slice()[i]) - .collect() - }) + // Public coordinates broadcast to packed base values. + let publics_packed: Vec> = lookup_publics + .iter() + .map(|&c| PackedVal::::from(c)) .collect(); + + let inner = |i_start: usize| { + quotient_values_inner::( + circuit, + &sels, + quotient_size, + preprocessed_on_quotient_domain, + stage_1_on_quotient_domain, + stage_2_on_quotient_domain, + main_width, + stage_2_width, + preprocessed_width, + &publics_packed, + &decomposed_alpha_powers, + next_step, + i_start, + ) + }; #[cfg(feature = "parallel")] { (0..quotient_size) .into_par_iter() .step_by(PackedVal::::WIDTH) - .flat_map_iter(|i_start| { - quotient_values_inner::( - air, - public_values, - stage_2_public_values, - &sels, - quotient_size, - preprocessed_on_quotient_domain, - stage_1_on_quotient_domain, - stage_2_on_quotient_domain, - stage_1_width, - stage_2_width, - preprocessed_width, - &alpha_powers, - &decomposed_alpha_powers, - next_step, - i_start, - ) - }) + .flat_map_iter(inner) .collect() } #[cfg(not(feature = "parallel"))] { (0..quotient_size) .step_by(PackedVal::::WIDTH) - .flat_map_iter(|i_start| { - quotient_values_inner::( - air, - public_values, - stage_2_public_values, - &sels, - quotient_size, - preprocessed_on_quotient_domain, - stage_1_on_quotient_domain, - stage_2_on_quotient_domain, - stage_1_width, - stage_2_width, - preprocessed_width, - &alpha_powers, - &decomposed_alpha_powers, - next_step, - i_start, - ) - }) + .flat_map_iter(inner) .collect() } } #[allow(clippy::too_many_arguments)] -fn quotient_values_inner( - air: &A, - stage_1_public_values: &[Val], - stage_2_public_values: &[SC::Challenge], +fn quotient_values_inner( + circuit: &crate::system::Circuit>, sels: &LagrangeSelectors>>, quotient_size: usize, preprocessed_on_quotient_domain: &Option>, stage_1_on_quotient_domain: &EvaluationsOnDomain<'_, SC>, stage_2_on_quotient_domain: &EvaluationsOnDomain<'_, SC>, - stage_1_width: usize, + main_width: usize, stage_2_width: usize, preprocessed_width: usize, - alpha_powers: &[SC::Challenge], + publics_packed: &[PackedVal], decomposed_alpha_powers: &[Vec>], next_step: usize, i_start: usize, ) -> impl Iterator where SC: StarkGenericConfig, - A: for<'a> Air, SC::Challenge>>, + Val: TwoAdicField + Ord, { let i_range = i_start..i_start + PackedVal::::WIDTH; - let is_first_row = *PackedVal::::from_slice(&sels.is_first_row[i_range.clone()]); let is_last_row = *PackedVal::::from_slice(&sels.is_last_row[i_range.clone()]); let is_transition = *PackedVal::::from_slice(&sels.is_transition[i_range.clone()]); let inv_vanishing = *PackedVal::::from_slice(&sels.inv_vanishing[i_range]); - let preprocessed = preprocessed_on_quotient_domain + // Packed two-row windows of each trace, as base columns. + let preprocessed_pair: Option>> = preprocessed_on_quotient_domain .as_ref() - .map(|on_quotient_domain| { - RowMajorMatrix::new( - on_quotient_domain.vertically_packed_row_pair(i_start, next_step), - preprocessed_width, - ) - }); - let stage_1 = RowMajorMatrix::new( - stage_1_on_quotient_domain.vertically_packed_row_pair::>(i_start, next_step), - stage_1_width, - ); - let extension_d = as BasedVectorSpace>>::DIMENSION; - let stage_2 = RowMajorMatrix::new( - pack::( - stage_2_on_quotient_domain.width(), - stage_2_on_quotient_domain.wrapping_row_slices(i_start, PackedVal::::WIDTH), - ) - .chain(pack::( - stage_2_on_quotient_domain.width(), - stage_2_on_quotient_domain - .wrapping_row_slices(i_start + next_step, PackedVal::::WIDTH), - )) - .collect::>(), - stage_2_width / extension_d, - ); - - let accumulator = PackedChallenge::::ZERO; - let mut folder = ProverConstraintFolder { - preprocessed: match preprocessed.as_ref() { - Some(mat) => RowWindow::from_view(&mat.as_view()), - None => RowWindow::from_two_rows(&[], &[]), - }, - stage_1: stage_1.as_view(), - stage_2: stage_2.as_view(), - stage_1_public_values, - stage_2_public_values, + .map(|m| m.vertically_packed_row_pair::>(i_start, next_step)); + let stage_1_pair = + stage_1_on_quotient_domain.vertically_packed_row_pair::>(i_start, next_step); + let stage_2_pair = + stage_2_on_quotient_domain.vertically_packed_row_pair::>(i_start, next_step); + + let (stage_1_cur, stage_1_next) = stage_1_pair.split_at(main_width); + let (stage_2_cur, stage_2_next) = stage_2_pair.split_at(stage_2_width); + let (preprocessed_cur, preprocessed_next): (&[PackedVal], &[PackedVal]) = + match &preprocessed_pair { + Some(pair) => pair.split_at(preprocessed_width), + None => (&[], &[]), + }; + + let view = VarValues { + preprocessed: [preprocessed_cur, preprocessed_next], + main: [stage_1_cur, stage_1_next], + stage2: [stage_2_cur, stage_2_next], + publics: publics_packed, is_first_row, is_last_row, is_transition, - alpha_powers, - decomposed_alpha_powers, - accumulator, - constraint_index: 0, }; - air.eval(&mut folder); - - let quotient = folder.accumulator * inv_vanishing; + let mut buf = Vec::new(); + circuit.graph.sweep(&view, &mut buf); + let constraint_values = circuit.graph.constraint_values(&buf); + + // Fold the base constraint values with α through the decomposed path, + // reassembling one packed extension accumulator. + let accumulator = PackedChallenge::::from_basis_coefficients_fn(|coeff_idx| { + PackedVal::::batched_linear_combination( + &constraint_values, + &decomposed_alpha_powers[coeff_idx], + ) + }); + let quotient = accumulator * inv_vanishing; (0..quotient_size.min(PackedVal::::WIDTH)).map(move |idx_in_packing| { SC::Challenge::from_basis_coefficients_fn(|coeff_idx| { @@ -778,15 +756,3 @@ where }) }) } - -fn pack( - n: usize, - rows: Vec]>>, -) -> impl Iterator> { - let extension_d = as BasedVectorSpace>>::DIMENSION; - (0..n).step_by(extension_d).map(move |c| { - PackedChallenge::::from_basis_coefficients_fn(|j| { - PackedVal::::from_fn(|i| rows[i][c + j]) - }) - }) -} diff --git a/src/system.rs b/src/system.rs index 3d89bd4..7f4af80 100644 --- a/src/system.rs +++ b/src/system.rs @@ -1,48 +1,151 @@ -use crate::{ - builder::symbolic::{SymbolicAirBuilder, get_max_constraint_degree, get_symbolic_constraints}, - config::{Com, PcsData, StarkGenericConfig, Val}, - lookup::{LOOKUP_PUBLIC_SIZE, LookupAir, LookupValues}, -}; -use p3_air::{Air, BaseAir}; +//! Multi-circuit STARK system. +//! +//! Circuits are compiled constraint data ([`crate::graph::ConstraintGraph`]) rather +//! than AIRs, so there is no `A` type parameter. `System::new` derives each +//! circuit's stage-2 and public-input layout from its lookups and the +//! challenge field's extension degree, appends the synthesized lookup +//! constraints, and compiles. + use p3_challenger::CanObserve; use p3_commit::{Pcs, PolynomialSpace}; -use p3_field::{ExtensionField, Field, PrimeCharacteristicRing}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing}; use p3_matrix::{Matrix, dense::RowMajorMatrix}; -/// A multi-circuit STARK system. Contains all circuits together with their -/// shared preprocessed commitment and the protocol configuration. -pub struct System { +use crate::config::{Com, PcsData, StarkGenericConfig, Val}; +use crate::lookup::LookupValues; + +use crate::eval::VarValues; +use crate::expr::{CircuitSpec, Expr, ExtExpr}; +use crate::graph::{ConstraintGraph, ExtensionParams, compile}; +use crate::lookup::{Lookup, num_publics, stage2_width, synthesize_lookups}; + +/// User-facing definition of one circuit: main-trace width, optional +/// preprocessed trace, base and extension constraints, and lookups. The +/// stage-2 width and public-input count are derived (from the lookups and +/// the challenge field's extension degree), not supplied here. +pub struct CircuitInputs { + pub main_width: usize, + pub preprocessed: Option>, + pub constraints: Vec>, + pub ext_constraints: Vec>, + pub lookups: Vec>>, +} + +impl Default for CircuitInputs { + fn default() -> Self { + Self { + main_width: 0, + preprocessed: None, + constraints: vec![], + ext_constraints: vec![], + lookups: vec![], + } + } +} + +/// A compiled circuit within the system, with the metadata the prover and +/// verifier need. The preprocessed trace is retained for witness-time +/// lookup evaluation (it is also committed at setup). +pub struct Circuit { + pub graph: ConstraintGraph, + pub main_width: usize, + pub preprocessed: Option>, + pub preprocessed_width: usize, + pub preprocessed_height: usize, + pub num_lookups: usize, + /// Stage-2 trace width in flattened base columns: `(1 + num_lookups)·D`. + pub stage_2_width: usize, + /// Public-input width in base coordinates: `4·D`. + pub num_publics: usize, +} + +impl Circuit { + /// Number of constraint roots (after canonicalization). + pub fn constraint_count(&self) -> usize { + self.graph.zeros.len() + } + + pub fn max_constraint_degree(&self) -> usize { + self.graph.max_constraint_degree as usize + } + + /// Degree of the quotient polynomial as a multiple of the trace degree. + /// Division by the vanishing polynomial reduces the composition + /// polynomial's degree by 1; the result is padded to a power of two so + /// the quotient can be split into equally-sized chunks. + pub fn quotient_degree(&self) -> usize { + (self.max_constraint_degree().max(2) - 1).next_power_of_two() + } +} + +/// A multi-circuit STARK system over compiled constraint circuits. Contains +/// all circuits together with their shared preprocessed commitment and the +/// protocol configuration. +pub struct System { pub config: SC, - pub circuits: Vec>>, + pub circuits: Vec>>, /// Commitment to all preprocessed traces (if any circuit has one). pub preprocessed_commit: Option>, - /// Maps each circuit index to its position within the preprocessed commitment. - /// `None` if that circuit has no preprocessed trace. + /// Maps each circuit index to its position within the preprocessed + /// commitment; `None` if that circuit has no preprocessed trace. pub preprocessed_indices: Vec>, } -/// Prover-side data that must be retained between system setup and proving. +/// Prover-side data retained between system setup and proving. pub struct ProverKey { /// PCS prover data for the preprocessed traces. pub preprocessed_data: Option>, } -impl System -where - SC: StarkGenericConfig, - A: BaseAir> + Air, SC::Challenge>>, -{ - #[inline] +impl System { + /// Builds the system from per-circuit inputs. + /// + /// # Panics + /// Panics if a circuit fails to compile, or if its constraint degree + /// exceeds what the PCS can serve (`max_quotient_degree`). pub fn new( config: SC, - airs: impl IntoIterator>>, + inputs: impl IntoIterator>>>, ) -> (Self, ProverKey) { let pcs = config.pcs(); + let params = extension_params::(); + let d = params.degree; + let mut circuits = vec![]; let mut preprocessed_traces = vec![]; let mut preprocessed_indices = vec![]; - for (circuit_idx, air) in airs.into_iter().enumerate() { - let (circuit, maybe_preprocessed_trace) = Circuit::from_air::(air); + for (i, input) in inputs.into_iter().enumerate() { + let input = input.into(); + let num_lookups = input.lookups.len(); + let preprocessed_width = input.preprocessed.as_ref().map_or(0, |m| m.width()); + let preprocessed_height = input.preprocessed.as_ref().map_or(0, |m| m.height()); + let s2_width = stage2_width(num_lookups, d); + let n_publics = num_publics(d); + + let mut ext_constraints = input.ext_constraints; + ext_constraints.extend(synthesize_lookups(&input.lookups, d)); + let spec = CircuitSpec { + main_width: input.main_width, + preprocessed_width, + stage2_width: s2_width, + num_publics: n_publics, + constraints: input.constraints, + ext_constraints, + lookups: input.lookups, + }; + let graph = compile(&spec, ¶ms) + .unwrap_or_else(|e| panic!("circuit {i}: constraint compilation failed: {e:?}")); + + let circuit = Circuit { + graph, + main_width: input.main_width, + preprocessed: input.preprocessed, + preprocessed_width, + preprocessed_height, + num_lookups, + stage_2_width: s2_width, + num_publics: n_publics, + }; // The prover obtains trace evaluations on the quotient domain // from the PCS, which can only serve domains up to // `max_quotient_degree` times the trace domain (the blowup @@ -50,26 +153,28 @@ where // invalid proofs, so reject the circuit upfront. assert!( circuit.quotient_degree() <= config.max_quotient_degree(), - "circuit {circuit_idx}: constraint degree {} needs quotient degree {}, but the \ - PCS only supports {}; increase log_blowup or lower the constraint degree", - circuit.max_constraint_degree, + "circuit {i}: constraint degree {} needs quotient degree {}, but the PCS only \ + supports {}; increase log_blowup or lower the constraint degree", + circuit.max_constraint_degree(), circuit.quotient_degree(), config.max_quotient_degree(), ); - circuits.push(circuit); - if let Some(preprocessed_trace) = maybe_preprocessed_trace { + + if let Some(preprocessed) = &circuit.preprocessed { preprocessed_indices.push(Some(preprocessed_traces.len())); - let domain = pcs.natural_domain_for_degree(preprocessed_trace.height()); - preprocessed_traces.push((domain, preprocessed_trace)); + let domain = pcs.natural_domain_for_degree(preprocessed.height()); + preprocessed_traces.push((domain, preprocessed.clone())); } else { preprocessed_indices.push(None); } + circuits.push(circuit); } - let (preprocessed_commit, preprocessed_data) = if !preprocessed_traces.is_empty() { + + let (preprocessed_commit, preprocessed_data) = if preprocessed_traces.is_empty() { + (None, None) + } else { let (commit, data) = pcs.commit(preprocessed_traces); (Some(commit), Some(data)) - } else { - (None, None) }; let system = Self { config, @@ -77,12 +182,9 @@ where preprocessed_commit, preprocessed_indices, }; - let key = ProverKey { preprocessed_data }; - (system, key) + (system, ProverKey { preprocessed_data }) } -} -impl System { /// Binds the system shape into the Fiat-Shamir transcript. The prover and /// the verifier must call this identically, before observing any /// commitment, so that transcripts of systems with different circuit @@ -90,43 +192,21 @@ impl System { /// via the challenger seed (see /// [`StarkGenericConfig::initialise_challenger`]). pub fn observe_shape(&self, challenger: &mut SC::Challenger) { - let mut observe_usize = |x: usize| challenger.observe(Val::::from_usize(x)); - observe_usize(self.circuits.len()); + let mut observe = |x: usize| challenger.observe(Val::::from_usize(x)); + observe(self.circuits.len()); for circuit in &self.circuits { - observe_usize(circuit.constraint_count); - observe_usize(circuit.max_constraint_degree); - observe_usize(circuit.preprocessed_height); - observe_usize(circuit.preprocessed_width); - observe_usize(circuit.stage_1_width); - observe_usize(circuit.stage_2_width); + observe(circuit.constraint_count()); + observe(circuit.max_constraint_degree()); + observe(circuit.preprocessed_height); + observe(circuit.preprocessed_width); + observe(circuit.main_width); + observe(circuit.stage_2_width); } } } -/// A single circuit within the system, wrapping an AIR together with -/// precomputed metadata used by the prover and verifier. -pub struct Circuit { - pub air: LookupAir, - pub constraint_count: usize, - pub max_constraint_degree: usize, - pub preprocessed_height: usize, - pub preprocessed_width: usize, - pub stage_1_width: usize, - pub stage_2_width: usize, -} - -impl Circuit { - /// Degree of the quotient polynomial as a multiple of the trace degree. - /// Division by the vanishing polynomial reduces the composition - /// polynomial's degree by 1; the result is padded to a power of two so - /// the quotient can be split into equally-sized chunks. - pub fn quotient_degree(&self) -> usize { - (self.max_constraint_degree.max(2) - 1).next_power_of_two() - } -} - -/// Witness data for the multi-circuit system, comprising stage 1 traces and -/// the concrete lookup values derived from them. +/// Witness data for the system: stage-1 traces and the concrete lookup +/// values derived from them. #[derive(Clone)] pub struct SystemWitness { /// Stage 1 (main) execution traces, one per circuit. @@ -136,15 +216,15 @@ pub struct SystemWitness { } impl SystemWitness { - /// Builds the witness from the stage 1 traces, computing the concrete - /// lookup values for each row of each circuit. + /// Builds the witness from the stage-1 traces, computing each circuit's + /// lookup values by sweeping the compiled lookup prefix over its rows. /// /// # Panics /// Panics if the number of traces differs from the number of circuits, or /// if a circuit with a preprocessed trace receives a main trace of a /// different height (both traces are opened on the same domain, so their /// heights must match; the rows would otherwise be silently truncated). - pub fn from_stage_1(traces: Vec>, system: &System) -> Self + pub fn from_stage_1(traces: Vec>, system: &System) -> Self where SC: StarkGenericConfig, SC::Pcs: Pcs>, @@ -156,65 +236,107 @@ impl SystemWitness { ); let lookups = traces .iter() - .zip(system.circuits.iter()) + .zip(&system.circuits) .enumerate() - .map(|(circuit_idx, (trace, circuit))| { - let preprocessed = circuit.air.preprocessed.as_ref(); - if let Some(preprocessed) = preprocessed { + .map(|(i, (trace, circuit))| { + if let Some(preprocessed) = &circuit.preprocessed { assert_eq!( trace.height(), preprocessed.height(), - "circuit {circuit_idx}: main trace height must equal preprocessed trace height" + "circuit {i}: main trace height must equal preprocessed trace height" ); } - LookupValues::compute(&circuit.air.lookups, trace, preprocessed) + compute_lookup_values(circuit, trace) }) - .collect::>(); + .collect(); Self { traces, lookups } } } -impl Circuit { - pub fn from_air(air: LookupAir) -> (Self, Option>) - where - EF: ExtensionField, - A: BaseAir + Air>, - { - // as of now, we assume no public values apart from the lookup values - let io_size = 0; - let stage_1_width = air.inner_air.width(); - let stage_2_width = air.stage_2_width(); - let preprocessed_trace = air.preprocessed_trace(); - let preprocessed_height = preprocessed_trace.as_ref().map_or(0, |mat| mat.height()); - let preprocessed_width = preprocessed_trace.as_ref().map_or(0, |mat| mat.width()); - let symbolic_constraints = get_symbolic_constraints::( - &air, - preprocessed_width, - stage_1_width, - stage_2_width, - io_size, - LOOKUP_PUBLIC_SIZE, - ); - let constraint_count = symbolic_constraints.len(); - let max_constraint_degree = get_max_constraint_degree(&symbolic_constraints); - let circuit = Self { - air, - max_constraint_degree, - preprocessed_height, - preprocessed_width, - constraint_count, - stage_1_width, - stage_2_width, +/// Computes a circuit's concrete lookup values by sweeping the compiled +/// lookup prefix over each row (with wrap-around for the next-row window). +fn compute_lookup_values( + circuit: &Circuit, + trace: &RowMajorMatrix, +) -> LookupValues { + let height = trace.height(); + let slot_widths: Vec = circuit + .graph + .lookups + .iter() + .map(|lookup| lookup.args.len()) + .collect(); + // No rows, or no lookups: nothing to sweep, but preserve num_lookups. + if height == 0 || slot_widths.is_empty() { + return LookupValues::builder(height, &slot_widths).finish(); + } + + let preprocessed = circuit.preprocessed.as_ref(); + let empty: [F; 0] = []; + let mut builder = LookupValues::builder(height, &slot_widths); + let mut buf = Vec::new(); + let mut args = Vec::new(); + let mut writers = builder.rows_mut(); + for (r, writer) in writers.iter_mut().enumerate() { + let r_next = (r + 1) % height; + let main_cur = trace.row_slice(r).unwrap(); + let main_next = trace.row_slice(r_next).unwrap(); + let preprocessed_rows = + preprocessed.map(|pp| (pp.row_slice(r).unwrap(), pp.row_slice(r_next).unwrap())); + let (pp_cur, pp_next): (&[F], &[F]) = match &preprocessed_rows { + Some((cur, next)) => (cur, next), + None => (&empty, &empty), }; - (circuit, preprocessed_trace) + let view = VarValues { + preprocessed: [pp_cur, pp_next], + main: [&main_cur, &main_next], + stage2: [&empty, &empty], + publics: &empty, + is_first_row: if r == 0 { F::ONE } else { F::ZERO }, + is_last_row: if r == height - 1 { F::ONE } else { F::ZERO }, + is_transition: if r == height - 1 { F::ZERO } else { F::ONE }, + }; + circuit.graph.sweep_lookup_prefix(&view, &mut buf); + for (slot, lookup) in circuit.graph.lookups.iter().enumerate() { + let multiplicity = buf[lookup.multiplicity.index()]; + args.clear(); + args.extend(lookup.args.iter().map(|a| buf[a.index()])); + // The multiplicity expression already carries its sign, so store + // it as-is (push semantics). + writer.push(slot, multiplicity, &args); + } + } + drop(writers); + builder.finish() +} + +/// Extracts the binomial extension parameters of the challenge field +/// generically: the degree is `Challenge::DIMENSION`, and the modulus +/// constant `W` (with `X^D = W`) is recovered by evaluating `X^D` and +/// reading its base coordinate — no dependence on a concrete field type. +fn extension_params() -> ExtensionParams> { + let d = >>::DIMENSION; + let x = >>::ith_basis_element(1) + .expect("challenge field must have extension degree >= 2"); + let x_pow_d = x.powers().nth(d).expect("powers iterator is infinite"); + let coords = x_pow_d.as_basis_coefficients_slice(); + debug_assert!( + coords[1..].iter().all(|c| c.is_zero()), + "challenge field is not a binomial extension: X^D is not a base element" + ); + ExtensionParams { + degree: d, + w: coords[0], + karatsuba: d == 2, } } #[cfg(test)] mod tests { use super::*; + use crate::p3_adapter::LookupAir; use crate::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; - use p3_air::{AirBuilder, WindowAccess}; + use p3_air::{Air, AirBuilder, BaseAir, WindowAccess}; /// A trivial AIR with a preprocessed trace of 4 rows and no constraints. struct Preprocessed; diff --git a/src/test_circuits/baby_bear_config.rs b/src/test_circuits/baby_bear_config.rs index 8e6d0dc..ccf019e 100644 --- a/src/test_circuits/baby_bear_config.rs +++ b/src/test_circuits/baby_bear_config.rs @@ -6,9 +6,9 @@ //! change compiles only for the reference config, the smoke test here //! catches it. -use crate::builder::symbolic::{SymbolicExpression, var}; use crate::config::StarkGenericConfig; -use crate::lookup::{Lookup, LookupAir}; +use crate::lookup::Lookup; +use crate::p3_adapter::{LookupAir, SymbolicExpression, var}; use crate::system::{System, SystemWitness}; use crate::types::{CommitmentParameters, FriParameters}; use p3_air::{Air, AirBuilder, BaseAir, WindowAccess}; diff --git a/src/test_circuits/blake3.rs b/src/test_circuits/blake3.rs index 999afbc..2cdcd52 100644 --- a/src/test_circuits/blake3.rs +++ b/src/test_circuits/blake3.rs @@ -1,7 +1,7 @@ #[cfg(test)] mod tests { - use crate::builder::symbolic::{preprocessed_var, var}; - use crate::lookup::{Lookup, LookupAir}; + use crate::lookup::Lookup; + use crate::p3_adapter::{LookupAir, preprocessed_var, var}; use crate::system::{System, SystemWitness}; use crate::test_circuits::SymbExpr; use crate::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; @@ -1515,7 +1515,7 @@ mod tests { impl Blake3CompressionClaims { fn witness( &self, - system: &System, + system: &System, ) -> (Vec>, SystemWitness) { // Grabbing values from a claims diff --git a/src/test_circuits/byte_operations.rs b/src/test_circuits/byte_operations.rs index 213419c..58238ac 100644 --- a/src/test_circuits/byte_operations.rs +++ b/src/test_circuits/byte_operations.rs @@ -1,7 +1,7 @@ #[cfg(test)] mod tests { - use crate::builder::symbolic::{preprocessed_var, var}; - use crate::lookup::{Lookup, LookupAir}; + use crate::lookup::Lookup; + use crate::p3_adapter::{LookupAir, preprocessed_var, var}; use crate::system::{System, SystemWitness}; use crate::test_circuits::SymbExpr; use crate::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; @@ -107,7 +107,7 @@ mod tests { } impl ByteCalls { - fn witness(&self, system: &System) -> SystemWitness { + fn witness(&self, system: &System) -> SystemWitness { let mut byte_trace = RowMajorMatrix::new(vec![Val::ZERO; TRACE_WIDTH * 256 * 256], TRACE_WIDTH); for (op, x, y) in self.calls.iter() { diff --git a/src/test_circuits/mod.rs b/src/test_circuits/mod.rs index d6a2e0d..b33fbf9 100644 --- a/src/test_circuits/mod.rs +++ b/src/test_circuits/mod.rs @@ -3,7 +3,7 @@ mod blake3; mod byte_operations; mod u32_add; -use crate::builder::symbolic::SymbolicExpression; +use crate::p3_adapter::SymbolicExpression; use crate::types::Val; type SymbExpr = SymbolicExpression; diff --git a/src/test_circuits/u32_add.rs b/src/test_circuits/u32_add.rs index fb4fc02..c6d435c 100644 --- a/src/test_circuits/u32_add.rs +++ b/src/test_circuits/u32_add.rs @@ -6,8 +6,8 @@ mod tests { use crate::test_circuits::SymbExpr; use crate::{ - builder::symbolic::{preprocessed_var, var}, - lookup::{Lookup, LookupAir}, + lookup::Lookup, + p3_adapter::{LookupAir, preprocessed_var, var}, system::{ProverKey, System, SystemWitness}, types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}, }; @@ -128,7 +128,7 @@ mod tests { fn byte_system( config: GoldilocksBlake3Config, ) -> ( - System, + System, ProverKey, ) { let byte_table = LookupAir::new(U32CS::ByteTable, U32CS::ByteTable.lookups()); @@ -141,7 +141,7 @@ mod tests { } impl AddCalls { - fn witness(&self, system: &System) -> SystemWitness { + fn witness(&self, system: &System) -> SystemWitness { let byte_width = 1; let add_width = 14; let mut byte_trace = RowMajorMatrix::new(vec![Val::ZERO; byte_width * 256], byte_width); diff --git a/src/verifier.rs b/src/verifier.rs index adb2ec5..b26d1ac 100644 --- a/src/verifier.rs +++ b/src/verifier.rs @@ -1,4 +1,11 @@ -//! Multi-circuit STARK verifier. +//! Multi-circuit STARK verifier (constraint-IR version). +//! +//! The out-of-domain check recomputes each circuit's composition polynomial by +//! running the compiled circuit's dense constraint sweep in the challenge field +//! at ζ, then folding the constraint values with the constraint challenge. +//! Opened stage-2 columns are fed as base columns directly (no `from_ext_basis` +//! reassembly): the coordinate-expanded constraints are base-field polynomial +//! identities, so evaluating them at the ζ-openings is correct. //! //! # Verification steps //! @@ -42,7 +49,7 @@ //! would be far too small. //! - ρ = 2^(-log_blowup) — FRI rate parameter (inverse of the blowup factor) //! - n — number of FRI queries (`num_queries`) -//! - k — number of AIR constraints (after lookup expansion) +//! - k — number of constraints (after lookup expansion) //! - N — total number of lookup rows across all circuits //! - D — maximum degree of the quotient polynomial (trace_degree × quotient_degree) //! @@ -71,7 +78,7 @@ //! //! ## Constraint folding (α) //! -//! All k AIR constraints are folded into a single composition polynomial using +//! All k constraints are folded into a single composition polynomial using //! powers of a random challenge α. The folded polynomial Σ α^i · C_i(x) has degree //! k - 1 in α. By the Schwartz-Zippel lemma, if any individual constraint C_i is //! nonzero, the folded sum is nonzero with probability at least @@ -140,7 +147,7 @@ //! is observed before any commitment or challenge, and the verifier takes //! every active circuit's widths, constraints, and lookup structure from //! the canonical system by bitmap index. The one contract change relative -//! to dense proofs: an AIR can no longer force its own presence through +//! to dense proofs: a circuit can no longer force its own presence through //! must-hold padding rows — sparse verification guarantees the constraints //! of ACTIVE circuits plus global lookup balance, and nothing about //! inactive ones. @@ -152,19 +159,16 @@ //! actual low-degree-extension values of the witness. Do not use it when the //! witness must remain hidden from the verifier. -use crate::{ - builder::folder::VerifierConstraintFolder, - config::{PcsError, StarkGenericConfig, Val}, - ensure, ensure_eq, - lookup::fingerprint, - prover::Proof, - system::System, -}; -use p3_air::{Air, BaseAir, RowWindow}; +use crate::config::{PcsError, StarkGenericConfig, Val}; +use crate::ensure_eq; +use crate::eval::VarValues; +use crate::lookup::fingerprint; +use crate::prover::Proof; +use crate::system::System; + use p3_challenger::{CanObserve, FieldChallenger}; use p3_commit::{Pcs, PolynomialSpace}; use p3_field::{BasedVectorSpace, ExtensionField, Field, PrimeCharacteristicRing}; -use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair}; use p3_util::log2_strict_usize; /// Errors that can occur during proof verification. @@ -187,15 +191,8 @@ pub enum VerificationError { UnbalancedChannel, } -impl System -where - SC: StarkGenericConfig, - A: BaseAir> + for<'a> Air, SC::Challenge>>, -{ +impl System { /// Verifies a STARK proof against a single claim. - /// - /// Returns `Ok(())` if the proof is valid, or a [`VerificationError`] describing - /// the first check that failed. pub fn verify( &self, claim: &[Val], @@ -321,6 +318,8 @@ where // Soundness: OOD evaluation. ζ is sampled after all commitments are fixed. // A nonzero polynomial of degree ≤ D vanishes at ζ with probability ≤ D/|F_ext|. let zeta: SC::Challenge = challenger.sample_algebra_element(); + + // Reconstruct the PCS opening rounds (identical to the prover). let mut stage_1_trace_evaluations = vec![]; let mut stage_2_trace_evaluations = vec![]; let mut quotient_chunks_evaluations = vec![]; @@ -396,10 +395,10 @@ where ), ]; if let Some(preprocessed_commitment) = &self.preprocessed_commit { - coms_to_verify.extend([( + coms_to_verify.push(( preprocessed_commitment.clone(), preprocessed_trace_evaluations, - )]) + )); } // Soundness: FRI proximity test. Verifies that the committed polynomials // are close to low-degree polynomials and that the claimed evaluations are @@ -411,76 +410,72 @@ where // use the opened values to compute the composition polynomial for each circuit // and check that the evaluation of the composition polynomial equals the // product of the zerofier with the quotient + let extension_d = >>::DIMENSION; + let empty: [SC::Challenge; 0] = []; for (pos, &ci) in active_indices.iter().enumerate() { let circuit = &self.circuits[ci]; let degree = 1 << log_degrees[pos]; let quotient_degree = quotient_degrees[pos]; let next_acc = intermediate_accumulators[pos]; - let stage_1_row = &stage_1_opened_values[pos][0]; - let stage_1_next_row = &stage_1_opened_values[pos][1]; - let stage_2_row = &stage_2_opened_values[pos][0]; - let stage_2_next_row = &stage_2_opened_values[pos][1]; - let quotient_row = "ient_opened_values[pos][0]; - - // compute the composition polynomial evaluation let trace_domain = pcs.natural_domain_for_degree(degree); let sels = trace_domain.selectors_at_point(zeta); - let preprocessed = if let Some(slot) = self.preprocessed_indices[ci] { - let preprocessed_opened_values = preprocessed_opened_values.as_ref().unwrap(); - let preprocessed_row = &preprocessed_opened_values[slot][0]; - let preprocessed_next_row = &preprocessed_opened_values[slot][1]; - RowWindow::from_two_rows(preprocessed_row, preprocessed_next_row) - } else { - RowWindow::from_two_rows(&[], &[]) - }; - let stage_1 = VerticalPair::new( - RowMajorMatrixView::new_row(stage_1_row), - RowMajorMatrixView::new_row(stage_1_next_row), - ); - let extension_d = >>::DIMENSION; - let stage_2_row = &stage_2_row - .chunks_exact(extension_d) - .map(from_ext_basis::, SC::Challenge>) - .collect::>(); - let stage_2_next_row = &stage_2_next_row - .chunks_exact(extension_d) - .map(from_ext_basis::, SC::Challenge>) - .collect::>(); - let stage_2 = VerticalPair::new( - RowMajorMatrixView::new_row(stage_2_row), - RowMajorMatrixView::new_row(stage_2_next_row), - ); - let stage_1_public_values = &[]; - let stage_2_public_values = &[ + + // The four lookup publics (β, γ, acc, next_acc) as base + // coordinates embedded into the challenge field. + let mut publics: Vec = Vec::with_capacity(4 * extension_d); + for ef in [ lookup_argument_challenge, fingerprint_challenge, acc, next_acc, - ]; - let mut folder = VerifierConstraintFolder { - preprocessed, - stage_1, - stage_2, - stage_1_public_values, - stage_2_public_values, + ] { + for &coord in ef.as_basis_coefficients_slice() { + publics.push(SC::Challenge::from(coord)); + } + } + + let (preprocessed_cur, preprocessed_next): (&[SC::Challenge], &[SC::Challenge]) = + if let Some(slot) = self.preprocessed_indices[ci] { + let values = preprocessed_opened_values.as_ref().unwrap(); + (&values[slot][0], &values[slot][1]) + } else { + (&empty, &empty) + }; + let view = VarValues { + preprocessed: [preprocessed_cur, preprocessed_next], + main: [ + &stage_1_opened_values[pos][0], + &stage_1_opened_values[pos][1], + ], + stage2: [ + &stage_2_opened_values[pos][0], + &stage_2_opened_values[pos][1], + ], + publics: &publics, is_first_row: sels.is_first_row, is_last_row: sels.is_last_row, is_transition: sels.is_transition, - alpha: constraint_challenge, - accumulator: SC::Challenge::ZERO, }; - circuit.air.eval(&mut folder); - let composition_polynomial = folder.accumulator; + let constraint_values = circuit.graph.evaluate_constraints(&view); + // Fold with α (Horner): Σ_i α^{k-1-i} · value_i, matching the + // prover's reversed α-power weighting. + let composition = constraint_values + .iter() + .fold(SC::Challenge::ZERO, |acc, &v| { + acc * constraint_challenge + v + }); + // Recombine the quotient from its coefficient slices: // `Q(ζ) = Σᵢ ζ^{i·n}·cᵢ(ζ)`, with each slice's value read from // the wide quotient matrix opened at ζ. + let quotient_row = "ient_opened_values[pos][0]; + debug_assert_eq!(quotient_row.len(), quotient_degree * extension_d); let zeta_pow_n = zeta.exp_power_of_2(log2_strict_usize(degree)); let quotient = quotient_row .chunks_exact(extension_d) .zip(zeta_pow_n.powers()) .map(|(chunk, zeta_pow)| zeta_pow * from_ext_basis::, SC::Challenge>(chunk)) .sum::(); - debug_assert_eq!(quotient_row.len(), quotient_degree * extension_d); // Soundness: OOD check. If any constraint is violated on the trace // domain, the composition polynomial is not divisible by the vanishing @@ -490,19 +485,18 @@ where // this ensures that the opened values are consistent with actually // low-degree polynomials satisfying all constraints. ensure_eq!( - composition_polynomial * sels.inv_vanishing, + composition * sels.inv_vanishing, quotient, VerificationError::OodEvaluationMismatch ); // the accumulator must become the next accumulator for the next iteration acc = next_acc; } - Ok(()) } /// Validates the structural shape of the proof without checking any cryptographic - /// properties. Returns the quotient degrees per circuit on success. + /// properties. Returns the quotient degrees per active circuit on success. pub fn verify_shape( &self, proof: &Proof, @@ -517,11 +511,8 @@ where stage_2_opened_values, .. } = proof; - // The following are proof shape checks let num_circuits = self.circuits.len(); - // there must be at least one circuit - ensure!(num_circuits > 0, VerificationError::InvalidSystem); - // the activation bitmap covers the canonical circuit set + crate::ensure!(num_circuits > 0, VerificationError::InvalidSystem); ensure_eq!( active.len(), num_circuits, @@ -533,15 +524,12 @@ where .filter_map(|(i, &a)| a.then_some(i)) .collect(); let num_active = active_indices.len(); - // an empty activation proves nothing - ensure!(num_active > 0, VerificationError::InvalidProofShape); - // there must be one log degree per active circuit + crate::ensure!(num_active > 0, VerificationError::InvalidProofShape); ensure_eq!( log_degrees.len(), num_active, VerificationError::InvalidProofShape ); - // the preprocessed commitment is empty if and only if there are zero preprocessed circuits let num_preprocessed = self .preprocessed_indices .iter() @@ -574,13 +562,11 @@ where ); } } - // stage 1 round: one entry per active circuit ensure_eq!( stage_1_opened_values.len(), num_active, VerificationError::InvalidProofShape ); - // stage 2 round: one entry per active circuit ensure_eq!( stage_2_opened_values.len(), num_active, @@ -618,18 +604,18 @@ where } ensure_eq!( stage_1_opened_values[pos][j].len(), - circuit.stage_1_width, + circuit.main_width, VerificationError::InvalidProofShape ); - let extension_d = >>::DIMENSION; + // Stage-2 is committed as flattened base columns, so the + // opened width is already the flattened width. ensure_eq!( stage_2_opened_values[pos][j].len(), - circuit.stage_2_width * extension_d, + circuit.stage_2_width, VerificationError::InvalidProofShape ); } } - // quotient round: per active circuit let mut quotient_degrees = vec![]; for (&ci, log_degree) in active_indices.iter().zip(log_degrees) { let quotient_degree = self.circuits[ci].quotient_degree(); @@ -637,7 +623,7 @@ where // domain can still be committed and opened by the PCS. This also // guards the `1 << log_degree` shifts used during verification // against overflow on adversarial proofs. - ensure!( + crate::ensure!( usize::from(*log_degree) + log2_strict_usize(quotient_degree) <= self.config.max_log_degree(), VerificationError::InvalidProofShape @@ -651,21 +637,19 @@ where num_active, VerificationError::InvalidProofShape ); + let extension_d = >>::DIMENSION; for (pos, quotient_degree) in quotient_degrees.iter().enumerate() { - // zeta - let num_openings = 1; ensure_eq!( quotient_opened_values[pos].len(), - num_openings, + 1, VerificationError::InvalidProofShape ); ensure_eq!( quotient_opened_values[pos][0].len(), - quotient_degree * >>::DIMENSION, + quotient_degree * extension_d, VerificationError::InvalidProofShape ); } - // there must be as many intermediate accumulators as active circuits ensure_eq!( intermediate_accumulators.len(), num_active, @@ -675,6 +659,7 @@ where } } +/// Reassembles an extension element from its base coordinates. fn from_ext_basis>(coeffs: &[EF]) -> EF { coeffs .iter() @@ -687,12 +672,12 @@ fn from_ext_basis>(coeffs: &[EF]) -> EF { mod tests { use super::*; use crate::{ - lookup::LookupAir, + p3_adapter::LookupAir, prover::Proof, system::{ProverKey, SystemWitness}, types::{CommitmentParameters, ExtVal, FriParameters, GoldilocksBlake3Config, Val}, }; - use p3_air::{AirBuilder, BaseAir, WindowAccess}; + use p3_air::{Air, AirBuilder, BaseAir, WindowAccess}; use p3_matrix::dense::RowMajorMatrix; enum CS { @@ -750,7 +735,7 @@ mod tests { }; fn system() -> ( - System, + System, ProverKey, ) { let config = GoldilocksBlake3Config::new(COMMITMENT_PARAMETERS, FRI_PARAMETERS); @@ -808,7 +793,7 @@ mod tests { /// Helper: creates a small system and valid proof for negative tests. fn small_system_and_proof() -> ( - System, + System, Proof, ) { let (system, key) = system(); From b76e43b388c4a0a60553d3b4d5d8ac2cbf772959 Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Tue, 28 Jul 2026 20:48:45 -0300 Subject: [PATCH 2/3] Evaluate logUp constraints directly; stop compiling them into the graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The logUp constraints are protocol machinery, fully determined by a circuit's lookups, yet they were synthesized as ExtExpr data and compiled into the constraint graph — coordinate-expanded, Karatsuba- expanded, materialized as nodes. Measured on the IxVM kernel system they were 91% of the compiled graph (626k of 682k nodes), which dominated the serialized verifying key (5.9 MB vs 1.1 MB on main) and dragged protocol boilerplate through the generic node sweep on the whole quotient domain. Now the graph carries only the user constraints and the lookup- expression prefix. lookup::logup_constraint_values evaluates the logUp formulas directly — one generic implementation over Algebra, run in PackedVal on the prover's quotient domain and in the challenge field at zeta by the verifier, so both sides share the exact same code. Values fold after the user roots in a fixed protocol order (per lookup the D message-inverse coordinates, then first/transition/last row), replacing the old sorted-and-deduplicated placement of the compiled roots. Circuit now stores constraint_count (user roots + (L+3)*D) and max_constraint_degree (user graph max combined with the analytic logUp degrees, derived from the compiled prefix degrees by the same rules the graph compiler uses). synthesize_lookups remains as the executable specification: a pin test checks logup_constraint_values against the reference ExtExpr evaluation of the synthesized constraints at pseudo-random points, coordinate for coordinate, in order. --- src/lookup.rs | 249 ++++++++++++++++++++++++++++++++++++++++++++++-- src/prover.rs | 25 ++++- src/system.rs | 37 +++++-- src/verifier.rs | 22 ++++- 4 files changed, 313 insertions(+), 20 deletions(-) diff --git a/src/lookup.rs b/src/lookup.rs index d27974f..6d6c8ac 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -1,13 +1,14 @@ -//! Lookups: the [`Lookup`] type, the logUp stage-2 constraint synthesis, and +//! Lookups: the [`Lookup`] type, the direct logUp constraint evaluation, and //! the concrete lookup witness ([`LookupValues`]) plus its stage-2 trace //! construction. //! -//! Synthesis is a frontend-to-frontend function: given the lookups, it -//! produces the extension-field constraints for the running-accumulator -//! argument as ordinary [`ExtExpr`] data. [`crate::system::System::new`] -//! appends these to a circuit's `ext_constraints` before compilation, so they -//! are interned, coordinate-expanded and canonicalized like any other -//! constraint. +//! The logUp constraints are protocol machinery, not user data, so they are +//! NOT compiled into the circuit graph: [`logup_constraint_values`] evaluates +//! them directly — in `PackedVal` on the prover's quotient domain and in the +//! challenge field at ζ for the verifier — and their values are folded after +//! the user-constraint roots in the canonical protocol order. +//! [`synthesize_lookups`] remains as the executable specification the direct +//! evaluation is pinned against in tests. //! //! # Layout conventions (owned here) //! @@ -18,7 +19,9 @@ //! running accumulator, slots `1..=L` = the message inverse of each lookup, //! in lookup order. So `stage2_width = (1 + L)·d` base columns. -use p3_field::{ExtensionField, Field, PrimeCharacteristicRing, batch_multiplicative_inverse}; +use p3_field::{ + Algebra, ExtensionField, Field, PrimeCharacteristicRing, batch_multiplicative_inverse, +}; use p3_matrix::dense::RowMajorMatrix; use p3_maybe_rayon::prelude::*; @@ -82,9 +85,143 @@ pub fn stage2_width(num_lookups: usize, d: usize) -> usize { (1 + num_lookups) * d } +/// Number of base-field constraint values the logUp argument contributes: +/// one extension constraint per lookup (message · inverse = 1) plus the +/// three accumulator constraints (first/transition/last row), each expanded +/// into `d` coordinates. +pub fn logup_constraint_count(num_lookups: usize, d: usize) -> usize { + (num_lookups + 3) * d +} + +/// Coordinate product in the binomial extension `X^d = w`, schoolbook: +/// `out[k] = Σ_{i+j=k} a_i·b_j + w · Σ_{i+j=k+d} a_i·b_j`. +fn coord_mul + Copy>(a: &[A], b: &[A], w: F) -> Vec { + let d = a.len(); + let mut out = vec![A::ZERO; d]; + for i in 0..d { + for j in 0..d { + let prod = a[i] * b[j]; + if i + j < d { + out[i + j] += prod; + } else { + out[i + j - d] += prod * w; + } + } + } + out +} + +/// Directly evaluates the logUp constraint values at one evaluation context, +/// in the canonical protocol order: for each lookup (in order) the `d` +/// coordinates of `(β + fingerprint(γ, args)) · inv − 1`, then the `d` +/// coordinates each of the first-row, transition, and last-row accumulator +/// constraints. Semantically identical to compiling +/// [`synthesize_lookups`]'s constraints and evaluating their roots (see the +/// pin test), without materializing them. +/// +/// Everything is base-field coordinate arithmetic, so the working type `A` +/// is generic: `PackedVal` on the prover's quotient domain, the challenge +/// field at ζ for the verifier — both sides share this one implementation. +/// +/// Layout contracts (see the module docs): `publics` holds the 4·d base +/// coordinates of (β, γ, acc, next_acc); `stage2` / `stage2_next` are the +/// flattened stage-2 base columns, slot 0 the accumulator, slot `1+j` the +/// message inverse of lookup `j`. `lookups` carries the evaluated lookup +/// expressions (multiplicity and args embed in coordinate 0). +#[allow(clippy::too_many_arguments)] +pub fn logup_constraint_values + Copy>( + lookups: &[Lookup], + stage2: &[A], + stage2_next: &[A], + publics: &[A], + is_first_row: A, + is_last_row: A, + is_transition: A, + w: F, + d: usize, + out: &mut Vec, +) { + let beta = &publics[..d]; + let gamma = &publics[d..2 * d]; + let acc = &publics[2 * d..3 * d]; + let next_acc = &publics[3 * d..4 * d]; + let acc_col = &stage2[..d]; + let next_acc_col = &stage2_next[..d]; + + // acc_expr = acc_col + Σ_j multiplicity_j · inv_j, built alongside the + // message constraints exactly as `synthesize_lookups` does. + let mut acc_expr: Vec = acc_col.to_vec(); + for (j, lookup) in lookups.iter().enumerate() { + let inv = &stage2[(1 + j) * d..(2 + j) * d]; + + // fingerprint = Σ_i args[i] · γ^i, via Horner over the reversed args + // (base values embed in coordinate 0). + let mut fingerprint = vec![A::ZERO; d]; + for &arg in lookup.args.iter().rev() { + fingerprint = coord_mul(&fingerprint, gamma, w); + fingerprint[0] += arg; + } + + // message = β + fingerprint; constraint: message · inv − 1 = 0. + let mut message = fingerprint; + for (m, &b) in message.iter_mut().zip(beta) { + *m += b; + } + let mut constraint = coord_mul(&message, inv, w); + constraint[0] -= A::ONE; + out.extend_from_slice(&constraint); + + for (a, &i) in acc_expr.iter_mut().zip(inv) { + *a += i * lookup.multiplicity; + } + } + + // First row: acc_col = acc. + for k in 0..d { + out.push(is_first_row * (acc_col[k] - acc[k])); + } + // Transition: acc_expr = next row's accumulator column. + for k in 0..d { + out.push(is_transition * (acc_expr[k] - next_acc_col[k])); + } + // Last row: acc_expr = next_acc. + for k in 0..d { + out.push(is_last_row * (acc_expr[k] - next_acc[k])); + } +} + +/// The maximum degree multiple over the logUp constraints, computed +/// analytically from the compiled lookup-expression degrees (mirroring the +/// graph compiler's rules: columns degree 1, publics/IsTransition 0, +/// IsFirstRow/IsLastRow 1, add = max, mul = sum). +pub fn logup_max_degree(graph: &crate::graph::ConstraintGraph) -> u32 { + let node_degree = |id: crate::graph::NodeId| graph.degrees[id.index()]; + // acc_expr = acc_col (deg 1) + Σ mult·inv (deg mult + 1). + let acc_degree = graph + .lookups + .iter() + .map(|l| node_degree(l.multiplicity) + 1) + .max() + .unwrap_or(0) + .max(1); + // Per lookup: (β + fingerprint(args)) · inv, degree = max arg degree + 1. + let message_degree = graph + .lookups + .iter() + .map(|l| l.args.iter().map(|&a| node_degree(a)).max().unwrap_or(0) + 1) + .max() + .unwrap_or(0); + // First row: 1 + 1; transition: 0 + acc; last: 1 + acc. + message_degree.max(2).max(acc_degree + 1) +} + /// Builds the logUp stage-2 constraints for `lookups` at extension degree /// `d`. Order: the per-lookup message/inverse constraints in lookup order, /// then the first-row, transition, and last-row accumulator constraints. +/// +/// Since the direct-evaluation change this is no longer compiled into the +/// circuit graph; it remains as the executable specification that +/// [`logup_constraint_values`] is pinned against in tests. pub fn synthesize_lookups(lookups: &[Lookup>], d: usize) -> Vec> { let d = u32::try_from(d).expect("extension degree exceeds u32"); let beta = ExtExpr::public(0, d); @@ -438,6 +575,102 @@ mod tests { use super::*; + /// Pin: `logup_constraint_values` (the direct evaluation the prover and + /// verifier run) computes exactly the values of the constraints + /// `synthesize_lookups` specifies, coordinate for coordinate, in order — + /// checked against the reference `ExtExpr` evaluator at pseudo-random + /// points. This is the executable mirror contract between the two. + #[test] + fn direct_logup_matches_synthesized_reference() { + use crate::eval::{VarValues, eval_expr, eval_ext_expr}; + use crate::graph::ExtensionParams; + use p3_field::PrimeCharacteristicRing; + + let params = crate::system::extension_params::(); + let (w, d) = (params.w, params.degree); + + // Lookups with assorted shapes: multi-arg with a product, single + // arg, and the degenerate empty-args case. + let lookups = vec![ + Lookup::push( + Expr::main(0), + vec![ + Expr::constant(Val::from_u32(7)), + Expr::main(1), + Expr::main(2) * Expr::main(3), + ], + ), + Lookup::pull(Expr::main(4), vec![Expr::main(5)]), + Lookup { + multiplicity: Expr::constant(Val::ONE), + args: vec![], + }, + ]; + let synthesized = synthesize_lookups(&lookups, d); + + // Deterministic pseudo-random base-field values (an equality of + // polynomials checked at random points). + let mut seed = 0x1234_5678_9abc_def0u64; + let mut next = move || { + seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + Val::from_u64(seed >> 8) + }; + let main_cur: Vec = (0..6).map(|_| next()).collect(); + let main_next: Vec = (0..6).map(|_| next()).collect(); + let s2_width = stage2_width(lookups.len(), d); + let s2_cur: Vec = (0..s2_width).map(|_| next()).collect(); + let s2_next: Vec = (0..s2_width).map(|_| next()).collect(); + let publics: Vec = (0..num_publics(d)).map(|_| next()).collect(); + let (isf, isl, ist) = (next(), next(), next()); + + let empty: [Val; 0] = []; + let view = VarValues { + preprocessed: [&empty, &empty], + main: [&main_cur, &main_next], + stage2: [&s2_cur, &s2_next], + publics: &publics, + is_first_row: isf, + is_last_row: isl, + is_transition: ist, + }; + + let ref_params = ExtensionParams { + degree: d, + w, + karatsuba: false, + }; + let mut reference: Vec = vec![]; + for constraint in &synthesized { + reference.extend(eval_ext_expr(constraint, &view, &ref_params)); + } + + let lookup_vals: Vec> = lookups + .iter() + .map(|l| Lookup { + multiplicity: eval_expr(&l.multiplicity, &view), + args: l.args.iter().map(|a| eval_expr(a, &view)).collect(), + }) + .collect(); + let mut direct: Vec = vec![]; + logup_constraint_values( + &lookup_vals, + &s2_cur, + &s2_next, + &publics, + isf, + isl, + ist, + w, + d, + &mut direct, + ); + + assert_eq!(direct.len(), logup_constraint_count(lookups.len(), d)); + assert_eq!(direct, reference); + } + enum CS { Even, Odd, diff --git a/src/prover.rs b/src/prover.rs index 9c1a3fe..b01433c 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -647,6 +647,7 @@ where .map(|&c| PackedVal::::from(c)) .collect(); + let ext_params = crate::system::extension_params::(); let inner = |i_start: usize| { quotient_values_inner::( circuit, @@ -662,6 +663,8 @@ where &decomposed_alpha_powers, next_step, i_start, + ext_params.w, + ext_params.degree, ) }; #[cfg(feature = "parallel")] @@ -696,6 +699,8 @@ fn quotient_values_inner( decomposed_alpha_powers: &[Vec>], next_step: usize, i_start: usize, + ext_w: Val, + ext_degree: usize, ) -> impl Iterator where SC: StarkGenericConfig, @@ -735,7 +740,25 @@ where }; let mut buf = Vec::new(); circuit.graph.sweep(&view, &mut buf); - let constraint_values = circuit.graph.constraint_values(&buf); + let mut constraint_values = circuit.graph.constraint_values(&buf); + // The logUp constraint values are evaluated directly (they are not + // compiled into the graph), appended after the user roots in the + // canonical protocol order. Coordinate-expanded logUp constraints are + // base-field-only, so they evaluate in `PackedVal` like everything else. + let lookup_vals = circuit.graph.lookup_values(&buf); + crate::lookup::logup_constraint_values( + &lookup_vals, + stage_2_cur, + stage_2_next, + publics_packed, + is_first_row, + is_last_row, + is_transition, + ext_w, + ext_degree, + &mut constraint_values, + ); + debug_assert_eq!(constraint_values.len(), circuit.constraint_count()); // Fold the base constraint values with α through the decomposed path, // reassembling one packed extension accumulator. diff --git a/src/system.rs b/src/system.rs index 7f4af80..cc51f42 100644 --- a/src/system.rs +++ b/src/system.rs @@ -3,8 +3,11 @@ //! Circuits are compiled constraint data ([`crate::graph::ConstraintGraph`]) rather //! than AIRs, so there is no `A` type parameter. `System::new` derives each //! circuit's stage-2 and public-input layout from its lookups and the -//! challenge field's extension degree, appends the synthesized lookup -//! constraints, and compiles. +//! challenge field's extension degree, and compiles the user constraints +//! plus the lookup-expression prefix. The logUp constraints are NOT +//! compiled: prover and verifier evaluate them directly +//! ([`crate::lookup::logup_constraint_values`]), folding their values after +//! the user roots. use p3_challenger::CanObserve; use p3_commit::{Pcs, PolynomialSpace}; @@ -17,7 +20,7 @@ use crate::lookup::LookupValues; use crate::eval::VarValues; use crate::expr::{CircuitSpec, Expr, ExtExpr}; use crate::graph::{ConstraintGraph, ExtensionParams, compile}; -use crate::lookup::{Lookup, num_publics, stage2_width, synthesize_lookups}; +use crate::lookup::{Lookup, logup_constraint_count, logup_max_degree, num_publics, stage2_width}; /// User-facing definition of one circuit: main-trace width, optional /// preprocessed trace, base and extension constraints, and lookups. The @@ -57,16 +60,22 @@ pub struct Circuit { pub stage_2_width: usize, /// Public-input width in base coordinates: `4·D`. pub num_publics: usize, + /// Total constraint count folded with α: the graph's user-constraint + /// roots plus the directly-evaluated logUp values, `(L + 3)·D`. + pub constraint_count: usize, + /// Maximum degree multiple over user constraints AND the logUp + /// constraints (the latter computed analytically, not compiled). + pub max_constraint_degree: usize, } impl Circuit { - /// Number of constraint roots (after canonicalization). + /// Number of constraint values folded with α (user roots + logUp). pub fn constraint_count(&self) -> usize { - self.graph.zeros.len() + self.constraint_count } pub fn max_constraint_degree(&self) -> usize { - self.graph.max_constraint_degree as usize + self.max_constraint_degree } /// Degree of the quotient polynomial as a multiple of the trace degree. @@ -122,20 +131,26 @@ impl System { let s2_width = stage2_width(num_lookups, d); let n_publics = num_publics(d); - let mut ext_constraints = input.ext_constraints; - ext_constraints.extend(synthesize_lookups(&input.lookups, d)); + // The logUp constraints are NOT compiled into the graph: the + // prover and verifier evaluate them directly (see + // `lookup::logup_constraint_values`), folded after the user + // roots in the canonical protocol order. The graph carries only + // the user constraints and the lookup-expression prefix. let spec = CircuitSpec { main_width: input.main_width, preprocessed_width, stage2_width: s2_width, num_publics: n_publics, constraints: input.constraints, - ext_constraints, + ext_constraints: input.ext_constraints, lookups: input.lookups, }; let graph = compile(&spec, ¶ms) .unwrap_or_else(|e| panic!("circuit {i}: constraint compilation failed: {e:?}")); + let constraint_count = graph.zeros.len() + logup_constraint_count(num_lookups, d); + let max_constraint_degree = + (graph.max_constraint_degree.max(logup_max_degree(&graph))) as usize; let circuit = Circuit { graph, main_width: input.main_width, @@ -145,6 +160,8 @@ impl System { num_lookups, stage_2_width: s2_width, num_publics: n_publics, + constraint_count, + max_constraint_degree, }; // The prover obtains trace evaluations on the quotient domain // from the PCS, which can only serve domains up to @@ -314,7 +331,7 @@ fn compute_lookup_values( /// generically: the degree is `Challenge::DIMENSION`, and the modulus /// constant `W` (with `X^D = W`) is recovered by evaluating `X^D` and /// reading its base coordinate — no dependence on a concrete field type. -fn extension_params() -> ExtensionParams> { +pub(crate) fn extension_params() -> ExtensionParams> { let d = >>::DIMENSION; let x = >>::ith_basis_element(1) .expect("challenge field must have extension degree >= 2"); diff --git a/src/verifier.rs b/src/verifier.rs index b26d1ac..aacf88e 100644 --- a/src/verifier.rs +++ b/src/verifier.rs @@ -456,7 +456,27 @@ impl System { is_last_row: sels.is_last_row, is_transition: sels.is_transition, }; - let constraint_values = circuit.graph.evaluate_constraints(&view); + // One sweep serves both the user-constraint roots and the + // evaluated lookup expressions; the logUp constraint values are + // then computed directly (never compiled) and folded after the + // user roots, in the canonical protocol order. + let mut buf = Vec::new(); + circuit.graph.sweep(&view, &mut buf); + let mut constraint_values = circuit.graph.constraint_values(&buf); + let lookup_vals = circuit.graph.lookup_values(&buf); + crate::lookup::logup_constraint_values( + &lookup_vals, + view.stage2[0], + view.stage2[1], + &publics, + view.is_first_row, + view.is_last_row, + view.is_transition, + crate::system::extension_params::().w, + extension_d, + &mut constraint_values, + ); + debug_assert_eq!(constraint_values.len(), circuit.constraint_count()); // Fold with α (Horner): Σ_i α^{k-1-i} · value_i, matching the // prover's reversed α-power weighting. let composition = constraint_values From c4355cb9140ece62fbbb41610203b52f641b9a92 Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Wed, 29 Jul 2026 14:08:13 -0300 Subject: [PATCH 3/3] logup: allocation-free Karatsuba fast path for the quotient hot loop The direct logUp evaluation regressed the prover's quotient constraint evaluation vs the compiled-graph path it replaced: coord_mul allocated a Vec per Horner step per point-packet (~hundreds of thousands of allocations per prove), lookup_values materialized a Vec per packet, and the schoolbook product paid 4 multiplications where the compiled expansion used Karatsuba's 3. logup_constraint_values now reads lookup expressions straight out of the sweep buffer (it takes the graph's Lookup list and the node-value slice; no materialization) and takes a fully allocation-free Karatsuba path for the reference degree-2 extension, with the generic Vec-based schoolbook retained for other degrees. Karatsuba is value-identical in exact arithmetic, so the protocol, transcript, and the Lean mirror are untouched (the pin test against synthesize_lookups still passes bit-for-bit). Measured on a lookup-heavy circuit (16 lookups x 4 args, width 40, 2^14 rows; quotient constraint eval only, release, medians of 3): compiled logUp (1c6ef62): ~24.5 ms direct logUp (before): ~28.8 ms (+17% regression) direct logUp (this commit): ~19.5 ms (-20% vs compiled) 28/28 tests, clippy/fmt clean. --- src/lookup.rs | 75 +++++++++++++++++++++++++++++++++++++++++++------ src/prover.rs | 7 +++-- src/verifier.rs | 4 +-- 3 files changed, 73 insertions(+), 13 deletions(-) diff --git a/src/lookup.rs b/src/lookup.rs index 6d6c8ac..06314ed 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -94,7 +94,8 @@ pub fn logup_constraint_count(num_lookups: usize, d: usize) -> usize { } /// Coordinate product in the binomial extension `X^d = w`, schoolbook: -/// `out[k] = Σ_{i+j=k} a_i·b_j + w · Σ_{i+j=k+d} a_i·b_j`. +/// `out[k] = Σ_{i+j=k} a_i·b_j + w · Σ_{i+j=k+d} a_i·b_j`. Generic-degree +/// fallback; the hot D=2 path uses [`mul2`]. fn coord_mul + Copy>(a: &[A], b: &[A], w: F) -> Vec { let d = a.len(); let mut out = vec![A::ZERO; d]; @@ -111,6 +112,16 @@ fn coord_mul + Copy>(a: &[A], b: &[A], w: F) -> Vec { out } +/// Degree-2 coordinate product `X² = w`, Karatsuba (3 multiplications, +/// matching the compiled graph's expansion): allocation-free. +#[inline] +fn mul2 + Copy>(a: (A, A), b: (A, A), w: F) -> (A, A) { + let v0 = a.0 * b.0; + let v1 = a.1 * b.1; + let cross = (a.0 + a.1) * (b.0 + b.1) - v0 - v1; + (v0 + v1 * w, cross) +} + /// Directly evaluates the logUp constraint values at one evaluation context, /// in the canonical protocol order: for each lookup (in order) the `d` /// coordinates of `(β + fingerprint(γ, args)) · inv − 1`, then the `d` @@ -130,7 +141,8 @@ fn coord_mul + Copy>(a: &[A], b: &[A], w: F) -> Vec { /// expressions (multiplicity and args embed in coordinate 0). #[allow(clippy::too_many_arguments)] pub fn logup_constraint_values + Copy>( - lookups: &[Lookup], + lookups: &[Lookup], + node_vals: &[A], stage2: &[A], stage2_next: &[A], publics: &[A], @@ -141,6 +153,44 @@ pub fn logup_constraint_values + Copy>( d: usize, out: &mut Vec, ) { + // Allocation-free Karatsuba fast path for the reference degree-2 + // extension; this runs per point-packet on the prover's quotient + // domain, so it must not touch the heap. + if d == 2 { + let nv = |id: crate::graph::NodeId| node_vals[id.index()]; + let beta = (publics[0], publics[1]); + let gamma = (publics[2], publics[3]); + let acc = (publics[4], publics[5]); + let next_acc = (publics[6], publics[7]); + let acc_col = (stage2[0], stage2[1]); + let next_acc_col = (stage2_next[0], stage2_next[1]); + + let mut acc_expr = acc_col; + for (j, lookup) in lookups.iter().enumerate() { + let inv = (stage2[2 + 2 * j], stage2[3 + 2 * j]); + // fingerprint = Σ_i args[i] · γ^i, Horner over the reversed args. + let mut f = (A::ZERO, A::ZERO); + for &arg in lookup.args.iter().rev() { + f = mul2(f, gamma, w); + f.0 += nv(arg); + } + // message = β + fingerprint; constraint: message · inv − 1 = 0. + let c = mul2((f.0 + beta.0, f.1 + beta.1), inv, w); + out.push(c.0 - A::ONE); + out.push(c.1); + let m = nv(lookup.multiplicity); + acc_expr.0 += inv.0 * m; + acc_expr.1 += inv.1 * m; + } + out.push(is_first_row * (acc_col.0 - acc.0)); + out.push(is_first_row * (acc_col.1 - acc.1)); + out.push(is_transition * (acc_expr.0 - next_acc_col.0)); + out.push(is_transition * (acc_expr.1 - next_acc_col.1)); + out.push(is_last_row * (acc_expr.0 - next_acc.0)); + out.push(is_last_row * (acc_expr.1 - next_acc.1)); + return; + } + let beta = &publics[..d]; let gamma = &publics[d..2 * d]; let acc = &publics[2 * d..3 * d]; @@ -159,7 +209,7 @@ pub fn logup_constraint_values + Copy>( let mut fingerprint = vec![A::ZERO; d]; for &arg in lookup.args.iter().rev() { fingerprint = coord_mul(&fingerprint, gamma, w); - fingerprint[0] += arg; + fingerprint[0] += node_vals[arg.index()]; } // message = β + fingerprint; constraint: message · inv − 1 = 0. @@ -171,8 +221,9 @@ pub fn logup_constraint_values + Copy>( constraint[0] -= A::ONE; out.extend_from_slice(&constraint); + let mv = node_vals[lookup.multiplicity.index()]; for (a, &i) in acc_expr.iter_mut().zip(inv) { - *a += i * lookup.multiplicity; + *a += i * mv; } } @@ -646,16 +697,24 @@ mod tests { reference.extend(eval_ext_expr(constraint, &view, &ref_params)); } - let lookup_vals: Vec> = lookups + // The direct evaluator reads lookup expressions out of a node-value + // buffer by id; build a synthetic buffer with consecutive ids. + let mut node_vals: Vec = vec![]; + let mut fresh = |v: Val| { + node_vals.push(v); + crate::graph::NodeId(u32::try_from(node_vals.len() - 1).unwrap()) + }; + let lookup_ids: Vec> = lookups .iter() .map(|l| Lookup { - multiplicity: eval_expr(&l.multiplicity, &view), - args: l.args.iter().map(|a| eval_expr(a, &view)).collect(), + multiplicity: fresh(eval_expr(&l.multiplicity, &view)), + args: l.args.iter().map(|a| fresh(eval_expr(a, &view))).collect(), }) .collect(); let mut direct: Vec = vec![]; logup_constraint_values( - &lookup_vals, + &lookup_ids, + &node_vals, &s2_cur, &s2_next, &publics, diff --git a/src/prover.rs b/src/prover.rs index b01433c..c18993a 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -744,10 +744,11 @@ where // The logUp constraint values are evaluated directly (they are not // compiled into the graph), appended after the user roots in the // canonical protocol order. Coordinate-expanded logUp constraints are - // base-field-only, so they evaluate in `PackedVal` like everything else. - let lookup_vals = circuit.graph.lookup_values(&buf); + // base-field-only, so they evaluate in `PackedVal` like everything else; + // the lookup expressions are read straight out of the sweep buffer. crate::lookup::logup_constraint_values( - &lookup_vals, + &circuit.graph.lookups, + &buf, stage_2_cur, stage_2_next, publics_packed, diff --git a/src/verifier.rs b/src/verifier.rs index aacf88e..b113c37 100644 --- a/src/verifier.rs +++ b/src/verifier.rs @@ -463,9 +463,9 @@ impl System { let mut buf = Vec::new(); circuit.graph.sweep(&view, &mut buf); let mut constraint_values = circuit.graph.constraint_values(&buf); - let lookup_vals = circuit.graph.lookup_values(&buf); crate::lookup::logup_constraint_values( - &lookup_vals, + &circuit.graph.lookups, + &buf, view.stage2[0], view.stage2[1], &publics,