From 34dc1673624ca2eceff077a6277daefd4770d04c Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:23:16 -0300 Subject: [PATCH 01/20] feat(stark): port GKR, sumcheck, and Lagrange kernel modules from feat/gkr-logup Verbatim from the branch tip (78be304f), which includes the batch GKR protocol with SVO sumcheck, all soundness-fix commits, and in-module tests (61 passing). Not yet wired into the prover/verifier. --- crypto/stark/src/gkr.rs | 3188 +++++++++++++++++++++++++++ crypto/stark/src/lagrange_kernel.rs | 311 +++ crypto/stark/src/lib.rs | 3 + crypto/stark/src/sumcheck.rs | 830 +++++++ 4 files changed, 4332 insertions(+) create mode 100644 crypto/stark/src/gkr.rs create mode 100644 crypto/stark/src/lagrange_kernel.rs create mode 100644 crypto/stark/src/sumcheck.rs diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs new file mode 100644 index 000000000..e00d64543 --- /dev/null +++ b/crypto/stark/src/gkr.rs @@ -0,0 +1,3188 @@ +use crate::sumcheck::{RoundPoly, SumcheckProof}; +use core::fmt; +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +/// Minimum parent_num_vars for enabling Split-Value Optimization (SVO). +/// Below this threshold, the standard flat eq_table approach is used. +const SVO_THRESHOLD: usize = 8; + +// ============================================================================= +// Layer enum for gate-specialized GKR +// ============================================================================= + +/// A layer in the GKR binary-tree circuit. +/// +/// Each layer has half the size of the layer below it. The leaves (input layer) +/// are at the bottom, and the root (output layer) has a single element. +/// +/// Different gate types allow specialized inner loops: +/// - `LogUpGeneric`: explicit numerators and denominators (6 muls/pair in sumcheck) +/// - `LogUpSingles`: numerators are implicitly 1 (2 muls/pair — ~50% savings at leaf) +#[derive(Debug, Clone)] +pub enum Layer { + /// LogUp with explicit numerators and denominators. + LogUpGeneric { + numerators: Vec>, + denominators: Vec>, + }, + /// LogUp where all numerators are implicitly 1. + /// Saves ~50% muls in the sumcheck inner loop (no nl/nr tables needed). + LogUpSingles { denominators: Vec> }, +} + +impl Layer { + /// Number of variables: log2 of the layer size. + pub fn n_variables(&self) -> usize { + let len = match self { + Self::LogUpGeneric { denominators, .. } => denominators.len(), + Self::LogUpSingles { denominators } => denominators.len(), + }; + debug_assert!(len.is_power_of_two()); + len.trailing_zeros() as usize + } + + /// Whether this is the root layer (single value, 0 variables). + pub fn is_output_layer(&self) -> bool { + self.n_variables() == 0 + } + + /// Returns the root (numerator, denominator) for a single-element output layer. + pub fn try_into_output_values(&self) -> Option<(FieldElement, FieldElement)> { + if !self.is_output_layer() { + return None; + } + Some(match self { + Layer::LogUpGeneric { + numerators, + denominators, + } => (numerators[0].clone(), denominators[0].clone()), + Layer::LogUpSingles { denominators } => (FieldElement::one(), denominators[0].clone()), + }) + } + + /// Computes the next (parent) layer by pairwise fraction addition. + /// Returns `None` if already at the output layer. + /// + /// Both Singles and Generic produce Generic output (since 1/a + 1/b = (a+b)/(a*b) + /// requires an explicit numerator). + pub fn next_layer(&self) -> Option { + if self.is_output_layer() { + return None; + } + Some(match self { + Self::LogUpGeneric { + numerators, + denominators, + } => next_logup_layer(Some(numerators), denominators), + Self::LogUpSingles { denominators } => next_logup_layer(None, denominators), + }) + } +} + +/// Pairwise fraction addition for LogUp layers. +/// +/// If `numerators` is `None`, all numerators are implicitly 1 (singles case): +/// 1/d[2j] + 1/d[2j+1] = (d[2j+1] + d[2j]) / (d[2j] * d[2j+1]) +/// +/// Otherwise: n[2j]/d[2j] + n[2j+1]/d[2j+1] = cross-multiply. +fn next_logup_layer( + numerators: Option<&[FieldElement]>, + denominators: &[FieldElement], +) -> Layer { + let half_n = denominators.len() / 2; + let mut next_numerators = Vec::with_capacity(half_n); + let mut next_denominators = Vec::with_capacity(half_n); + + for j in 0..half_n { + let dl = &denominators[2 * j]; + let dr = &denominators[2 * j + 1]; + let (num, den) = match numerators { + Some(nums) => { + let nl = &nums[2 * j]; + let nr = &nums[2 * j + 1]; + // nl/dl + nr/dr = (nl*dr + nr*dl) / (dl*dr) + (&(nl * dr) + &(nr * dl), dl * dr) + } + None => { + // 1/dl + 1/dr = (dr + dl) / (dl*dr) + (dl + dr, dl * dr) + } + }; + next_numerators.push(num); + next_denominators.push(den); + } + + Layer::LogUpGeneric { + numerators: next_numerators, + denominators: next_denominators, + } +} + +/// Generates all layers from the input (leaves) to the output (root). +/// +/// Returns layers[0] = input (leaves), layers[last] = output (root, 1 element). +pub fn gen_layers(input_layer: Layer) -> Vec> { + let n_variables = input_layer.n_variables(); + let mut layers = vec![input_layer]; + while let Some(next) = layers.last().unwrap().next_layer() { + layers.push(next); + } + assert_eq!(layers.len(), n_variables + 1); + layers +} + +// ============================================================================= +// Batch GKR proof types +// ============================================================================= + +/// Proof for a single layer in a batch GKR reduction. +/// +/// Contains the shared sumcheck proof and per-instance child claims (masks). +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct BatchGkrLayerProof { + /// Shared sumcheck proof for this layer (combined across all active instances). + pub sumcheck_proof: SumcheckProof, + /// Per-active-instance child claims: [n_left, n_right, d_left, d_right]. + /// Order matches the order instances became active. + pub child_claims_by_instance: Vec<[FieldElement; 4]>, +} + +/// Complete batch GKR proof for multiple fractional summation trees. +/// +/// All instances share one sumcheck per layer via random linear combination. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct BatchGkrProof { + /// Per-instance root claims as (numerator, denominator) pairs. + /// The claimed sum for instance i is root_claims[i].0 / root_claims[i].1. + pub root_claims: Vec<(FieldElement, FieldElement)>, + /// One layer proof per reduction step (from root towards leaves). + /// The number of layers equals max(n_variables) across all instances. + pub layer_proofs: Vec>, +} + +/// A rational number `numerator / denominator` in a field. +/// +/// Used in the GKR protocol to represent LogUp contributions before they are +/// reduced to a single field element via batch inversion. Fraction addition +/// uses cross-multiplication to avoid per-addition inversions. +pub struct Fraction { + pub numerator: FieldElement, + pub denominator: FieldElement, +} + +impl Fraction { + /// Create a new fraction `numerator / denominator`. + pub fn new(numerator: FieldElement, denominator: FieldElement) -> Self { + Self { + numerator, + denominator, + } + } + + /// Add two fractions via cross-multiplication: + /// a/b + c/d = (a*d + c*b) / (b*d) + /// + /// No reduction or normalization is performed. + pub fn add(&self, other: &Fraction) -> Fraction { + let numerator = + &(&self.numerator * &other.denominator) + &(&other.numerator * &self.denominator); + let denominator = &self.denominator * &other.denominator; + Fraction { + numerator, + denominator, + } + } +} + +/// One layer of the summation tree used in the GKR protocol. +/// +/// Each layer stores parallel arrays of numerators and denominators representing +/// fractions at that level. The leaf layer has N fractions; each subsequent layer +/// halves the count by pairwise fraction addition, until the root layer has 1. +pub struct SummationLayer { + pub numerators: Vec>, + pub denominators: Vec>, +} + +/// Build a summation tree from leaf fractions. +/// +/// Takes N leaf fractions (as parallel numerator/denominator vectors) and +/// returns layers from leaves (index 0, size N) to root (last index, size 1). +/// Each layer is built by pairwise fraction addition: +/// parent_n = left_n * right_d + right_n * left_d +/// parent_d = left_d * right_d +/// +/// # Panics +/// Panics if `leaf_numerators` and `leaf_denominators` have different lengths, +/// or if the length is not a power of 2. +pub fn build_summation_tree( + leaf_numerators: Vec>, + leaf_denominators: Vec>, +) -> Vec> { + let n = leaf_numerators.len(); + assert_eq!( + n, + leaf_denominators.len(), + "numerators and denominators must have the same length" + ); + assert!(n.is_power_of_two(), "number of leaves must be a power of 2"); + + // Number of layers: log2(n) + 1 (leaves + intermediate + root) + let num_layers = n.trailing_zeros() as usize + 1; + let mut layers = Vec::with_capacity(num_layers); + + // Layer 0: the leaves themselves + layers.push(SummationLayer { + numerators: leaf_numerators, + denominators: leaf_denominators, + }); + + // Build each subsequent layer by pairwise fraction addition + for layer_idx in 1..num_layers { + let prev = &layers[layer_idx - 1]; + let prev_len = prev.numerators.len(); + let new_len = prev_len / 2; + + let compute_pair = |i: usize| -> (FieldElement, FieldElement) { + let left_n = &prev.numerators[2 * i]; + let left_d = &prev.denominators[2 * i]; + let right_n = &prev.numerators[2 * i + 1]; + let right_d = &prev.denominators[2 * i + 1]; + + // Cross-multiply: (left_n * right_d + right_n * left_d) / (left_d * right_d) + let parent_n = &(left_n * right_d) + &(right_n * left_d); + let parent_d = left_d * right_d; + (parent_n, parent_d) + }; + + #[cfg(feature = "parallel")] + let (numerators, denominators): (Vec<_>, Vec<_>) = if new_len >= 256 { + (0..new_len).into_par_iter().map(compute_pair).unzip() + } else { + (0..new_len).map(compute_pair).unzip() + }; + + #[cfg(not(feature = "parallel"))] + let (numerators, denominators): (Vec<_>, Vec<_>) = (0..new_len).map(compute_pair).unzip(); + + layers.push(SummationLayer { + numerators, + denominators, + }); + } + + layers +} + +/// Proof for a single GKR layer reduction. +/// +/// Contains the sumcheck proof that reduces claims about a parent layer's MLEs +/// to claims about the children layer's MLEs, plus the four claimed evaluations +/// of the children's numerator/denominator MLEs at the reduced point. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct GkrLayerProof { + /// Sumcheck proof for the layer reduction (degree-3 round polynomials). + pub sumcheck_proof: SumcheckProof, + /// Claimed evaluations at the children layer: [n_left, n_right, d_left, d_right]. + /// These are the children MLE values at the point (r', 0) and (r', 1) where + /// r' is the sumcheck challenge point and the last coordinate selects left/right. + pub child_claims: [FieldElement; 4], +} + +/// Complete GKR proof for a fractional summation tree. +/// +/// Proves that the root of the summation tree has a specific value by +/// layer-by-layer reduction from the root to the leaves via fractional sumcheck. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct GkrProof { + /// The claimed sum at the root: numerator / denominator as a field element. + pub claimed_sum: FieldElement, + /// One layer proof per reduction step (from root towards leaves). + pub layer_proofs: Vec>, +} + +/// Compute the equality polynomial evaluations eq(point, b) for all b in {0,1}^n. +/// +/// The equality polynomial is defined as: +/// eq(x, y) = prod_{i=0}^{n-1} (x_i * y_i + (1 - x_i) * (1 - y_i)) +/// +/// For a fixed point r = (r_0, ..., r_{n-1}), this computes eq(r, b) for every +/// Boolean vector b, yielding 2^n values. Uses the standard butterfly/tensor +/// product construction: +/// - Start with [1] +/// - For each coordinate r_i, double the table: +/// existing entries are scaled by (1 - r_i), and new entries by r_i +/// +/// Returns a vector of length 2^n in little-endian bit order. +pub fn compute_eq_evals(point: &[FieldElement]) -> Vec> { + let n = point.len(); + let size = 1 << n; + let mut evals = Vec::with_capacity(size); + evals.push(FieldElement::one()); + + for r_i in point.iter() { + let one_minus_ri = &FieldElement::::one() - r_i; + let prev_len = evals.len(); + // Extend: for each existing entry e, push e * r_i, then scale e by (1 - r_i) + for j in 0..prev_len { + let new_val = &evals[j] * r_i; + evals.push(new_val); + } + // Scale existing entries by (1 - r_i) + for eval in evals[..prev_len].iter_mut() { + *eval = &*eval * &one_minus_ri; + } + } + + evals +} + +/// Evaluate the MLE (multilinear extension) of a table at a given point. +/// +/// Given evaluations `table[b]` for b in {0,1}^n, computes: +/// MLE(point) = sum_{b in {0,1}^n} table[b] * eq(point, b) +/// +/// This is equivalent to multilinear interpolation of the table at the point. +#[cfg(test)] +fn evaluate_mle( + table: &[FieldElement], + point: &[FieldElement], +) -> FieldElement { + let eq_evals = compute_eq_evals(point); + assert_eq!(table.len(), eq_evals.len()); + table + .iter() + .zip(eq_evals.iter()) + .fold(FieldElement::zero(), |acc, (t, e)| &acc + &(t * e)) +} + +/// Run the GKR prover on a fractional summation tree. +/// +/// Proves that the root of the tree has a specific numerator/denominator by +/// layer-by-layer reduction from the root to the leaves. At each layer, a +/// fractional sumcheck proves that the parent layer's claims are consistent +/// with the children layer via the gate equation: +/// parent_n(j) = child_n(2j) * child_d(2j+1) + child_n(2j+1) * child_d(2j) +/// parent_d(j) = child_d(2j) * child_d(2j+1) +/// +/// The sumcheck at each layer operates on a degree-3 function (product of eq +/// weight with two child values), so round polynomials have degree 3 with +/// 4 evaluation points each. +/// +/// # Arguments +/// - `tree`: the summation tree layers from leaves (index 0) to root (last index) +/// - `transcript`: Fiat-Shamir transcript for challenge sampling +/// +/// # Returns +/// A `GkrProof` containing the claimed sum and layer proofs, plus the final +/// evaluation point and claims at the leaf layer. +/// +/// # Panics +/// Panics if the tree is empty or has inconsistent layer sizes. +#[allow(clippy::type_complexity)] +pub fn gkr_prove( + tree: &[SummationLayer], + transcript: &mut impl IsTranscript, +) -> Result< + ( + GkrProof, + Vec>, + FieldElement, + FieldElement, + ), + GkrError, +> { + assert!(!tree.is_empty(), "tree must have at least one layer"); + + let num_layers = tree.len(); // layers 0..num_layers-1, root is at num_layers-1 + let root = &tree[num_layers - 1]; + assert_eq!( + root.numerators.len(), + 1, + "root layer must have exactly 1 element" + ); + + let root_n = &root.numerators[0]; + let root_d = &root.denominators[0]; + + // Compute claimed_sum = root_n / root_d. root_d = 0 means the LogUp challenge z + // collided with a fingerprint denominator (probability ~1/p ≈ 2^{-64}); return + // an error rather than panicking so the caller can retry with a fresh transcript. + let root_d_inv = root_d.inv().map_err(|_| GkrError::ZeroDenominator)?; + let claimed_sum = root_n * &root_d_inv; + + // Append the claimed sum to the transcript + transcript.append_field_element(&claimed_sum); + + // If the tree has only 1 layer (just leaves = root), no reductions needed + if num_layers == 1 { + return Ok(( + GkrProof { + claimed_sum, + layer_proofs: vec![], + }, + vec![], + root_n.clone(), + root_d.clone(), + )); + } + + let mut layer_proofs = Vec::with_capacity(num_layers - 1); + let mut n_claim = root_n.clone(); + let mut d_claim = root_d.clone(); + let mut current_point: Vec> = vec![]; + + // Reduce from root towards leaves: for each layer l from (num_layers-2) down to 0, + // the children layer is tree[l] and the parent is tree[l+1]. + for l in (0..num_layers - 1).rev() { + let child_n = &tree[l].numerators; + let child_d = &tree[l].denominators; + let parent_size = child_n.len() / 2; // = tree[l+1].numerators.len() + let parent_num_vars = parent_size.trailing_zeros() as usize; + + // Sample lambda to combine numerator and denominator claims + let lambda: FieldElement = transcript.sample_field_element(); + let combined_claim = &n_claim + &(&lambda * &d_claim); + + if parent_num_vars == 0 { + // Trivial case: parent has 1 element (0 variables), no sumcheck needed. + // The "sumcheck" is just a direct check: combined_claim must equal + // child_n[0]*child_d[1] + child_n[1]*child_d[0] + lambda * child_d[0]*child_d[1] + let nl = &child_n[0]; + let nr = &child_n[1]; + let dl = &child_d[0]; + let dr = &child_d[1]; + + // Provide the 4 child claims (they are just the raw values since there's + // no random point yet) + let child_claims = [nl.clone(), nr.clone(), dl.clone(), dr.clone()]; + + // Append child claims to transcript + for claim in &child_claims { + transcript.append_field_element(claim); + } + + // Sample eta to combine left/right into new claims for next layer + let eta: FieldElement = transcript.sample_field_element(); + + // New claims for the next layer down: fold left and right using eta + // The children MLE at point (eta) for a 2-element table [a, b] is: + // a*(1-eta) + b*eta + n_claim = &child_n[0] + &(&eta * &(&child_n[1] - &child_n[0])); + d_claim = &child_d[0] + &(&eta * &(&child_d[1] - &child_d[0])); + current_point = vec![eta]; + + layer_proofs.push(GkrLayerProof { + sumcheck_proof: SumcheckProof { + round_polys: vec![], + }, + child_claims, + }); + } else { + // Non-trivial case: run the fractional sumcheck over parent_num_vars variables. + // + // The function to sum over {0,1}^parent_num_vars is: + // f(b) = eq(current_point, b) * [n_left(b)*d_right(b) + n_right(b)*d_left(b) + lambda*d_left(b)*d_right(b)] + // where left/right are the even/odd children indexed by b. + // + // This function has degree 3 in each variable (product of eq * two child values), + // so the round polynomial needs 4 evaluation points (degree 3). + + // Build the four gate "bookkeeping" tables for the sumcheck: + // - nl_table: child_n[2b] (left numerators) + // - nr_table: child_n[2b+1] (right numerators) + // - dl_table: child_d[2b] (left denominators) + // - dr_table: child_d[2b+1] (right denominators) + let mut nl_table: Vec> = + (0..parent_size).map(|j| child_n[2 * j].clone()).collect(); + let mut nr_table: Vec> = (0..parent_size) + .map(|j| child_n[2 * j + 1].clone()) + .collect(); + let mut dl_table: Vec> = + (0..parent_size).map(|j| child_d[2 * j].clone()).collect(); + let mut dr_table: Vec> = (0..parent_size) + .map(|j| child_d[2 * j + 1].clone()) + .collect(); + + let mut round_polys = Vec::with_capacity(parent_num_vars); + let mut challenges = Vec::with_capacity(parent_num_vars); + let mut round_combined_claim = combined_claim.clone(); + // Eq correction factor: accumulates eq(r_k, c_k) from previous rounds. + // Instead of folding eq_table (N/2 multiplications per round), we halve + // it (N/2 additions) and track the missing fold factors in this scalar. + let mut eq_correction = FieldElement::::one(); + + // SVO (Split-Value Optimization, ePrint 2025/1117 Algorithm 5): + // For large tables, split eq(w, x) into prefix and suffix halves to + // reduce memory from 2^l to 2 * 2^{l/2}. + // + // eq(w, b) = eq_prefix(w_prefix, b_prefix) * eq_suffix(w_suffix, b_suffix) + // + // During the first suffix_len rounds, eq_suffix is halved each round + // while eq_prefix stays fixed. The inner loop restructures as: + // h_raw(t) = Σ_s eq_suffix[s] * Σ_p eq_prefix[p] * gate(t, p*suffix_half + s) + // + // After suffix rounds, eq_suffix is absorbed into eq_correction, and + // eq_prefix becomes the eq_table for the remaining prefix rounds. + let use_svo = parent_num_vars >= SVO_THRESHOLD; + + if use_svo { + // --- SVO path: split eq into prefix + suffix --- + let suffix_len = parent_num_vars / 2; + let _prefix_len = parent_num_vars - suffix_len; + let mut eq_suffix = compute_eq_evals(¤t_point[..suffix_len]); + let eq_prefix = compute_eq_evals(¤t_point[suffix_len..parent_num_vars]); + let prefix_size = eq_prefix.len(); // = 2^prefix_len, constant during suffix rounds + + // Verify initial consistency using the split eq tables + debug_assert!({ + let mut check_sum = FieldElement::::zero(); + for j in 0..parent_size { + let suffix_idx = j & (eq_suffix.len() - 1); + let prefix_idx = j >> suffix_len; + let eq_val = &eq_suffix[suffix_idx] * &eq_prefix[prefix_idx]; + let gate_val = &(&nl_table[j] * &dr_table[j]) + + &(&nr_table[j] * &dl_table[j]) + + &(&lambda * &(&dl_table[j] * &dr_table[j])); + check_sum = &check_sum + &(&eq_val * &gate_val); + } + check_sum == combined_claim + }); + + // Phase 1: Suffix rounds (first suffix_len rounds) + // Process variables current_point[0..suffix_len]. + // eq_suffix is halved each round; eq_prefix stays fixed. + #[allow(clippy::needless_range_loop)] + for round_idx in 0..suffix_len { + let r_round = ¤t_point[round_idx]; + let half = nl_table.len() / 2; + let suffix_half = eq_suffix.len() / 2; + + // Pre-halve eq_suffix (same as Dao-Thaler halving for eq_table) + for j in 0..suffix_half { + eq_suffix[j] = &eq_suffix[2 * j] + &eq_suffix[2 * j + 1]; + } + eq_suffix.truncate(suffix_half); + + let one = FieldElement::::one(); + + // Inner loop: for each suffix index, accumulate gate contributions + // weighted by eq_prefix, then weight by eq_suffix. + // + // h_raw(t) = Σ_s eq_suffix[s] * Σ_p eq_prefix[p] * gate(t, p*suffix_half + s) + let compute_suffix_contribution = |suffix_idx: usize| -> [FieldElement; 2] { + let eq_s = &eq_suffix[suffix_idx]; + let mut contrib_h0 = FieldElement::::zero(); + let mut contrib_h2 = FieldElement::::zero(); + + #[allow(clippy::needless_range_loop)] + for prefix_idx in 0..prefix_size { + let j = prefix_idx * suffix_half + suffix_idx; + let eq_p = &eq_prefix[prefix_idx]; + + let nl_l = &nl_table[2 * j]; + let nl_r = &nl_table[2 * j + 1]; + let nr_l = &nr_table[2 * j]; + let nr_r = &nr_table[2 * j + 1]; + let dl_l = &dl_table[2 * j]; + let dl_r = &dl_table[2 * j + 1]; + let dr_l = &dr_table[2 * j]; + let dr_r = &dr_table[2 * j + 1]; + + // t=0: gate = nl*dr + dl*(nr + lambda*dr) + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(&lambda * dr_l))); + contrib_h0 = &contrib_h0 + &(eq_p * &gate_0); + + // t=2: val = 2*right - left + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = + &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(&lambda * &dr_2))); + contrib_h2 = &contrib_h2 + &(eq_p * &gate_2); + } + + [eq_s * &contrib_h0, eq_s * &contrib_h2] + }; + + let zero2 = || [FieldElement::::zero(), FieldElement::::zero()]; + let add2 = |a: [FieldElement; 2], b: [FieldElement; 2]| { + [&a[0] + &b[0], &a[1] + &b[1]] + }; + + #[cfg(feature = "parallel")] + let totals: [FieldElement; 2] = if suffix_half >= 256 { + (0..suffix_half) + .into_par_iter() + .fold(zero2, |acc, s| add2(acc, compute_suffix_contribution(s))) + .reduce(zero2, add2) + } else { + (0..suffix_half) + .fold(zero2(), |acc, s| add2(acc, compute_suffix_contribution(s))) + }; + + #[cfg(not(feature = "parallel"))] + let totals: [FieldElement; 2] = (0..suffix_half) + .fold(zero2(), |acc, s| add2(acc, compute_suffix_contribution(s))); + + // Phase 2: Recover S(t) from h(t) and eq_round(r, t). + let [raw_h0, raw_h2] = totals; + let total_h0 = &eq_correction * &raw_h0; + let total_h2 = &eq_correction * &raw_h2; + + let one_minus_r = &one - r_round; + let s0 = &one_minus_r * &total_h0; + let s1 = &round_combined_claim - &s0; + + let r_inv = r_round + .inv() + .expect("r_round = 0 is probability 2^{-64} for random challenges"); + let h1 = &s1 * &r_inv; + + let three = FieldElement::::from(3u64); + let h3 = &(&(&three * &total_h2) - &(&three * &h1)) + &total_h0; + + let eq_at_2 = &(&three * r_round) - &one; + let s2 = &eq_at_2 * &total_h2; + + let eq_at_3 = &(&FieldElement::::from(5u64) * r_round) + - &FieldElement::::from(2u64); + let s3 = &eq_at_3 * &h3; + + let poly_evals = vec![s0, s1, s2, s3]; + let round_poly = RoundPoly::new(poly_evals); + + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + let challenge: FieldElement = transcript.sample_field_element(); + round_combined_claim = round_poly.evaluate(&challenge); + + let eq_update = + &(r_round * &challenge) + &(&one_minus_r * &(&one - &challenge)); + eq_correction = &eq_correction * &eq_update; + + // Fold the four gate tables + let fold_table = |table: &mut Vec>| { + #[cfg(feature = "parallel")] + if half >= 256 { + let folded: Vec> = table + .par_chunks(2) + .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) + .collect(); + *table = folded; + return; + } + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + }; + + fold_table(&mut nl_table); + fold_table(&mut nr_table); + fold_table(&mut dl_table); + fold_table(&mut dr_table); + + round_polys.push(round_poly); + challenges.push(challenge); + } + + // Transition: absorb the remaining eq_suffix scalar into eq_correction. + // After suffix_len halvings, eq_suffix has been reduced to a single entry. + debug_assert_eq!(eq_suffix.len(), 1); + eq_correction = &eq_correction * &eq_suffix[0]; + + // Phase 2: Prefix rounds (remaining prefix_len rounds) + // Use eq_prefix as the eq_table for standard Dao-Thaler halving. + let mut eq_table = eq_prefix; + + for r_round in current_point.iter().take(parent_num_vars).skip(suffix_len) { + let half = nl_table.len() / 2; + let one = FieldElement::::one(); + + // Pre-halve eq_table + for j in 0..half { + eq_table[j] = &eq_table[2 * j] + &eq_table[2 * j + 1]; + } + eq_table.truncate(half); + + let compute_pair_sums = |j: usize| -> [FieldElement; 2] { + let eq_rem = &eq_table[j]; + + let nl_l = &nl_table[2 * j]; + let nl_r = &nl_table[2 * j + 1]; + let nr_l = &nr_table[2 * j]; + let nr_r = &nr_table[2 * j + 1]; + let dl_l = &dl_table[2 * j]; + let dl_r = &dl_table[2 * j + 1]; + let dr_l = &dr_table[2 * j]; + let dr_r = &dr_table[2 * j + 1]; + + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(&lambda * dr_l))); + let h0 = eq_rem * &gate_0; + + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(&lambda * &dr_2))); + let h2 = eq_rem * &gate_2; + + [h0, h2] + }; + + let zero2 = || [FieldElement::::zero(), FieldElement::::zero()]; + let add2 = |a: [FieldElement; 2], b: [FieldElement; 2]| { + [&a[0] + &b[0], &a[1] + &b[1]] + }; + + #[cfg(feature = "parallel")] + let totals: [FieldElement; 2] = if half >= 256 { + (0..half) + .into_par_iter() + .fold(zero2, |acc, j| add2(acc, compute_pair_sums(j))) + .reduce(zero2, add2) + } else { + (0..half).fold(zero2(), |acc, j| add2(acc, compute_pair_sums(j))) + }; + + #[cfg(not(feature = "parallel"))] + let totals: [FieldElement; 2] = + (0..half).fold(zero2(), |acc, j| add2(acc, compute_pair_sums(j))); + + // Phase 2: Recover S(t) from h(t) and eq_round(r, t). + let [raw_h0, raw_h2] = totals; + let total_h0 = &eq_correction * &raw_h0; + let total_h2 = &eq_correction * &raw_h2; + + let one_minus_r = &one - r_round; + let s0 = &one_minus_r * &total_h0; + let s1 = &round_combined_claim - &s0; + + let r_inv = r_round + .inv() + .expect("r_round = 0 is probability 2^{-64} for random challenges"); + let h1 = &s1 * &r_inv; + + let three = FieldElement::::from(3u64); + let h3 = &(&(&three * &total_h2) - &(&three * &h1)) + &total_h0; + + let eq_at_2 = &(&three * r_round) - &one; + let s2 = &eq_at_2 * &total_h2; + + let eq_at_3 = &(&FieldElement::::from(5u64) * r_round) + - &FieldElement::::from(2u64); + let s3 = &eq_at_3 * &h3; + + let poly_evals = vec![s0, s1, s2, s3]; + let round_poly = RoundPoly::new(poly_evals); + + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + let challenge: FieldElement = transcript.sample_field_element(); + round_combined_claim = round_poly.evaluate(&challenge); + + let eq_update = + &(r_round * &challenge) + &(&one_minus_r * &(&one - &challenge)); + eq_correction = &eq_correction * &eq_update; + + let fold_table = |table: &mut Vec>| { + #[cfg(feature = "parallel")] + if half >= 256 { + let folded: Vec> = table + .par_chunks(2) + .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) + .collect(); + *table = folded; + return; + } + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + }; + + fold_table(&mut nl_table); + fold_table(&mut nr_table); + fold_table(&mut dl_table); + fold_table(&mut dr_table); + + round_polys.push(round_poly); + challenges.push(challenge); + } + } else { + // --- Standard path for small tables (parent_num_vars < SVO_THRESHOLD) --- + let mut eq_table = compute_eq_evals(¤t_point); + + // Verify initial consistency + debug_assert!({ + let mut check_sum = FieldElement::::zero(); + for j in 0..parent_size { + let gate_val = &(&nl_table[j] * &dr_table[j]) + + &(&nr_table[j] * &dl_table[j]) + + &(&lambda * &(&dl_table[j] * &dr_table[j])); + check_sum = &check_sum + &(&eq_table[j] * &gate_val); + } + check_sum == combined_claim + }); + + for r_round in current_point.iter().take(parent_num_vars) { + let half = nl_table.len() / 2; + let one = FieldElement::::one(); + + // Pre-halve eq_table + for j in 0..half { + eq_table[j] = &eq_table[2 * j] + &eq_table[2 * j + 1]; + } + eq_table.truncate(half); + + let compute_pair_sums = |j: usize| -> [FieldElement; 2] { + let eq_rem = &eq_table[j]; + + let nl_l = &nl_table[2 * j]; + let nl_r = &nl_table[2 * j + 1]; + let nr_l = &nr_table[2 * j]; + let nr_r = &nr_table[2 * j + 1]; + let dl_l = &dl_table[2 * j]; + let dl_r = &dl_table[2 * j + 1]; + let dr_l = &dr_table[2 * j]; + let dr_r = &dr_table[2 * j + 1]; + + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(&lambda * dr_l))); + let h0 = eq_rem * &gate_0; + + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(&lambda * &dr_2))); + let h2 = eq_rem * &gate_2; + + [h0, h2] + }; + + let zero2 = || [FieldElement::::zero(), FieldElement::::zero()]; + let add2 = |a: [FieldElement; 2], b: [FieldElement; 2]| { + [&a[0] + &b[0], &a[1] + &b[1]] + }; + + #[cfg(feature = "parallel")] + let totals: [FieldElement; 2] = if half >= 256 { + (0..half) + .into_par_iter() + .fold(zero2, |acc, j| add2(acc, compute_pair_sums(j))) + .reduce(zero2, add2) + } else { + (0..half).fold(zero2(), |acc, j| add2(acc, compute_pair_sums(j))) + }; + + #[cfg(not(feature = "parallel"))] + let totals: [FieldElement; 2] = + (0..half).fold(zero2(), |acc, j| add2(acc, compute_pair_sums(j))); + + let [raw_h0, raw_h2] = totals; + let total_h0 = &eq_correction * &raw_h0; + let total_h2 = &eq_correction * &raw_h2; + + let one_minus_r = &one - r_round; + let s0 = &one_minus_r * &total_h0; + let s1 = &round_combined_claim - &s0; + + let r_inv = r_round + .inv() + .expect("r_round = 0 is probability 2^{-64} for random challenges"); + let h1 = &s1 * &r_inv; + + let three = FieldElement::::from(3u64); + let h3 = &(&(&three * &total_h2) - &(&three * &h1)) + &total_h0; + + let eq_at_2 = &(&three * r_round) - &one; + let s2 = &eq_at_2 * &total_h2; + + let eq_at_3 = &(&FieldElement::::from(5u64) * r_round) + - &FieldElement::::from(2u64); + let s3 = &eq_at_3 * &h3; + + let poly_evals = vec![s0, s1, s2, s3]; + let round_poly = RoundPoly::new(poly_evals); + + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + let challenge: FieldElement = transcript.sample_field_element(); + round_combined_claim = round_poly.evaluate(&challenge); + + let eq_update = + &(r_round * &challenge) + &(&one_minus_r * &(&one - &challenge)); + eq_correction = &eq_correction * &eq_update; + + let fold_table = |table: &mut Vec>| { + #[cfg(feature = "parallel")] + if half >= 256 { + let folded: Vec> = table + .par_chunks(2) + .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) + .collect(); + *table = folded; + return; + } + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + }; + + fold_table(&mut nl_table); + fold_table(&mut nr_table); + fold_table(&mut dl_table); + fold_table(&mut dr_table); + + round_polys.push(round_poly); + challenges.push(challenge); + } + } + + // After all rounds, each table has a single entry: the MLE evaluated + // at the sumcheck challenge point. + let child_claims = [ + nl_table[0].clone(), + nr_table[0].clone(), + dl_table[0].clone(), + dr_table[0].clone(), + ]; + + // Append child claims to transcript + for claim in &child_claims { + transcript.append_field_element(claim); + } + + // Sample eta to fold left/right children into new claims + let eta: FieldElement = transcript.sample_field_element(); + + // New claims: fold left and right using eta + // n_claim = nl + eta*(nr - nl), d_claim = dl + eta*(dr - dl) + n_claim = &child_claims[0] + &(&eta * &(&child_claims[1] - &child_claims[0])); + d_claim = &child_claims[2] + &(&eta * &(&child_claims[3] - &child_claims[2])); + + // Update current_point: eta corresponds to x_0 (the even/odd selector), + // followed by the sumcheck challenges for the remaining parent variables. + // In little-endian convention: point[i] = value of x_i. + let mut new_point = Vec::with_capacity(challenges.len() + 1); + new_point.push(eta); + new_point.extend(challenges); + current_point = new_point; + + layer_proofs.push(GkrLayerProof { + sumcheck_proof: SumcheckProof { round_polys }, + child_claims, + }); + } + } + + Ok(( + GkrProof { + claimed_sum, + layer_proofs, + }, + current_point, + n_claim, + d_claim, + )) +} + +/// Errors that can occur during GKR proving or verification. +#[derive(Debug, Clone)] +pub enum GkrError { + /// The summation tree structure is invalid. + InvalidTree { reason: String }, + /// A sumcheck round failed verification. + SumcheckFailed { layer: usize, reason: String }, + /// The gate equation check failed at a layer. + GateCheckFailed { layer: usize }, + /// The claimed sum does not match (unused in the verifier itself, + /// but available for callers that compare the claimed sum to an external value). + ClaimedSumMismatch, + /// The root denominator is zero (LogUp challenge z collided with a fingerprint + /// denominator). Probability ~1/p ≈ 2^{-64}; prover should sample a new transcript. + ZeroDenominator, +} + +impl fmt::Display for GkrError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + GkrError::InvalidTree { reason } => { + write!(f, "invalid GKR tree: {}", reason) + } + GkrError::SumcheckFailed { layer, reason } => { + write!(f, "sumcheck failed at layer {}: {}", layer, reason) + } + GkrError::GateCheckFailed { layer } => { + write!(f, "gate check failed at layer {}", layer) + } + GkrError::ClaimedSumMismatch => { + write!(f, "claimed sum mismatch") + } + GkrError::ZeroDenominator => { + write!( + f, + "GKR root denominator is zero (LogUp challenge z collided with a fingerprint denominator; probability ~1/p ≈ 2^{{-64}})" + ) + } + } + } +} + +/// Verify a GKR proof for a fractional summation tree. +/// +/// Replays the Fiat-Shamir transcript identically to `gkr_prove` and checks: +/// - At each non-trivial layer: sumcheck round consistency (p(0)+p(1) = current_sum) +/// and the gate equation at the final evaluation point. +/// - At trivial layers (0-variable parent): no sumcheck check; soundness is +/// enforced by subsequent layers. +/// +/// # Arguments +/// - `proof`: the GKR proof produced by `gkr_prove` +/// - `transcript`: Fiat-Shamir transcript (must use the same seed as the prover) +/// +/// # Returns +/// `Ok((final_point, n_claim, d_claim))` where `final_point` is the random +/// evaluation point at the leaf layer, and `n_claim`/`d_claim` are the claimed +/// MLE evaluations of the leaf numerators/denominators at that point. +#[allow(clippy::type_complexity)] +pub fn gkr_verify( + proof: &GkrProof, + transcript: &mut impl IsTranscript, +) -> Result<(Vec>, FieldElement, FieldElement), GkrError> { + // Step 1: Append claimed_sum to transcript (mirrors prover line 239) + transcript.append_field_element(&proof.claimed_sum); + + // If there are no layer proofs, the tree had a single leaf (root = leaf). + // Return empty point and the claimed_sum as n_claim with d_claim = 1. + if proof.layer_proofs.is_empty() { + return Ok((vec![], proof.claimed_sum.clone(), FieldElement::one())); + } + + // Step 2: Initialize claims. + // The verifier sets n_claim = claimed_sum, d_claim = 1. + // This represents the same rational value as root_n/root_d. + // Trivial layers (0 sumcheck rounds) are gate-checked directly in the loop below. + let mut n_claim = proof.claimed_sum.clone(); + let mut d_claim = FieldElement::::one(); + let mut current_point: Vec> = vec![]; + + for (layer_idx, layer_proof) in proof.layer_proofs.iter().enumerate() { + // Step 4: Sample lambda (mirrors prover line 268) + let lambda: FieldElement = transcript.sample_field_element(); + + // Step 5: combined_claim = n_claim + lambda * d_claim + let combined_claim = &n_claim + &(&lambda * &d_claim); + + let round_polys = &layer_proof.sumcheck_proof.round_polys; + + if round_polys.is_empty() { + // Trivial layer (0 variables in parent): no sumcheck rounds. + // Gate check: verify that n_claim * (dl·dr) = nl·dr + nr·dl. + // + // The verifier works in normalized form: n_claim = root_n/root_d, d_claim = 1. + // The prover's gate equation root_n + λ·root_d = nl·dr + nr·dl + λ·dl·dr, + // divided by root_d (= dl·dr), becomes n_claim·(dl·dr) = nl·dr + nr·dl + // (the λ terms cancel). This binds claimed_sum to the actual tree structure. + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims; + let lhs = &n_claim * &(dl * dr); + let rhs = &(nl * dr) + &(nr * dl); + if lhs != rhs { + return Err(GkrError::GateCheckFailed { layer: layer_idx }); + } + } else { + // Non-trivial layer: verify sumcheck inline. + let num_rounds = round_polys.len(); + let mut current_sum = combined_claim; + let mut challenges = Vec::with_capacity(num_rounds); + + for (round, round_poly) in round_polys.iter().enumerate() { + // Check p(0) + p(1) == current_sum + if round_poly.sum_at_binary() != current_sum { + return Err(GkrError::SumcheckFailed { + layer: layer_idx, + reason: format!("round {} sum mismatch: p(0)+p(1) != expected sum", round), + }); + } + + // Append round poly evals to transcript (mirrors prover lines 388-389) + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + // Sample challenge (mirrors prover line 393) + let challenge: FieldElement = transcript.sample_field_element(); + + // Update current_sum to p(challenge) + current_sum = round_poly.evaluate(&challenge); + + challenges.push(challenge); + } + + // Gate check: verify that the final sumcheck evaluation equals + // eq(current_point, challenges) * gate(child_claims, lambda) + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims; + + // Compute eq(current_point, challenges) as a single field element. + // eq(a, b) = prod_i (a_i*b_i + (1-a_i)*(1-b_i)) + let eq_val = compute_eq_at_point(¤t_point, &challenges); + + // gate_combined = nl*dr + nr*dl + lambda*dl*dr + let gate_combined = &(&(nl * dr) + &(nr * dl)) + &(&lambda * &(dl * dr)); + + let expected = &eq_val * &gate_combined; + + if current_sum != expected { + return Err(GkrError::GateCheckFailed { layer: layer_idx }); + } + + // Build the new current_point from eta (below) and sumcheck challenges. + // We need to store challenges for constructing current_point after eta is sampled. + // Store them temporarily. + // (We'll construct the point after sampling eta below.) + + // Append child claims to transcript (mirrors prover lines 431-433) + for claim in &layer_proof.child_claims { + transcript.append_field_element(claim); + } + + // Sample eta (mirrors prover line 436) + let eta: FieldElement = transcript.sample_field_element(); + + // Update claims: fold left/right using eta + n_claim = nl + &(&eta * &(nr - nl)); + d_claim = dl + &(&eta * &(dr - dl)); + + // Update current_point = [eta] ++ challenges (mirrors prover lines 447-450) + let mut new_point = Vec::with_capacity(challenges.len() + 1); + new_point.push(eta); + new_point.extend(challenges); + current_point = new_point; + + continue; + } + + // Trivial layer path: append child claims and sample eta + // (mirrors prover lines 285-290) + for claim in &layer_proof.child_claims { + transcript.append_field_element(claim); + } + let eta: FieldElement = transcript.sample_field_element(); + + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims; + n_claim = nl + &(&eta * &(nr - nl)); + d_claim = dl + &(&eta * &(dr - dl)); + current_point = vec![eta.clone()]; + } + + Ok((current_point, n_claim, d_claim)) +} + +// ============================================================================= +// Batch GKR prover +// ============================================================================= + +/// Run the batch GKR prover on multiple instances (each a Vec). +/// +/// All instances share one sumcheck per layer via random linear combination +/// (sumcheck_alpha). Instances can have different numbers of layers (different +/// trace lengths). Smaller instances start participating at later layers. +/// +/// # Arguments +/// - `layers_per_instance`: for each instance, its layer tree from leaves (index 0) +/// to root (last index), as produced by `gen_layers`. +/// - `transcript`: Fiat-Shamir transcript for challenge sampling +/// +/// # Returns +/// A `BatchGkrProof`, the shared `random_point`, and per-instance `(n_claim, d_claim)`. +#[allow(clippy::type_complexity)] +pub fn gkr_prove_batch( + layers_per_instance: Vec>>, + transcript: &mut impl IsTranscript, +) -> ( + BatchGkrProof, + Vec>, + Vec<(FieldElement, FieldElement)>, +) { + let n_instances = layers_per_instance.len(); + + // Domain separation + transcript.append_bytes(b"gkr_batch"); + transcript.append_bytes(&(n_instances as u64).to_le_bytes()); + + if n_instances == 0 { + return ( + BatchGkrProof { + root_claims: vec![], + layer_proofs: vec![], + }, + vec![], + vec![], + ); + } + + // n_layers_by_instance[i] = number of tree layers - 1 = number of reduction steps + let n_layers_by_instance: Vec = layers_per_instance + .iter() + .map(|layers| layers.len() - 1) + .collect(); + let max_layers = *n_layers_by_instance.iter().max().unwrap(); + + // Extract root (numerator, denominator) for each instance + let root_claims: Vec<(FieldElement, FieldElement)> = layers_per_instance + .iter() + .map(|layers| { + let root = layers.last().unwrap(); + root.try_into_output_values() + .expect("root layer must be output") + }) + .collect(); + + // Track per-instance state + // n_claim[i], d_claim[i] start as root values and get updated each layer + let mut n_claims: Vec>> = vec![None; n_instances]; + let mut d_claims: Vec>> = vec![None; n_instances]; + + let mut current_point: Vec> = vec![]; + let mut layer_proofs = Vec::with_capacity(max_layers); + + for layer in 0..max_layers { + let n_remaining = max_layers - layer; + + // Detect output layers: instances whose tree has exactly n_remaining reduction steps + // start participating now (their root is the output layer). + // Use actual (root_n, root_d) so the gate check at the end of the sumcheck + // matches: gate(children) = nl*dr + nr*dl + lambda*dl*dr uses the true root values. + for i in 0..n_instances { + if n_layers_by_instance[i] == n_remaining { + n_claims[i] = Some(root_claims[i].0.clone()); + d_claims[i] = Some(root_claims[i].1.clone()); + } + } + + // Append active claims to transcript (matches verifier) + for i in 0..n_instances { + if let (Some(n), Some(d)) = (&n_claims[i], &d_claims[i]) { + transcript.append_field_element(n); + transcript.append_field_element(d); + } + } + + // Sample randomness + let sumcheck_alpha: FieldElement = transcript.sample_field_element(); + let lambda: FieldElement = transcript.sample_field_element(); + + // Collect active instances (those that have claims and have reduction steps remaining) + let mut active_instances: Vec = Vec::new(); + let mut combined_claims: Vec> = Vec::new(); + + for i in 0..n_instances { + if n_claims[i].is_some() && n_layers_by_instance[i] > 0 { + active_instances.push(i); + let n = n_claims[i].as_ref().unwrap(); + let d = d_claims[i].as_ref().unwrap(); + let claim = n + &(&lambda * d); + + // Apply doubling factor for instances with fewer layers + let n_unused = max_layers - n_layers_by_instance[i]; + if n_unused > 0 { + let doubling = FieldElement::::from(1u64 << n_unused); + combined_claims.push(&claim * &doubling); + } else { + combined_claims.push(claim); + } + } + } + + if active_instances.is_empty() { + break; + } + + // The child layer index for instance i: layers[tree_layer_idx] where + // tree_layer_idx = n_layers_by_instance[i] - 1 - (layer - n_unused) + // Equivalently: layers[n_remaining - 1 - n_unused] but let's compute directly. + + // Compute the max parent_num_vars across active instances + // (this determines the sumcheck dimension for this batch layer) + let parent_num_vars_by_instance: Vec = active_instances + .iter() + .map(|&i| { + let tree_layer_idx = n_remaining - 1; + // tree_layer_idx is the child layer; the parent has half the size + // child has 2^k elements → parent has 2^{k-1} → parent_num_vars = k-1 + let child_n_vars = layers_per_instance[i][tree_layer_idx].n_variables(); + debug_assert!( + child_n_vars >= 1, + "child of a non-output layer must have >= 2 elements" + ); + child_n_vars - 1 + }) + .collect(); + let max_parent_vars = *parent_num_vars_by_instance.iter().max().unwrap(); + + if max_parent_vars == 0 { + // All active instances have trivial layers (0 variables in parent). + // No sumcheck needed — provide child claims directly. + let mut child_claims_by_instance = Vec::new(); + + for (idx, &i) in active_instances.iter().enumerate() { + let tree_layer_idx = n_remaining - 1; + let child = &layers_per_instance[i][tree_layer_idx]; + + let (child_n, child_d) = match child { + Layer::LogUpGeneric { + numerators, + denominators, + } => (numerators.as_slice(), denominators.as_slice()), + Layer::LogUpSingles { denominators } => { + // For singles at size 2: numerators are [1, 1] + // We need to handle this specially. Since the child has 2 elements, + // and numerators are implicit 1s, we need explicit values. + // Actually, if the leaf is Singles with 2 elements, after next_layer + // it becomes Generic. But at the leaf level with 2 elements, + // the tree has 2 layers (leaf + root). So the "child layer" could + // be Singles with 2 elements. Let's handle it. + // For the trivial case, we just provide the claims directly. + // n_left=1, n_right=1, d_left=d[0], d_right=d[1] + child_claims_by_instance.push([ + FieldElement::one(), + FieldElement::one(), + denominators[0].clone(), + denominators[1].clone(), + ]); + continue; + } + }; + + let _ = idx; // suppress unused warning + child_claims_by_instance.push([ + child_n[0].clone(), + child_n[1].clone(), + child_d[0].clone(), + child_d[1].clone(), + ]); + } + + // Append child claims to transcript + for claims in &child_claims_by_instance { + for claim in claims { + transcript.append_field_element(claim); + } + } + + // Sample eta to fold left/right + let eta: FieldElement = transcript.sample_field_element(); + + // Update per-instance claims + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = child_claims_by_instance[idx]; + n_claims[i] = Some(nl + &(&eta * &(nr - nl))); + d_claims[i] = Some(dl + &(&eta * &(dr - dl))); + } + + current_point = vec![eta]; + + layer_proofs.push(BatchGkrLayerProof { + sumcheck_proof: SumcheckProof { + round_polys: vec![], + }, + child_claims_by_instance, + }); + } else { + // Non-trivial case: run shared sumcheck over max_parent_vars variables. + // For each active instance, build bookkeeping tables. + + // We run the sumcheck manually, combining round polynomials across instances. + let mut per_instance_tables: Vec> = active_instances + .iter() + .map(|&i| { + let tree_layer_idx = n_remaining - 1; + let child = &layers_per_instance[i][tree_layer_idx]; + + let parent_size = match child { + Layer::LogUpGeneric { denominators, .. } + | Layer::LogUpSingles { denominators } => denominators.len() / 2, + }; + + let (nl_table, nr_table, is_singles) = match child { + Layer::LogUpGeneric { numerators, .. } => { + let nl: Vec<_> = (0..parent_size) + .map(|j| numerators[2 * j].clone()) + .collect(); + let nr: Vec<_> = (0..parent_size) + .map(|j| numerators[2 * j + 1].clone()) + .collect(); + (nl, nr, false) + } + Layer::LogUpSingles { .. } => { + // Singles: numerators are all 1 + (vec![], vec![], true) + } + }; + + let denominators = match child { + Layer::LogUpGeneric { denominators, .. } + | Layer::LogUpSingles { denominators } => denominators, + }; + + let dl_table: Vec<_> = (0..parent_size) + .map(|j| denominators[2 * j].clone()) + .collect(); + let dr_table: Vec<_> = (0..parent_size) + .map(|j| denominators[2 * j + 1].clone()) + .collect(); + + let my_parent_num_vars = parent_num_vars_by_instance + [active_instances.iter().position(|&x| x == i).unwrap()]; + // current_point must have at least parent_num_vars coordinates + // (it grows by max_parent_vars+1 each layer, matching the tree doubling). + debug_assert!( + current_point.len() >= my_parent_num_vars, + "current_point.len()={} < parent_num_vars={}", + current_point.len(), + my_parent_num_vars + ); + let inst_point = instance_eval_point(¤t_point, my_parent_num_vars); + let use_svo = my_parent_num_vars >= SVO_THRESHOLD; + let svo_suffix_len = if use_svo { my_parent_num_vars / 2 } else { 0 }; + + let (eq_table, eq_prefix, eq_suffix) = if use_svo { + let suffix = compute_eq_evals(&inst_point[..svo_suffix_len]); + let prefix = + compute_eq_evals(&inst_point[svo_suffix_len..my_parent_num_vars]); + (Vec::new(), prefix, suffix) + } else { + (compute_eq_evals(&inst_point), Vec::new(), Vec::new()) + }; + + PerInstanceTables { + nl_table, + nr_table, + dl_table, + dr_table, + eq_table, + eq_correction: FieldElement::one(), + is_singles, + parent_num_vars: my_parent_num_vars, + instance_point: inst_point, + use_svo, + svo_suffix_len, + eq_prefix, + eq_suffix, + } + }) + .collect(); + + let mut round_polys = Vec::with_capacity(max_parent_vars); + let mut challenges = Vec::with_capacity(max_parent_vars); + // Combined claim across instances (via sumcheck_alpha) + let mut _round_combined_claim = { + let mut sum = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + for claim in &combined_claims { + sum = &sum + &(&alpha_pow * claim); + alpha_pow = &alpha_pow * &sumcheck_alpha; + } + sum + }; + + // Per-instance round polynomial evals, used to update combined_claims + // via O(1) interpolation instead of O(N) table scan. + let mut per_instance_evals: Vec<[FieldElement; 4]> = vec![ + [ + FieldElement::zero(), + FieldElement::zero(), + FieldElement::zero(), + FieldElement::zero() + ]; + per_instance_tables.len() + ]; + + for round_idx in 0..max_parent_vars { + // For each active instance, compute the round poly contribution + let mut batch_s0 = FieldElement::::zero(); + let mut batch_s1 = FieldElement::::zero(); + let mut batch_s2 = FieldElement::::zero(); + let mut batch_s3 = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + + for (idx, tables) in per_instance_tables.iter_mut().enumerate() { + let n_unused = max_parent_vars - tables.parent_num_vars; + + if round_idx < n_unused { + // This instance hasn't started yet — constant polynomial = claim/2 + let half_claim = + &combined_claims[idx] * &FieldElement::::from(2u64).inv().unwrap(); + // S(0) = S(1) = half_claim, S(2) = S(3) = half_claim (constant) + per_instance_evals[idx] = [ + half_claim.clone(), + half_claim.clone(), + half_claim.clone(), + half_claim.clone(), + ]; + batch_s0 = &batch_s0 + &(&alpha_pow * &half_claim); + batch_s1 = &batch_s1 + &(&alpha_pow * &half_claim); + batch_s2 = &batch_s2 + &(&alpha_pow * &half_claim); + batch_s3 = &batch_s3 + &(&alpha_pow * &half_claim); + alpha_pow = &alpha_pow * &sumcheck_alpha; + continue; + } + + let instance_round = round_idx - n_unused; + let half = tables.nl_table.len().max(tables.dl_table.len()) / 2; + + if half == 0 { + // Already reduced to constant + let half_claim = + &combined_claims[idx] * &FieldElement::::from(2u64).inv().unwrap(); + per_instance_evals[idx] = [ + half_claim.clone(), + half_claim.clone(), + half_claim.clone(), + half_claim.clone(), + ]; + batch_s0 = &batch_s0 + &(&alpha_pow * &half_claim); + batch_s1 = &batch_s1 + &(&alpha_pow * &half_claim); + batch_s2 = &batch_s2 + &(&alpha_pow * &half_claim); + batch_s3 = &batch_s3 + &(&alpha_pow * &half_claim); + alpha_pow = &alpha_pow * &sumcheck_alpha; + continue; + } + + // Eq polynomial factoring (same as single-instance prover). + let r_round = tables.instance_point[instance_round].clone(); + let one = FieldElement::::one(); + + // Helper: compute gate h(0) and h(2) for a generic pair at index j. + let gate_generic = |tables: &PerInstanceTables, + j: usize| + -> [FieldElement; 2] { + let nl_l = &tables.nl_table[2 * j]; + let nl_r = &tables.nl_table[2 * j + 1]; + let nr_l = &tables.nr_table[2 * j]; + let nr_r = &tables.nr_table[2 * j + 1]; + let dl_l = &tables.dl_table[2 * j]; + let dl_r = &tables.dl_table[2 * j + 1]; + let dr_l = &tables.dr_table[2 * j]; + let dr_r = &tables.dr_table[2 * j + 1]; + + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(&lambda * dr_l))); + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(&lambda * &dr_2))); + [gate_0, gate_2] + }; + + let gate_singles = + |tables: &PerInstanceTables, j: usize| -> [FieldElement; 2] { + let dl_l = &tables.dl_table[2 * j]; + let dl_r = &tables.dl_table[2 * j + 1]; + let dr_l = &tables.dr_table[2 * j]; + let dr_r = &tables.dr_table[2 * j + 1]; + + let gate_0 = &(dl_l + dr_l) + &(&lambda * &(dl_l * dr_l)); + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&dl_2 + &dr_2) + &(&lambda * &(&dl_2 * &dr_2)); + [gate_0, gate_2] + }; + + // Compute h(0) and h(2) using SVO or standard path. + let (raw_h0, raw_h2) = + if tables.use_svo && instance_round < tables.svo_suffix_len { + // SVO suffix round: nested eq_suffix × (eq_prefix × gate) loop + let suffix_half = tables.eq_suffix.len() / 2; + let prefix_size = tables.eq_prefix.len(); + + // Pre-halve eq_suffix + for j in 0..suffix_half { + tables.eq_suffix[j] = + &tables.eq_suffix[2 * j] + &tables.eq_suffix[2 * j + 1]; + } + tables.eq_suffix.truncate(suffix_half); + + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for suffix_idx in 0..suffix_half { + let eq_s = &tables.eq_suffix[suffix_idx]; + let mut ch0 = FieldElement::::zero(); + let mut ch2 = FieldElement::::zero(); + #[allow(clippy::needless_range_loop)] + for prefix_idx in 0..prefix_size { + let j = prefix_idx * suffix_half + suffix_idx; + let eq_p = &tables.eq_prefix[prefix_idx]; + let [g0, g2] = if tables.is_singles { + gate_singles(tables, j) + } else { + gate_generic(tables, j) + }; + ch0 = &ch0 + &(eq_p * &g0); + ch2 = &ch2 + &(eq_p * &g2); + } + h0 = &h0 + &(eq_s * &ch0); + h2 = &h2 + &(eq_s * &ch2); + } + (h0, h2) + } else { + // Standard path (non-SVO, or SVO prefix rounds after suffix is exhausted) + if tables.use_svo && instance_round == tables.svo_suffix_len { + // Transition: absorb remaining eq_suffix into eq_correction + // and switch to eq_prefix as eq_table. + debug_assert_eq!(tables.eq_suffix.len(), 1); + tables.eq_correction = &tables.eq_correction * &tables.eq_suffix[0]; + tables.eq_table = std::mem::take(&mut tables.eq_prefix); + tables.use_svo = false; + } + + // Pre-halve eq_table + for j in 0..half { + tables.eq_table[j] = + &tables.eq_table[2 * j] + &tables.eq_table[2 * j + 1]; + } + tables.eq_table.truncate(half); + + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for j in 0..half { + let eq_rem = &tables.eq_table[j]; + let [g0, g2] = if tables.is_singles { + gate_singles(tables, j) + } else { + gate_generic(tables, j) + }; + h0 = &h0 + &(eq_rem * &g0); + h2 = &h2 + &(eq_rem * &g2); + } + (h0, h2) + }; + + // Apply eq_correction + let total_h0 = &tables.eq_correction * &raw_h0; + let total_h2 = &tables.eq_correction * &raw_h2; + + // Recover S(t) from h(t) and eq_round(r, t) + let one_minus_r = &one - &r_round; + let s0 = &one_minus_r * &total_h0; + let s1 = &combined_claims[idx] - &s0; + + let r_inv = r_round.inv().expect("r_round = 0 is probability 2^{-64}"); + let h1 = &s1 * &r_inv; + + let three = FieldElement::::from(3u64); + let h3 = &(&(&three * &total_h2) - &(&three * &h1)) + &total_h0; + + let eq_at_2 = &(&three * &r_round) - &one; + let s2 = &eq_at_2 * &total_h2; + + let eq_at_3 = &(&FieldElement::::from(5u64) * &r_round) + - &FieldElement::::from(2u64); + let s3 = &eq_at_3 * &h3; + + per_instance_evals[idx] = [s0.clone(), s1.clone(), s2.clone(), s3.clone()]; + batch_s0 = &batch_s0 + &(&alpha_pow * &s0); + batch_s1 = &batch_s1 + &(&alpha_pow * &s1); + batch_s2 = &batch_s2 + &(&alpha_pow * &s2); + batch_s3 = &batch_s3 + &(&alpha_pow * &s3); + alpha_pow = &alpha_pow * &sumcheck_alpha; + } + + let round_poly = RoundPoly::new(vec![batch_s0, batch_s1, batch_s2, batch_s3]); + + // Append to transcript + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + // Sample challenge + let challenge: FieldElement = transcript.sample_field_element(); + _round_combined_claim = round_poly.evaluate(&challenge); + + // Update per-instance: fold tables, update eq_correction, and + // update combined_claims via O(1) polynomial evaluation (not O(N) table scan). + for (idx, tables) in per_instance_tables.iter_mut().enumerate() { + // Update combined_claims from saved per-instance round poly evals. + // S_i(challenge) via degree-3 Lagrange interpolation at {0,1,2,3}. + let [ref si0, ref si1, ref si2, ref si3] = per_instance_evals[idx]; + let instance_poly = + RoundPoly::new(vec![si0.clone(), si1.clone(), si2.clone(), si3.clone()]); + combined_claims[idx] = instance_poly.evaluate(&challenge); + + let n_unused = max_parent_vars - tables.parent_num_vars; + if round_idx < n_unused || tables.dl_table.len() / 2 == 0 { + continue; + } + + let instance_round = round_idx - n_unused; + let half = tables.dl_table.len() / 2; + + // Update eq_correction using instance-specific eval point + let r_round = &tables.instance_point[instance_round]; + let one = FieldElement::::one(); + let eq_update = + &(r_round * &challenge) + &(&(&one - r_round) * &(&one - &challenge)); + tables.eq_correction = &tables.eq_correction * &eq_update; + + // Fold gate tables + let fold_table = |table: &mut Vec>| { + #[cfg(feature = "parallel")] + if half >= 256 { + let folded: Vec> = table + .par_chunks(2) + .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) + .collect(); + *table = folded; + return; + } + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + }; + + if !tables.is_singles { + fold_table(&mut tables.nl_table); + fold_table(&mut tables.nr_table); + } + fold_table(&mut tables.dl_table); + fold_table(&mut tables.dr_table); + } + + round_polys.push(round_poly); + challenges.push(challenge); + } + + // After all sumcheck rounds, extract child claims from each instance's tables + let mut child_claims_by_instance = Vec::new(); + + for tables in &per_instance_tables { + if tables.is_singles { + child_claims_by_instance.push([ + FieldElement::one(), + FieldElement::one(), + tables.dl_table[0].clone(), + tables.dr_table[0].clone(), + ]); + } else { + child_claims_by_instance.push([ + tables.nl_table[0].clone(), + tables.nr_table[0].clone(), + tables.dl_table[0].clone(), + tables.dr_table[0].clone(), + ]); + } + } + + // Append child claims to transcript + for claims in &child_claims_by_instance { + for claim in claims { + transcript.append_field_element(claim); + } + } + + // Sample eta to fold left/right + let eta: FieldElement = transcript.sample_field_element(); + + // Update per-instance claims + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = child_claims_by_instance[idx]; + n_claims[i] = Some(nl + &(&eta * &(nr - nl))); + d_claims[i] = Some(dl + &(&eta * &(dr - dl))); + } + + let mut new_point = Vec::with_capacity(challenges.len() + 1); + new_point.push(eta); + new_point.extend(challenges); + current_point = new_point; + + layer_proofs.push(BatchGkrLayerProof { + sumcheck_proof: SumcheckProof { round_polys }, + child_claims_by_instance, + }); + } + } + + let final_claims: Vec<(FieldElement, FieldElement)> = (0..n_instances) + .map(|i| { + ( + n_claims[i] + .clone() + .unwrap_or_else(|| root_claims[i].0.clone()), + d_claims[i] + .clone() + .unwrap_or_else(|| root_claims[i].1.clone()), + ) + }) + .collect(); + + ( + BatchGkrProof { + root_claims, + layer_proofs, + }, + current_point, + final_claims, + ) +} + +/// Compute the per-instance evaluation point from the shared point. +/// +/// In batch GKR with mixed-size instances, smaller instances skip the first +/// rounds of each shared sumcheck. Their variables bind to the **last** +/// challenges, not the first. The eval point for instance i is: +/// `[shared_point[0] (eta)] ++ shared_point[len - (n_vars - 1)..]` +/// +/// Returns an empty vec for n_vars == 0 (trivial/output layers). +pub(crate) fn instance_eval_point( + shared_point: &[FieldElement], + n_vars: usize, +) -> Vec> { + if n_vars == 0 { + return vec![]; + } + if n_vars >= shared_point.len() { + return shared_point.to_vec(); + } + let mut point = Vec::with_capacity(n_vars); + point.push(shared_point[0].clone()); // eta (shared left/right selector) + let start = shared_point.len() - (n_vars - 1); + point.extend_from_slice(&shared_point[start..]); + point +} + +/// Per-instance bookkeeping tables for the batch sumcheck inner loop. +struct PerInstanceTables { + nl_table: Vec>, + nr_table: Vec>, + dl_table: Vec>, + dr_table: Vec>, + eq_table: Vec>, + eq_correction: FieldElement, + is_singles: bool, + parent_num_vars: usize, + /// The instance-specific evaluation point derived from the shared current_point. + /// Used for r_round lookups in the Dao-Thaler eq factoring. + instance_point: Vec>, + // SVO (Split-Value Optimization) fields. + // When use_svo is true, eq is split into prefix × suffix for sqrt memory. + use_svo: bool, + svo_suffix_len: usize, + eq_prefix: Vec>, + eq_suffix: Vec>, +} + +// ============================================================================= +// Batch GKR verifier +// ============================================================================= + +/// Verify a batch GKR proof. +/// +/// Replays the Fiat-Shamir transcript and checks sumcheck consistency and gate +/// equations for all instances simultaneously. +/// +/// # Returns +/// `Ok((shared_random_point, per_instance_claims))` where per_instance_claims[i] = (n_claim, d_claim). +#[allow(clippy::type_complexity)] +pub fn gkr_verify_batch( + proof: &BatchGkrProof, + n_layers_by_instance: &[usize], + transcript: &mut impl IsTranscript, +) -> Result< + ( + Vec>, + Vec<(FieldElement, FieldElement)>, + ), + GkrError, +> { + let n_instances = proof.root_claims.len(); + if n_layers_by_instance.len() != n_instances { + return Err(GkrError::InvalidTree { + reason: "n_layers_by_instance length mismatch".to_string(), + }); + } + + // Domain separation (must match prover) + transcript.append_bytes(b"gkr_batch"); + transcript.append_bytes(&(n_instances as u64).to_le_bytes()); + + if n_instances == 0 { + return Ok((vec![], vec![])); + } + + let max_layers = *n_layers_by_instance.iter().max().unwrap(); + + if proof.layer_proofs.len() != max_layers { + return Err(GkrError::InvalidTree { + reason: format!( + "expected {} layer proofs but got {}", + max_layers, + proof.layer_proofs.len(), + ), + }); + } + + // Track per-instance state + let mut n_claims: Vec>> = vec![None; n_instances]; + let mut d_claims: Vec>> = vec![None; n_instances]; + let mut current_point: Vec> = vec![]; + + for (layer_idx, layer_proof) in proof.layer_proofs.iter().enumerate() { + let n_remaining = max_layers - layer_idx; + + // Detect output layers — use actual (root_n, root_d) from the proof + for i in 0..n_instances { + if n_layers_by_instance[i] == n_remaining { + n_claims[i] = Some(proof.root_claims[i].0.clone()); + d_claims[i] = Some(proof.root_claims[i].1.clone()); + } + } + + // Append active claims to transcript + for i in 0..n_instances { + if let (Some(n), Some(d)) = (&n_claims[i], &d_claims[i]) { + transcript.append_field_element(n); + transcript.append_field_element(d); + } + } + + // Sample randomness + let sumcheck_alpha: FieldElement = transcript.sample_field_element(); + let lambda: FieldElement = transcript.sample_field_element(); + + // Collect active instances + let mut active_instances: Vec = Vec::new(); + let mut combined_claims: Vec> = Vec::new(); + + for i in 0..n_instances { + if n_claims[i].is_some() && n_layers_by_instance[i] > 0 { + active_instances.push(i); + let n = n_claims[i].as_ref().unwrap(); + let d = d_claims[i].as_ref().unwrap(); + let claim = n + &(&lambda * d); + let n_unused = max_layers - n_layers_by_instance[i]; + if n_unused > 0 { + let doubling = FieldElement::::from(1u64 << n_unused); + combined_claims.push(&claim * &doubling); + } else { + combined_claims.push(claim); + } + } + } + + let round_polys = &layer_proof.sumcheck_proof.round_polys; + + if layer_proof.child_claims_by_instance.len() < active_instances.len() { + return Err(GkrError::InvalidTree { + reason: format!( + "layer {}: child_claims_by_instance has {} entries but {} active instances", + layer_idx, + layer_proof.child_claims_by_instance.len(), + active_instances.len(), + ), + }); + } + + if round_polys.is_empty() { + // Trivial layer: no sumcheck rounds. + // Gate check: verify that the alpha-batched combined claims match the + // alpha-batched gate evaluations (gate_i = nl·dr + nr·dl + λ·dl·dr, scaled + // by 2^n_unused_i to match combined_claims). + { + let mut actual_sum = FieldElement::::zero(); + let mut expected_sum = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = + layer_proof.child_claims_by_instance[idx]; + let gate = &(&(nl * dr) + &(nr * dl)) + &(&lambda * &(dl * dr)); + let n_unused = max_layers - n_layers_by_instance[i]; + let gate_scaled = if n_unused > 0 { + &gate * &FieldElement::::from(1u64 << n_unused) + } else { + gate + }; + actual_sum = &actual_sum + &(&alpha_pow * &combined_claims[idx]); + expected_sum = &expected_sum + &(&alpha_pow * &gate_scaled); + alpha_pow = &alpha_pow * &sumcheck_alpha; + } + if actual_sum != expected_sum { + return Err(GkrError::GateCheckFailed { layer: layer_idx }); + } + } + + // Append child claims and sample eta (same as prover) + for claims in &layer_proof.child_claims_by_instance { + for claim in claims { + transcript.append_field_element(claim); + } + } + + let eta: FieldElement = transcript.sample_field_element(); + + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims_by_instance[idx]; + n_claims[i] = Some(nl + &(&eta * &(nr - nl))); + d_claims[i] = Some(dl + &(&eta * &(dr - dl))); + } + + current_point = vec![eta]; + } else { + // Non-trivial: verify sumcheck + let num_rounds = round_polys.len(); + + // Compute combined claim across instances via sumcheck_alpha + let mut current_sum = { + let mut sum = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + for claim in &combined_claims { + sum = &sum + &(&alpha_pow * claim); + alpha_pow = &alpha_pow * &sumcheck_alpha; + } + sum + }; + + let mut challenges = Vec::with_capacity(num_rounds); + + for (round, round_poly) in round_polys.iter().enumerate() { + if round_poly.sum_at_binary() != current_sum { + return Err(GkrError::SumcheckFailed { + layer: layer_idx, + reason: format!("round {} sum mismatch: p(0)+p(1) != expected sum", round), + }); + } + + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + let challenge: FieldElement = transcript.sample_field_element(); + current_sum = round_poly.evaluate(&challenge); + challenges.push(challenge); + } + + // Gate check: for each active instance, verify the gate equation + let mut expected_sum = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims_by_instance[idx]; + + // parent_num_vars for this instance at this layer: + // instance i's child layer has (n_layers[i] - n_remaining + 1) variables, + // so the parent has (n_layers[i] - n_remaining) variables. + let parent_num_vars_i = n_layers_by_instance[i] - n_remaining; + if num_rounds < parent_num_vars_i { + return Err(GkrError::InvalidTree { + reason: format!( + "layer {}: num_rounds ({}) < parent_num_vars for instance {} ({})", + layer_idx, num_rounds, i, parent_num_vars_i, + ), + }); + } + let sumcheck_n_unused = num_rounds - parent_num_vars_i; + + // eq evaluation: the prover builds eq from the instance-specific eval point + // (eta + last challenges), and the active sumcheck challenges are the last ones. + let eq_val = if parent_num_vars_i == 0 { + FieldElement::::one() + } else { + let inst_point = instance_eval_point(¤t_point, parent_num_vars_i); + compute_eq_at_point(&inst_point, &challenges[sumcheck_n_unused..]) + }; + + // gate_combined = nl*dr + nr*dl + lambda*dl*dr + let gate_combined = &(&(nl * dr) + &(nr * dl)) + &(&lambda * &(dl * dr)); + + // No doubling factor here: the 2^n_unused doubling is in the initial + // combined claim and gets cancelled by the unused rounds' halving. + // The point evaluation at the final challenge is just eq * gate. + let instance_eval = &eq_val * &gate_combined; + + expected_sum = &expected_sum + &(&alpha_pow * &instance_eval); + alpha_pow = &alpha_pow * &sumcheck_alpha; + } + + if current_sum != expected_sum { + return Err(GkrError::GateCheckFailed { layer: layer_idx }); + } + + // Append child claims to transcript + for claims in &layer_proof.child_claims_by_instance { + for claim in claims { + transcript.append_field_element(claim); + } + } + + let eta: FieldElement = transcript.sample_field_element(); + + for (idx, &i) in active_instances.iter().enumerate() { + let [ref nl, ref nr, ref dl, ref dr] = layer_proof.child_claims_by_instance[idx]; + n_claims[i] = Some(nl + &(&eta * &(nr - nl))); + d_claims[i] = Some(dl + &(&eta * &(dr - dl))); + } + + let mut new_point = Vec::with_capacity(challenges.len() + 1); + new_point.push(eta); + new_point.extend(challenges); + current_point = new_point; + } + } + + let final_claims: Vec<(FieldElement, FieldElement)> = (0..n_instances) + .map(|i| { + ( + n_claims[i] + .clone() + .unwrap_or_else(|| proof.root_claims[i].0.clone()), + d_claims[i] + .clone() + .unwrap_or_else(|| proof.root_claims[i].1.clone()), + ) + }) + .collect(); + + Ok((current_point, final_claims)) +} + +/// Compute eq(a, b) for two points of equal length. +/// +/// eq(a, b) = prod_i (a_i * b_i + (1 - a_i) * (1 - b_i)) +/// +/// This is a single field element, NOT the full eq table. +fn compute_eq_at_point( + a: &[FieldElement], + b: &[FieldElement], +) -> FieldElement { + assert_eq!(a.len(), b.len(), "eq points must have equal length"); + let one = FieldElement::::one(); + a.iter() + .zip(b.iter()) + .fold(FieldElement::one(), |acc, (ai, bi)| { + let term = &(ai * bi) + &(&(&one - ai) * &(&one - bi)); + &acc * &term + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + + #[test] + fn test_fraction_add() { + // (3/5) + (7/11) = (3*11 + 7*5) / (5*11) = (33 + 35) / 55 = 68/55 + let a = Fraction::new(FE::from(3u64), FE::from(5u64)); + let b = Fraction::new(FE::from(7u64), FE::from(11u64)); + let result = a.add(&b); + + assert_eq!(result.numerator, FE::from(68u64)); + assert_eq!(result.denominator, FE::from(55u64)); + } + + #[test] + fn test_build_summation_tree_4_leaves() { + // 4 leaves: 1/2, 3/4, 5/6, 7/8 + // Layer 0 (4 fractions): [1/2, 3/4, 5/6, 7/8] + // Layer 1 (2 fractions): + // pair 0: 1/2 + 3/4 = (1*4 + 3*2)/(2*4) = 10/8 + // pair 1: 5/6 + 7/8 = (5*8 + 7*6)/(6*8) = 82/48 + // Layer 2 (1 fraction, root): + // 10/8 + 82/48 = (10*48 + 82*8)/(8*48) = (480 + 656)/384 = 1136/384 + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + + let tree = build_summation_tree(nums, dens); + + // Should have 3 layers: 4 -> 2 -> 1 + assert_eq!(tree.len(), 3); + assert_eq!(tree[0].numerators.len(), 4); + assert_eq!(tree[1].numerators.len(), 2); + assert_eq!(tree[2].numerators.len(), 1); + + // Layer 1 checks + assert_eq!(tree[1].numerators[0], FE::from(10u64)); // 1*4 + 3*2 + assert_eq!(tree[1].denominators[0], FE::from(8u64)); // 2*4 + assert_eq!(tree[1].numerators[1], FE::from(82u64)); // 5*8 + 7*6 + assert_eq!(tree[1].denominators[1], FE::from(48u64)); // 6*8 + + // Root checks + assert_eq!(tree[2].numerators[0], FE::from(1136u64)); // 10*48 + 82*8 + assert_eq!(tree[2].denominators[0], FE::from(384u64)); // 8*48 + + // Verify the root fraction equals the actual sum: + // 1/2 + 3/4 + 5/6 + 7/8 = 12/24 + 18/24 + 20/24 + 21/24 = 71/24 + // Check: 1136/384 = 71/24 (both sides: 1136*24 = 27264, 71*384 = 27264) + let root_n = &tree[2].numerators[0]; + let root_d = &tree[2].denominators[0]; + assert_eq!(root_n * &FE::from(24u64), &FE::from(71u64) * root_d); + } + + #[test] + fn test_build_summation_tree_8_leaves() { + // 8 leaves: i/(i+1) for i in 1..=8, i.e., 1/2, 2/3, 3/4, 4/5, 5/6, 6/7, 7/8, 8/9 + let nums: Vec = (1..=8).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (2..=9).map(|i| FE::from(i as u64)).collect(); + + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Should have 4 layers: 8 -> 4 -> 2 -> 1 + assert_eq!(tree.len(), 4); + assert_eq!(tree[0].numerators.len(), 8); + assert_eq!(tree[1].numerators.len(), 4); + assert_eq!(tree[2].numerators.len(), 2); + assert_eq!(tree[3].numerators.len(), 1); + + // Compute expected root by sequential fraction addition + let mut acc = Fraction::new(nums[0], dens[0]); + for i in 1..8 { + acc = acc.add(&Fraction::new(nums[i], dens[i])); + } + + // The tree root and the sequential sum should represent the same rational number: + // tree_n / tree_d == acc_n / acc_d <=> tree_n * acc_d == acc_n * tree_d + let root_n = &tree[3].numerators[0]; + let root_d = &tree[3].denominators[0]; + assert_eq!( + root_n * &acc.denominator, + &acc.numerator * root_d, + "Tree root must equal sequential sum as a fraction" + ); + } + + #[test] + fn test_build_summation_tree_single_leaf() { + // Edge case: 1 leaf (2^0 = 1) + let nums = vec![FE::from(42u64)]; + let dens = vec![FE::from(7u64)]; + + let tree = build_summation_tree(nums, dens); + + // Should have 1 layer (just the leaf) + assert_eq!(tree.len(), 1); + assert_eq!(tree[0].numerators.len(), 1); + assert_eq!(tree[0].numerators[0], FE::from(42u64)); + assert_eq!(tree[0].denominators[0], FE::from(7u64)); + } + + #[test] + #[should_panic(expected = "number of leaves must be a power of 2")] + fn test_build_summation_tree_non_power_of_2_panics() { + let nums = vec![FE::from(1u64), FE::from(2u64), FE::from(3u64)]; + let dens = vec![FE::from(1u64), FE::from(1u64), FE::from(1u64)]; + let _ = build_summation_tree(nums, dens); + } + + #[test] + #[should_panic(expected = "numerators and denominators must have the same length")] + fn test_build_summation_tree_mismatched_lengths_panics() { + let nums = vec![FE::from(1u64), FE::from(2u64)]; + let dens = vec![FE::from(1u64)]; + let _ = build_summation_tree(nums, dens); + } + + // ==================== compute_eq_evals tests ==================== + + #[test] + fn test_compute_eq_evals_empty_point() { + // eq with 0 variables: single entry = 1 + let evals = compute_eq_evals::(&[]); + assert_eq!(evals.len(), 1); + assert_eq!(evals[0], FE::one()); + } + + #[test] + fn test_compute_eq_evals_1var() { + // eq((r,), (b,)) = r*b + (1-r)*(1-b) + // For r = 3: eq(3, 0) = (1-3) = -2, eq(3, 1) = 3 + let r = FE::from(3u64); + let evals = compute_eq_evals(std::slice::from_ref(&r)); + assert_eq!(evals.len(), 2); + assert_eq!(evals[0], FE::one() - r); // eq(r, 0) = 1-r = -2 + assert_eq!(evals[1], r); // eq(r, 1) = r = 3 + } + + #[test] + fn test_compute_eq_evals_2var() { + // eq((r0, r1), (b0, b1)) = [r0*b0 + (1-r0)*(1-b0)] * [r1*b1 + (1-r1)*(1-b1)] + let r0 = FE::from(2u64); + let r1 = FE::from(5u64); + let evals = compute_eq_evals(&[r0, r1]); + assert_eq!(evals.len(), 4); + + let one = FE::one(); + let one_minus_r0 = &one - &r0; + let one_minus_r1 = &one - &r1; + + // Index 0 = (b0=0, b1=0): (1-r0)*(1-r1) + assert_eq!(evals[0], &one_minus_r0 * &one_minus_r1); + // Index 1 = (b0=1, b1=0): r0*(1-r1) + assert_eq!(evals[1], &r0 * &one_minus_r1); + // Index 2 = (b0=0, b1=1): (1-r0)*r1 + assert_eq!(evals[2], &one_minus_r0 * &r1); + // Index 3 = (b0=1, b1=1): r0*r1 + assert_eq!(evals[3], &r0 * &r1); + } + + #[test] + fn test_compute_eq_evals_sum_to_one_on_booleans() { + // When point is Boolean, eq_evals should have exactly one 1 and rest 0 + let point = vec![FE::one(), FE::zero(), FE::one()]; // b = (1, 0, 1) = index 5 + let evals = compute_eq_evals(&point); + assert_eq!(evals.len(), 8); + for (i, e) in evals.iter().enumerate() { + if i == 5 { + assert_eq!(*e, FE::one()); + } else { + assert_eq!(*e, FE::zero()); + } + } + } + + // ==================== evaluate_mle tests ==================== + + #[test] + fn test_evaluate_mle_linear() { + // MLE of [3, 7] at point r: + // f(x) = 3*(1-x) + 7*x = 3 + 4x + // f(5) = 3 + 20 = 23 + let table = vec![FE::from(3u64), FE::from(7u64)]; + let result = evaluate_mle(&table, &[FE::from(5u64)]); + assert_eq!(result, FE::from(23u64)); + } + + #[test] + fn test_evaluate_mle_at_boolean() { + // MLE at a Boolean point should return the table entry + let table = vec![ + FE::from(10u64), + FE::from(20u64), + FE::from(30u64), + FE::from(40u64), + ]; + // Index 2 = (b0=0, b1=1) + let result = evaluate_mle(&table, &[FE::zero(), FE::one()]); + assert_eq!(result, FE::from(30u64)); + } + + // ==================== GKR prover tests ==================== + + #[test] + fn test_gkr_prove_2_leaves() { + // Simplest non-trivial case: 2 leaves + // Tree: layer 0 (2 fractions), layer 1 (root, 1 fraction) + // Leaves: 3/5, 7/11 + // Root: (3*11 + 7*5) / (5*11) = 68/55 + let nums = vec![FE::from(3u64), FE::from(7u64)]; + let dens = vec![FE::from(5u64), FE::from(11u64)]; + let tree = build_summation_tree(nums, dens); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, final_point, final_n_claim, final_d_claim) = + gkr_prove(&tree, &mut transcript).unwrap(); + + // claimed_sum = 68/55 + let expected_sum = &FE::from(68u64) * &FE::from(55u64).inv().unwrap(); + assert_eq!(proof.claimed_sum, expected_sum); + + // Should have 1 layer proof (root -> leaves) + assert_eq!(proof.layer_proofs.len(), 1); + + // The first (and only) layer proof has a trivial sumcheck (0 rounds) + assert_eq!(proof.layer_proofs[0].sumcheck_proof.round_polys.len(), 0); + + // The child claims should be the raw leaf values + assert_eq!(proof.layer_proofs[0].child_claims[0], FE::from(3u64)); // n_left + assert_eq!(proof.layer_proofs[0].child_claims[1], FE::from(7u64)); // n_right + assert_eq!(proof.layer_proofs[0].child_claims[2], FE::from(5u64)); // d_left + assert_eq!(proof.layer_proofs[0].child_claims[3], FE::from(11u64)); // d_right + + // final_point should have 1 element (eta) + assert_eq!(final_point.len(), 1); + + // Final claims should be the leaf MLE evaluated at final_point + // n_MLE(eta) = 3*(1-eta) + 7*eta = 3 + 4*eta + // d_MLE(eta) = 5*(1-eta) + 11*eta = 5 + 6*eta + let eta = &final_point[0]; + let expected_n = &FE::from(3u64) * &(&FE::one() - eta) + &(&FE::from(7u64) * eta); + let expected_d = &FE::from(5u64) * &(&FE::one() - eta) + &(&FE::from(11u64) * eta); + assert_eq!(final_n_claim, expected_n); + assert_eq!(final_d_claim, expected_d); + } + + #[test] + fn test_gkr_prove_4_leaves() { + // 4 leaves: 1/2, 3/4, 5/6, 7/8 + // Tree has 3 layers (0=leaves size 4, 1=size 2, 2=root size 1) + // GKR reduces: root -> layer 1 (trivial, 0 vars) -> layer 0 (1-var sumcheck) + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums.clone(), dens.clone()); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, final_point, final_n_claim, final_d_claim) = + gkr_prove(&tree, &mut transcript).unwrap(); + + // claimed_sum = root_n / root_d = 1136 / 384 + let expected_sum = &FE::from(1136u64) * &FE::from(384u64).inv().unwrap(); + assert_eq!(proof.claimed_sum, expected_sum); + + // Should have 2 layer proofs + assert_eq!(proof.layer_proofs.len(), 2); + + // First layer proof: root (1 elem) -> layer 1 (2 elems), trivial (0 rounds) + assert_eq!(proof.layer_proofs[0].sumcheck_proof.round_polys.len(), 0); + + // Second layer proof: layer 1 (2 elems) -> layer 0 (4 elems) + // This is a 1-variable sumcheck, so should have 1 round polynomial + assert_eq!(proof.layer_proofs[1].sumcheck_proof.round_polys.len(), 1); + + // The round polynomial should have 4 evaluations (degree 3) + assert_eq!( + proof.layer_proofs[1].sumcheck_proof.round_polys[0].num_evals(), + 4 + ); + + // final_point should have 2 elements (challenge from sumcheck + eta) + assert_eq!(final_point.len(), 2); + + // Verify final claims match the leaf MLEs at final_point + let expected_n = evaluate_mle(&nums, &final_point); + let expected_d = evaluate_mle(&dens, &final_point); + assert_eq!(final_n_claim, expected_n); + assert_eq!(final_d_claim, expected_d); + } + + #[test] + fn test_gkr_prove_claimed_sum() { + // Verify that proof.claimed_sum equals root_n * root_d.inv() + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums, dens); + + let root_n = &tree[2].numerators[0]; + let root_d = &tree[2].denominators[0]; + let expected_sum = root_n * &root_d.inv().unwrap(); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, _, _, _) = gkr_prove(&tree, &mut transcript).unwrap(); + + assert_eq!(proof.claimed_sum, expected_sum); + } + + #[test] + fn test_gkr_prove_8_leaves() { + // 8 leaves: i/(i+1) for i in 1..=8 + // Tree: 4 layers (sizes 8, 4, 2, 1) + // GKR reductions: root->layer2 (trivial), layer2->layer1 (1-var), layer1->layer0 (2-var) + let nums: Vec = (1..=8).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (2..=9).map(|i| FE::from(i as u64)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, final_point, final_n_claim, final_d_claim) = + gkr_prove(&tree, &mut transcript).unwrap(); + + // Should have 3 layer proofs + assert_eq!(proof.layer_proofs.len(), 3); + + // Layer 0: root (1) -> layer 2 (2): trivial + assert_eq!(proof.layer_proofs[0].sumcheck_proof.round_polys.len(), 0); + + // Layer 1: layer 2 (2) -> layer 1 (4): 1-variable sumcheck + assert_eq!(proof.layer_proofs[1].sumcheck_proof.round_polys.len(), 1); + + // Layer 2: layer 1 (4) -> layer 0 (8): 2-variable sumcheck + assert_eq!(proof.layer_proofs[2].sumcheck_proof.round_polys.len(), 2); + + // final_point should have 3 elements + assert_eq!(final_point.len(), 3); + + // Verify final claims match the leaf MLEs at final_point + let expected_n = evaluate_mle(&nums, &final_point); + let expected_d = evaluate_mle(&dens, &final_point); + assert_eq!(final_n_claim, expected_n); + assert_eq!(final_d_claim, expected_d); + } + + #[test] + fn test_gkr_prove_16_leaves() { + // 16 leaves with various fractions + // Tree: 5 layers (sizes 16, 8, 4, 2, 1) + let nums: Vec = (1..=16).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (17..=32).map(|i| FE::from(i as u64)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + let mut transcript = DefaultTranscript::::new(&[0xAB]); + let (proof, final_point, final_n_claim, final_d_claim) = + gkr_prove(&tree, &mut transcript).unwrap(); + + // Should have 4 layer proofs + assert_eq!(proof.layer_proofs.len(), 4); + + // final_point should have 4 elements (log2(16)) + assert_eq!(final_point.len(), 4); + + // Verify final claims match the leaf MLEs at final_point + let expected_n = evaluate_mle(&nums, &final_point); + let expected_d = evaluate_mle(&dens, &final_point); + assert_eq!(final_n_claim, expected_n); + assert_eq!(final_d_claim, expected_d); + } + + #[test] + fn test_gkr_prove_deterministic() { + // Same inputs and transcript seed should produce identical proofs + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums, dens); + + let mut t1 = DefaultTranscript::::new(&[0x42]); + let (proof1, point1, n1, d1) = gkr_prove(&tree, &mut t1).unwrap(); + + let mut t2 = DefaultTranscript::::new(&[0x42]); + let (proof2, point2, n2, d2) = gkr_prove(&tree, &mut t2).unwrap(); + + assert_eq!(proof1.claimed_sum, proof2.claimed_sum); + assert_eq!(point1, point2); + assert_eq!(n1, n2); + assert_eq!(d1, d2); + assert_eq!(proof1.layer_proofs.len(), proof2.layer_proofs.len()); + + for (lp1, lp2) in proof1.layer_proofs.iter().zip(proof2.layer_proofs.iter()) { + assert_eq!(lp1.child_claims, lp2.child_claims); + assert_eq!( + lp1.sumcheck_proof.round_polys.len(), + lp2.sumcheck_proof.round_polys.len() + ); + for (rp1, rp2) in lp1 + .sumcheck_proof + .round_polys + .iter() + .zip(lp2.sumcheck_proof.round_polys.iter()) + { + assert_eq!(rp1.evals(), rp2.evals()); + } + } + } + + #[test] + fn test_gkr_prove_sumcheck_consistency() { + // Verify that the round polynomial p(0) + p(1) matches the combined claim + // for the non-trivial sumcheck layers + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Re-run the protocol manually to extract the combined_claim at each layer + // and verify the sumcheck round poly sum + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, _, _, _) = gkr_prove(&tree, &mut transcript).unwrap(); + + // Replay transcript to get the same challenges + let mut replay = DefaultTranscript::::new(&[]); + replay.append_field_element(&proof.claimed_sum); + + let mut n_claim = tree[2].numerators[0]; + let mut d_claim = tree[2].denominators[0]; + + for (layer_idx, lp) in proof.layer_proofs.iter().enumerate() { + let lambda: FE = replay.sample_field_element(); + let combined_claim = &n_claim + &(&lambda * &d_claim); + + if lp.sumcheck_proof.round_polys.is_empty() { + // Trivial layer: verify gate equation directly + let nl = &lp.child_claims[0]; + let nr = &lp.child_claims[1]; + let dl = &lp.child_claims[2]; + let dr = &lp.child_claims[3]; + let gate_val = &(nl * dr) + &(nr * dl) + &(&lambda * &(dl * dr)); + assert_eq!( + combined_claim, gate_val, + "Gate equation failed at trivial layer {}", + layer_idx + ); + } else { + // Non-trivial layer: verify p(0) + p(1) = combined_claim + let first_round = &lp.sumcheck_proof.round_polys[0]; + assert_eq!( + first_round.sum_at_binary(), + combined_claim, + "Sumcheck round sum mismatch at layer {}", + layer_idx + ); + } + + // Replay transcript operations from the sumcheck + for rp in &lp.sumcheck_proof.round_polys { + for eval in rp.evals() { + replay.append_field_element(eval); + } + let _challenge: FE = replay.sample_field_element(); + } + + // Replay child claims and eta + for claim in &lp.child_claims { + replay.append_field_element(claim); + } + let eta: FE = replay.sample_field_element(); + + // Update claims for next layer + let one_minus_eta = &FE::one() - η + n_claim = &(&lp.child_claims[0] * &one_minus_eta) + &(&lp.child_claims[1] * &eta); + d_claim = &(&lp.child_claims[2] * &one_minus_eta) + &(&lp.child_claims[3] * &eta); + } + } + + #[test] + fn test_gkr_prove_single_leaf() { + // Edge case: single leaf (tree has 1 layer, no reductions) + let nums = vec![FE::from(42u64)]; + let dens = vec![FE::from(7u64)]; + let tree = build_summation_tree(nums, dens); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, final_point, final_n_claim, final_d_claim) = + gkr_prove(&tree, &mut transcript).unwrap(); + + assert_eq!(proof.claimed_sum, FE::from(6u64)); // 42/7 = 6 + assert_eq!(proof.layer_proofs.len(), 0); + assert!(final_point.is_empty()); + assert_eq!(final_n_claim, FE::from(42u64)); + assert_eq!(final_d_claim, FE::from(7u64)); + } + + // ==================== GKR verifier tests ==================== + + #[test] + fn test_gkr_prove_verify_roundtrip_4() { + // 4 leaves: 1/2, 3/4, 5/6, 7/8 + // Tree has 3 layers (0=leaves size 4, 1=size 2, 2=root size 1) + // GKR reduces: root -> layer 1 (trivial) -> layer 0 (1-var sumcheck) + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[0xAA]); + let (proof, prover_point, prover_n, prover_d) = + gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // Verify with a fresh transcript (same seed) + let mut verifier_transcript = DefaultTranscript::::new(&[0xAA]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + result.is_ok(), + "GKR verification should succeed for 4 leaves" + ); + let (verifier_point, verifier_n, verifier_d) = result.unwrap(); + + // The verifier's final point must match the prover's + assert_eq!( + verifier_point, prover_point, + "Verifier and prover must derive the same final point" + ); + + // The verifier's leaf claims must match the prover's + assert_eq!(verifier_n, prover_n, "n_claim must match"); + assert_eq!(verifier_d, prover_d, "d_claim must match"); + + // Additionally verify that the claims are consistent with the leaf MLEs + let expected_n = evaluate_mle(&nums, &verifier_point); + let expected_d = evaluate_mle(&dens, &verifier_point); + assert_eq!(verifier_n, expected_n, "n_claim must match leaf MLE"); + assert_eq!(verifier_d, expected_d, "d_claim must match leaf MLE"); + } + + #[test] + fn test_gkr_prove_verify_roundtrip_8() { + // 8 leaves: i/(i+1) for i in 1..=8 + // Tree: 4 layers (sizes 8, 4, 2, 1) + // GKR: root->layer2 (trivial), layer2->layer1 (1-var), layer1->layer0 (2-var) + let nums: Vec = (1..=8).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (2..=9).map(|i| FE::from(i as u64)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[0xAA]); + let (proof, prover_point, prover_n, prover_d) = + gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // Verify with a fresh transcript (same seed) + let mut verifier_transcript = DefaultTranscript::::new(&[0xAA]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + result.is_ok(), + "GKR verification should succeed for 8 leaves" + ); + let (verifier_point, verifier_n, verifier_d) = result.unwrap(); + + // The verifier's final point must match the prover's + assert_eq!(verifier_point, prover_point); + + // The verifier's leaf claims must match the prover's + assert_eq!(verifier_n, prover_n); + assert_eq!(verifier_d, prover_d); + + // Verify consistency with leaf MLEs + let expected_n = evaluate_mle(&nums, &verifier_point); + let expected_d = evaluate_mle(&dens, &verifier_point); + assert_eq!(verifier_n, expected_n); + assert_eq!(verifier_d, expected_d); + } + + #[test] + fn test_gkr_verify_wrong_claimed_sum() { + // Create a valid proof and then tamper with the claimed_sum. + // The verifier should fail (either at sumcheck or gate check). + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums, dens); + + // Prove with correct claimed_sum + let mut prover_transcript = DefaultTranscript::::new(&[0xAA]); + let (mut proof, _, _, _) = gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // Tamper with the claimed_sum + proof.claimed_sum = &proof.claimed_sum + &FE::one(); + + // Verify with the tampered proof + let mut verifier_transcript = DefaultTranscript::::new(&[0xAA]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + // The verification should fail. The tampered claimed_sum changes the + // transcript state, which changes lambda and eta, leading to a sumcheck + // or gate check failure at the non-trivial layer. + assert!( + result.is_err(), + "GKR verification should fail with tampered claimed_sum" + ); + } + + #[test] + fn test_gkr_verify_trivial_layer_gate_check_rejected() { + // A 4-leaf tree produces: + // layer_proofs[0]: trivial (root→layer1, parent_num_vars=0, round_polys=[]) + // layer_proofs[1]: non-trivial (layer1→leaves, parent_num_vars=1) + // + // Tamper with child_claims of the trivial layer_proofs[0] and assert + // that gkr_verify returns GkrError::GateCheckFailed. + let nums = vec![ + FE::from(1u64), + FE::from(3u64), + FE::from(5u64), + FE::from(7u64), + ]; + let dens = vec![ + FE::from(2u64), + FE::from(4u64), + FE::from(6u64), + FE::from(8u64), + ]; + let tree = build_summation_tree(nums, dens); + + let mut prover_transcript = DefaultTranscript::::new(&[0xC0]); + let (mut proof, _, _, _) = gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // layer_proofs[0] is the trivial layer: round_polys must be empty + assert!( + proof.layer_proofs[0].sumcheck_proof.round_polys.is_empty(), + "layer_proofs[0] must be the trivial layer for a 4-leaf tree" + ); + + // Corrupt nl (child_claims[0]) by adding 1 — this breaks the gate equation + // n_claim*(dl*dr) = nl*dr + nr*dl without altering the transcript ordering + proof.layer_proofs[0].child_claims[0] += FE::one(); + + let mut verifier_transcript = DefaultTranscript::::new(&[0xC0]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + matches!(result, Err(GkrError::GateCheckFailed { layer: 0 })), + "Tampered trivial-layer child_claims must be rejected with GateCheckFailed {{ layer: 0 }}, got: {:?}", + result + ); + } + + #[test] + fn test_gkr_prove_verify_roundtrip_16() { + // 16 leaves with various fractions + // Tree: 5 layers (sizes 16, 8, 4, 2, 1) + let nums: Vec = (1..=16).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (17..=32).map(|i| FE::from(i as u64)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[0xAA]); + let (proof, prover_point, prover_n, prover_d) = + gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // Verify with a fresh transcript (same seed) + let mut verifier_transcript = DefaultTranscript::::new(&[0xAA]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + result.is_ok(), + "GKR verification should succeed for 16 leaves" + ); + let (verifier_point, verifier_n, verifier_d) = result.unwrap(); + + // The verifier's final point must match the prover's + assert_eq!(verifier_point, prover_point); + + // The verifier's leaf claims must match the prover's + assert_eq!(verifier_n, prover_n); + assert_eq!(verifier_d, prover_d); + + // Verify consistency with leaf MLEs + let expected_n = evaluate_mle(&nums, &verifier_point); + let expected_d = evaluate_mle(&dens, &verifier_point); + assert_eq!(verifier_n, expected_n); + assert_eq!(verifier_d, expected_d); + } + + // ==================== SVO (Split-Value Optimization) tests ==================== + + #[test] + fn test_gkr_prove_verify_roundtrip_512_svo() { + // 512 leaves: exercises the SVO path (parent_num_vars = 8 >= SVO_THRESHOLD) + // Tree: 10 layers (sizes 512, 256, 128, 64, 32, 16, 8, 4, 2, 1) + let n = 512; + let nums: Vec = (1..=n).map(|i| FE::from(i as u64)).collect(); + let dens: Vec = (1..=n).map(|i| FE::from(i as u64 + 1000)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[0x5F, 0x00]); + let (proof, prover_point, prover_n, prover_d) = + gkr_prove(&tree, &mut prover_transcript).unwrap(); + + // Verify with a fresh transcript (same seed) + let mut verifier_transcript = DefaultTranscript::::new(&[0x5F, 0x00]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + result.is_ok(), + "GKR verification should succeed for 512 leaves (SVO path): {:?}", + result.err() + ); + let (verifier_point, verifier_n, verifier_d) = result.unwrap(); + + // The verifier's final point must match the prover's + assert_eq!(verifier_point, prover_point); + + // The verifier's leaf claims must match the prover's + assert_eq!(verifier_n, prover_n); + assert_eq!(verifier_d, prover_d); + + // Verify consistency with leaf MLEs + let expected_n = evaluate_mle(&nums, &verifier_point); + let expected_d = evaluate_mle(&dens, &verifier_point); + assert_eq!(verifier_n, expected_n); + assert_eq!(verifier_d, expected_d); + } + + #[test] + fn test_gkr_prove_verify_roundtrip_1024_svo() { + // 1024 leaves: ensures SVO path is exercised at multiple layers + // parent_num_vars = 9 at the bottom layer + let n = 1024; + let nums: Vec = (1..=n).map(|i| FE::from(i as u64 * 3 + 1)).collect(); + let dens: Vec = (1..=n).map(|i| FE::from(i as u64 * 7 + 5)).collect(); + let tree = build_summation_tree(nums.clone(), dens.clone()); + + let mut prover_transcript = DefaultTranscript::::new(&[0xBB]); + let (proof, prover_point, prover_n, prover_d) = + gkr_prove(&tree, &mut prover_transcript).unwrap(); + + let mut verifier_transcript = DefaultTranscript::::new(&[0xBB]); + let result = gkr_verify(&proof, &mut verifier_transcript); + + assert!( + result.is_ok(), + "GKR verification should succeed for 1024 leaves (SVO path): {:?}", + result.err() + ); + let (verifier_point, verifier_n, verifier_d) = result.unwrap(); + + assert_eq!(verifier_point, prover_point); + assert_eq!(verifier_n, prover_n); + assert_eq!(verifier_d, prover_d); + + let expected_n = evaluate_mle(&nums, &verifier_point); + let expected_d = evaluate_mle(&dens, &verifier_point); + assert_eq!(verifier_n, expected_n); + assert_eq!(verifier_d, expected_d); + } + + // ==================== Batch GKR tests ==================== + + /// Helper: create a Layer::LogUpGeneric leaf from random-ish fractions. + fn make_generic_leaf(n_vars: usize) -> Layer { + let size = 1 << n_vars; + let denominators: Vec = (1..=size).map(|i| FE::from(i as u64 + 100)).collect(); + let numerators: Vec = (1..=size).map(|i| FE::from(i as u64)).collect(); + Layer::LogUpGeneric { + numerators, + denominators, + } + } + + /// Helper: create a Layer::LogUpSingles leaf. + fn make_singles_leaf(n_vars: usize) -> Layer { + let size = 1 << n_vars; + let denominators: Vec = (1..=size).map(|i| FE::from(i as u64 + 200)).collect(); + Layer::LogUpSingles { denominators } + } + + #[test] + fn test_batch_gkr_same_size_instances() { + // 3 instances, all with 4 leaves (n_vars=2, 3 layers each) + let instances: Vec>> = + (0..3).map(|_| gen_layers(make_generic_leaf(2))).collect(); + + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, shared_point, final_claims) = + gkr_prove_batch(instances, &mut prover_transcript); + + assert_eq!(proof.root_claims.len(), 3); + assert_eq!(final_claims.len(), 3); + + let n_layers: Vec = proof + .root_claims + .iter() + .map(|_| 2) // all same: 2 reduction steps + .collect(); + + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); + assert!(result.is_ok(), "batch verify failed: {:?}", result.err()); + + let (v_point, v_claims) = result.unwrap(); + assert_eq!(v_point, shared_point); + assert_eq!(v_claims, final_claims); + } + + #[test] + fn test_batch_gkr_mixed_size_instances() { + // This is the key test: 3 instances with DIFFERENT sizes. + // n_vars = 2 (4 leaves), 4 (16 leaves), 6 (64 leaves) + // This exercises the instance_eval_point logic. + let instances: Vec>> = vec![ + gen_layers(make_generic_leaf(2)), // 3 layers, 2 reductions + gen_layers(make_generic_leaf(4)), // 5 layers, 4 reductions + gen_layers(make_generic_leaf(6)), // 7 layers, 6 reductions + ]; + + let n_layers: Vec = instances.iter().map(|l| l.len() - 1).collect(); + assert_eq!(n_layers, vec![2, 4, 6]); + + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, shared_point, final_claims) = + gkr_prove_batch(instances, &mut prover_transcript); + + assert_eq!(proof.root_claims.len(), 3); + + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); + assert!( + result.is_ok(), + "mixed-size batch verify failed: {:?}", + result.err() + ); + + let (v_point, v_claims) = result.unwrap(); + assert_eq!(v_point, shared_point); + assert_eq!(v_claims, final_claims); + } + + #[test] + fn test_batch_gkr_mixed_size_with_singles() { + // Mix of Generic and Singles leaves with different sizes + let instances: Vec>> = vec![ + gen_layers(make_singles_leaf(2)), // 3 layers + gen_layers(make_generic_leaf(3)), // 4 layers + gen_layers(make_singles_leaf(5)), // 6 layers + gen_layers(make_generic_leaf(4)), // 5 layers + ]; + + let n_layers: Vec = instances.iter().map(|l| l.len() - 1).collect(); + + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, shared_point, final_claims) = + gkr_prove_batch(instances, &mut prover_transcript); + + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); + assert!( + result.is_ok(), + "mixed singles/generic batch verify failed: {:?}", + result.err() + ); + + let (v_point, v_claims) = result.unwrap(); + assert_eq!(v_point, shared_point); + assert_eq!(v_claims, final_claims); + } + + #[test] + fn test_batch_gkr_many_mixed_instances() { + // Stress test: 20 instances with sizes varying from n_vars=1 to n_vars=6 + // Mimics the real VM scenario with many tables of different sizes + let instances: Vec>> = (0..20) + .map(|i| { + let n_vars = (i % 6) + 1; // 1, 2, 3, 4, 5, 6, 1, 2, ... + gen_layers(make_generic_leaf(n_vars)) + }) + .collect(); + + let n_layers: Vec = instances.iter().map(|l| l.len() - 1).collect(); + + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, shared_point, final_claims) = + gkr_prove_batch(instances, &mut prover_transcript); + + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); + assert!( + result.is_ok(), + "many mixed instances batch verify failed: {:?}", + result.err() + ); + + let (v_point, v_claims) = result.unwrap(); + assert_eq!(v_point, shared_point); + assert_eq!(v_claims, final_claims); + } + + #[test] + fn test_instance_eval_point() { + // Verify the instance_eval_point helper + let point: Vec = vec![ + FE::from(10u64), // eta + FE::from(20u64), // c_0 + FE::from(30u64), // c_1 + FE::from(40u64), // c_2 + ]; + + // n_vars = 0: empty + assert!(instance_eval_point::(&point, 0).is_empty()); + + // n_vars = 4 (full): entire point + assert_eq!(instance_eval_point(&point, 4), point); + + // n_vars = 1: [eta] + assert_eq!(instance_eval_point(&point, 1), vec![FE::from(10u64)]); + + // n_vars = 2: [eta, c_2] (eta + last 1) + assert_eq!( + instance_eval_point(&point, 2), + vec![FE::from(10u64), FE::from(40u64)] + ); + + // n_vars = 3: [eta, c_1, c_2] (eta + last 2) + assert_eq!( + instance_eval_point(&point, 3), + vec![FE::from(10u64), FE::from(30u64), FE::from(40u64)] + ); + } +} diff --git a/crypto/stark/src/lagrange_kernel.rs b/crypto/stark/src/lagrange_kernel.rs new file mode 100644 index 000000000..59d5f9d5e --- /dev/null +++ b/crypto/stark/src/lagrange_kernel.rs @@ -0,0 +1,311 @@ +use math::field::{ + element::FieldElement, + traits::{IsField, IsSubFieldOf}, +}; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +/// Compute the Lagrange kernel (eq polynomial) for a random point `r`. +/// +/// Given r = (r_0, r_1, ..., r_{n-1}), returns a vector s of length N = 2^n where: +/// +/// s[i] = eq(bits(i), r) = prod_{l=0}^{n-1} [r_l * b_l(i) + (1 - r_l) * (1 - b_l(i))] +/// +/// where b_l(i) = (i >> l) & 1 is the l-th bit of i. +/// +/// The Lagrange kernel is the auxiliary column that bridges GKR multilinear extension +/// claims back to committed univariate traces. It satisfies the partition-of-unity +/// property: sum_{i=0}^{N-1} s[i] = 1 for any r. +/// +/// Uses the butterfly algorithm in O(N) time and O(N) space. +pub fn compute_lagrange_kernel(r: &[FieldElement]) -> Vec> { + let n = r.len(); + let big_n = 1usize << n; + let one = FieldElement::::one(); + + let mut s = vec![one.clone(); big_n]; + + #[allow(clippy::needless_range_loop)] + for j in 0..n { + let rj = &r[j]; + let one_minus_rj = &one - rj; + + #[cfg(feature = "parallel")] + { + if big_n >= 1024 { + s.par_iter_mut().enumerate().for_each(|(i, s_i)| { + if (i >> j) & 1 == 1 { + *s_i = &*s_i * rj; + } else { + *s_i = &*s_i * &one_minus_rj; + } + }); + } else { + for i in 0..big_n { + if (i >> j) & 1 == 1 { + s[i] = &s[i] * rj; + } else { + s[i] = &s[i] * &one_minus_rj; + } + } + } + } + + #[cfg(not(feature = "parallel"))] + { + for i in 0..big_n { + if (i >> j) & 1 == 1 { + s[i] = &s[i] * rj; + } else { + s[i] = &s[i] * &one_minus_rj; + } + } + } + } + + s +} + +/// Evaluate the multilinear extension (MLE) of `values` at point `r`. +/// +/// Given a function f defined on {0,1}^n by its truth table `values` (length N = 2^n), +/// its multilinear extension is: +/// +/// MLE(r) = sum_{i=0}^{N-1} values[i] * eq(bits(i), r) +/// +/// This is the inner product of `values` with the Lagrange kernel at `r`. +/// +/// Panics if `values.len()` is not a power of 2 or if `r.len() != log2(values.len())`. +pub fn eval_mle(values: &[FieldElement], r: &[FieldElement]) -> FieldElement { + let n = r.len(); + let big_n = 1usize << n; + assert_eq!( + values.len(), + big_n, + "values length ({}) must equal 2^r.len() = 2^{} = {}", + values.len(), + n, + big_n + ); + + let kernel = compute_lagrange_kernel(r); + + let mut result = FieldElement::::zero(); + for i in 0..big_n { + result = &result + &(&values[i] * &kernel[i]); + } + + result +} + +/// Evaluate the multilinear extension (MLE) of base-field `values` at an extension-field point `r`. +/// +/// Same as `eval_mle` but accepts base-field values and an extension-field evaluation point. +/// Uses `F * E -> E` multiplication directly (no `to_extension()` conversion). +/// +/// Panics if `values.len()` is not a power of 2 or if `r.len() != log2(values.len())`. +pub fn eval_mle_base(values: &[FieldElement], r: &[FieldElement]) -> FieldElement +where + F: IsField + IsSubFieldOf, + E: IsField, +{ + let n = r.len(); + let big_n = 1usize << n; + assert_eq!( + values.len(), + big_n, + "values length ({}) must equal 2^r.len() = 2^{} = {}", + values.len(), + n, + big_n + ); + + let kernel = compute_lagrange_kernel(r); + + let mut result = FieldElement::::zero(); + for i in 0..big_n { + // F * E -> E multiplication (no to_extension) + result = &result + &(&values[i] * &kernel[i]); + } + + result +} + +/// Evaluate the MLE of base-field `values` at an extension-field point, using a pre-computed +/// Lagrange kernel. +/// +/// This avoids recomputing `compute_lagrange_kernel(r)` when evaluating multiple columns +/// at the same point `r`. +/// +/// Panics if `values.len() != kernel.len()`. +pub fn eval_mle_base_with_kernel( + values: &[FieldElement], + kernel: &[FieldElement], +) -> FieldElement +where + F: IsField + IsSubFieldOf, + E: IsField + Send + Sync, + FieldElement: Send + Sync, +{ + let big_n = values.len(); + assert_eq!( + big_n, + kernel.len(), + "values length ({}) must equal kernel length ({})", + big_n, + kernel.len() + ); + + #[cfg(feature = "parallel")] + { + if big_n >= 1024 { + return (0..big_n) + .into_par_iter() + .fold( + || FieldElement::::zero(), + |acc, i| acc + &values[i] * &kernel[i], + ) + .reduce(|| FieldElement::::zero(), |a, b| a + b); + } + } + + // Sequential fallback + let mut result = FieldElement::::zero(); + for i in 0..big_n { + result = &result + &(&values[i] * &kernel[i]); + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + + #[test] + fn test_compute_lagrange_kernel_n2() { + // n=2, r = (3, 7) + // s[0] = (1-3)*(1-7) = (-2)*(-6) = 12 + // s[1] = 3*(1-7) = 3*(-6) = -18 + // s[2] = (1-3)*7 = (-2)*7 = -14 + // s[3] = 3*7 = 21 + let r = vec![FE::from(3u64), FE::from(7u64)]; + let s = compute_lagrange_kernel(&r); + + assert_eq!(s.len(), 4); + + // In Goldilocks (p = 2^64 - 2^32 + 1), negative values wrap around. + // 12 mod p = 12 + assert_eq!(s[0], FE::from(12u64)); + // -18 mod p = p - 18 + let neg_18 = FE::zero() - FE::from(18u64); + assert_eq!(s[1], neg_18); + // -14 mod p = p - 14 + let neg_14 = FE::zero() - FE::from(14u64); + assert_eq!(s[2], neg_14); + // 21 + assert_eq!(s[3], FE::from(21u64)); + + // Verify partition of unity: sum = 12 + (-18) + (-14) + 21 = 1 + let sum = &(&(&s[0] + &s[1]) + &s[2]) + &s[3]; + assert_eq!(sum, FE::one()); + } + + #[test] + fn test_lagrange_kernel_partition_of_unity() { + // For any random r, the sum of kernel values must equal 1. + // Test with several different r vectors. + for &(n, r_vals) in &[ + (1usize, &[42u64][..]), + (2, &[100, 200][..]), + (3, &[5, 17, 99][..]), + (4, &[1, 2, 3, 4][..]), + ] { + let r: Vec = r_vals.iter().map(|&v| FE::from(v)).collect(); + let s = compute_lagrange_kernel(&r); + assert_eq!(s.len(), 1 << n); + + let mut sum = FE::zero(); + for val in &s { + sum = &sum + val; + } + assert_eq!( + sum, + FE::one(), + "Partition of unity failed for n={}, r={:?}", + n, + r_vals + ); + } + } + + #[test] + fn test_eval_mle_linear() { + // f(x1) = 3*x1 + 5 on {0, 1} + // f(0) = 5, f(1) = 8 + // MLE at r: 5*(1-r) + 8*r = 5 + 3r + let values = vec![FE::from(5u64), FE::from(8u64)]; + + // Test at r = 0: should be 5 + let result = eval_mle(&values, &[FE::from(0u64)]); + assert_eq!(result, FE::from(5u64)); + + // Test at r = 1: should be 8 + let result = eval_mle(&values, &[FE::from(1u64)]); + assert_eq!(result, FE::from(8u64)); + + // Test at r = 10: should be 5 + 3*10 = 35 + let result = eval_mle(&values, &[FE::from(10u64)]); + assert_eq!(result, FE::from(35u64)); + + // Test at r = 100: should be 5 + 3*100 = 305 + let result = eval_mle(&values, &[FE::from(100u64)]); + assert_eq!(result, FE::from(305u64)); + } + + #[test] + fn test_eval_mle_2var() { + // f(x1, x2) = 2*x1*x2 + 3*x1 + x2 + 7 + // Evaluate on {0,1}^2 (bit ordering: x1 = bit 0, x2 = bit 1): + // i=0 (x1=0, x2=0): 0 + 0 + 0 + 7 = 7 + // i=1 (x1=1, x2=0): 0 + 3 + 0 + 7 = 10 + // i=2 (x1=0, x2=1): 0 + 0 + 1 + 7 = 8 + // i=3 (x1=1, x2=1): 2 + 3 + 1 + 7 = 13 + let values = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(8u64), + FE::from(13u64), + ]; + + // Verify at Boolean points first (should recover original values) + assert_eq!( + eval_mle(&values, &[FE::from(0u64), FE::from(0u64)]), + FE::from(7u64) + ); + assert_eq!( + eval_mle(&values, &[FE::from(1u64), FE::from(0u64)]), + FE::from(10u64) + ); + assert_eq!( + eval_mle(&values, &[FE::from(0u64), FE::from(1u64)]), + FE::from(8u64) + ); + assert_eq!( + eval_mle(&values, &[FE::from(1u64), FE::from(1u64)]), + FE::from(13u64) + ); + + // Evaluate at (x1=5, x2=3): + // f(5, 3) = 2*5*3 + 3*5 + 3 + 7 = 30 + 15 + 3 + 7 = 55 + let result = eval_mle(&values, &[FE::from(5u64), FE::from(3u64)]); + assert_eq!(result, FE::from(55u64)); + + // Evaluate at (x1=2, x2=4): + // f(2, 4) = 2*2*4 + 3*2 + 4 + 7 = 16 + 6 + 4 + 7 = 33 + let result = eval_mle(&values, &[FE::from(2u64), FE::from(4u64)]); + assert_eq!(result, FE::from(33u64)); + } +} diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index 6f8e7c82e..c2bd817c4 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -15,11 +15,13 @@ pub mod domain; pub mod examples; pub mod frame; pub mod fri; +pub mod gkr; #[cfg(feature = "cuda")] pub mod gpu_lde; pub mod grinding; #[cfg(feature = "instruments")] pub mod instruments; +pub mod lagrange_kernel; #[cfg(feature = "cuda")] pub mod logup_gpu; pub mod lookup; @@ -31,6 +33,7 @@ pub mod prover; pub mod r4_denoms; #[cfg(feature = "disk-spill")] pub mod storage_mode; +pub mod sumcheck; pub mod table; pub mod trace; pub mod traits; diff --git a/crypto/stark/src/sumcheck.rs b/crypto/stark/src/sumcheck.rs new file mode 100644 index 000000000..83eb0a386 --- /dev/null +++ b/crypto/stark/src/sumcheck.rs @@ -0,0 +1,830 @@ +use core::fmt; +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; + +/// A degree-d univariate polynomial represented by its evaluations at the +/// integer nodes 0, 1, ..., d. This is the polynomial that the sumcheck prover +/// sends to the verifier in each round. +/// +/// Storing evaluations (rather than coefficients) is natural for the sumcheck +/// protocol because: +/// - The prover constructs the polynomial by evaluating it at small integer points. +/// - The verifier only needs to check p(0) + p(1) and evaluate at a random challenge. +/// - Lagrange interpolation over integer nodes is cheap (small denominators). +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct RoundPoly { + /// Evaluations at x = 0, 1, ..., d where d = evals.len() - 1. + evals: Vec>, +} + +impl RoundPoly { + /// Create a new `RoundPoly` from evaluations at x = 0, 1, ..., d. + /// + /// `evals[i]` is the polynomial evaluated at x = i. + /// The polynomial has degree `evals.len() - 1`. + pub fn new(evals: Vec>) -> Self { + assert!( + !evals.is_empty(), + "RoundPoly must have at least one evaluation" + ); + Self { evals } + } + + /// Returns p(0) + p(1). + /// + /// In the sumcheck protocol, the verifier checks that this sum matches the + /// claimed sum for the current round. This is the key consistency check. + pub fn sum_at_binary(&self) -> FieldElement { + assert!( + self.evals.len() >= 2, + "sum_at_binary requires at least 2 evaluations (at 0 and 1)" + ); + &self.evals[0] + &self.evals[1] + } + + /// Evaluate the polynomial at an arbitrary field element using Lagrange + /// interpolation over integer nodes 0, 1, ..., d. + /// + /// Given evaluations y_0, ..., y_d at nodes 0, 1, ..., d, the Lagrange + /// interpolation formula is: + /// + /// p(x) = sum_{i=0}^{d} y_i * prod_{j != i} (x - j) / (i - j) + /// + /// The denominators (i - j) for integer nodes are just products of small + /// integers, so we precompute them as field elements. + pub fn evaluate(&self, point: &FieldElement) -> FieldElement { + let d = self.evals.len() - 1; + + // Special case: constant polynomial (degree 0) + if d == 0 { + return self.evals[0].clone(); + } + + // Precompute (point - j) for j = 0, ..., d + let point_minus_j: Vec> = (0..=d) + .map(|j| point - &FieldElement::from(j as u64)) + .collect(); + + // Check if point is one of the integer nodes (avoid division by zero) + for (j, pm) in point_minus_j.iter().enumerate() { + if *pm == FieldElement::zero() { + return self.evals[j].clone(); + } + } + + // Barycentric Lagrange interpolation with batch inversion. + // + // Barycentric weights: w_i = 1 / prod_{j≠i}(i - j) for integer nodes. + // We batch-invert the weights together with point_minus_j to use only + // 1 field inversion total (instead of 2(d+1) individual inversions). + + // Compute barycentric weight denominators: prod_{j≠i}(i - j) + let mut to_invert: Vec> = (0..=d) + .map(|i| { + let mut denom = FieldElement::::one(); + for j in 0..=d { + if j != i { + denom = &denom * &FieldElement::from((i as i64) - (j as i64)); + } + } + denom + }) + .collect(); + + // Append point_minus_j values; batch-invert everything in one call + to_invert.extend(point_minus_j.iter().cloned()); + FieldElement::inplace_batch_inverse(&mut to_invert).expect("All values are nonzero"); + + let (w_inv, pm_inv) = to_invert.split_at(d + 1); + + // master_product = prod_{j=0}^{d} (point - j) + let master_product = point_minus_j + .iter() + .fold(FieldElement::one(), |acc, v| &acc * v); + + // result = master_product * Σ_i (evals[i] * w_inv[i] * pm_inv[i]) + let mut sum = FieldElement::zero(); + for i in 0..=d { + sum = &sum + &(&self.evals[i] * &(&w_inv[i] * &pm_inv[i])); + } + + &master_product * &sum + } + + /// Returns the number of evaluations (degree + 1). + pub fn num_evals(&self) -> usize { + self.evals.len() + } + + /// Returns a reference to the evaluations. + pub fn evals(&self) -> &[FieldElement] { + &self.evals + } +} + +/// Proof produced by the sumcheck prover: one round polynomial per variable. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct SumcheckProof { + pub round_polys: Vec>, +} + +/// Run the sumcheck interactive proof (made non-interactive via Fiat-Shamir). +/// +/// Proves that `claimed_sum == sum_{x in {0,1}^num_vars} f(x)` where `evals` +/// holds the evaluations of f over the Boolean hypercube in little-endian bit +/// order (index = b_0 + 2*b_1 + ...). +/// +/// # Arguments +/// - `evals`: evaluations of f over {0,1}^num_vars, length must be 2^num_vars +/// - `claimed_sum`: the asserted sum +/// - `num_vars`: number of Boolean variables +/// - `max_degree`: maximum degree of round polynomials (need max_degree+1 evaluation points) +/// - `transcript`: Fiat-Shamir transcript for deterministic challenge sampling +/// +/// # Returns +/// `(proof, challenges)` where `proof` contains one round polynomial per variable +/// and `challenges` is the vector of random points sampled during the protocol. +pub fn sumcheck_prove( + evals: &[FieldElement], + claimed_sum: &FieldElement, + num_vars: usize, + max_degree: usize, + transcript: &mut impl IsTranscript, +) -> (SumcheckProof, Vec>) { + assert_eq!( + evals.len(), + 1 << num_vars, + "evals length must be 2^num_vars" + ); + assert!( + max_degree >= 1, + "max_degree must be at least 1 for a meaningful sumcheck" + ); + + let mut table = evals.to_vec(); + let mut round_polys = Vec::with_capacity(num_vars); + let mut challenges = Vec::with_capacity(num_vars); + let mut round_claimed_sum = claimed_sum.clone(); + + for _round in 0..num_vars { + let half = table.len() / 2; + + // p(0) = sum of left halves (zero multiplications) + let mut p0 = FieldElement::::zero(); + for j in 0..half { + p0 = &p0 + &table[2 * j]; + } + + // p(1) = round_claimed_sum - p(0) (free, from sumcheck identity p(0)+p(1)=sum) + let p1 = &round_claimed_sum - &p0; + + let mut poly_evals = Vec::with_capacity(max_degree + 1); + poly_evals.push(p0); + poly_evals.push(p1); + + // For t >= 2: use delta form. + // Linear interpolation at t: table[2j] + t * (table[2j+1] - table[2j]) + // p(t) = Σ_j [table[2j] + t * delta_j] where delta_j = table[2j+1] - table[2j] + for t in 2..=max_degree { + let t_fe = FieldElement::::from(t as u64); + let mut sum = FieldElement::::zero(); + for j in 0..half { + let delta = &table[2 * j + 1] - &table[2 * j]; + sum = &sum + &(&table[2 * j] + &(&t_fe * &delta)); + } + poly_evals.push(sum); + } + + let round_poly = RoundPoly::new(poly_evals); + + // Append all evaluations of the round polynomial to the transcript + for eval in round_poly.evals() { + transcript.append_field_element(eval); + } + + // Sample the challenge for this round + let challenge: FieldElement = transcript.sample_field_element(); + + // Update round_claimed_sum for next round: p(challenge) + round_claimed_sum = round_poly.evaluate(&challenge); + + // Bind the table: fold using the challenge + // table[j] = table[2j] + challenge * (table[2j+1] - table[2j]) + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + + round_polys.push(round_poly); + challenges.push(challenge); + } + + (SumcheckProof { round_polys }, challenges) +} + +/// Errors that can occur during sumcheck verification. +#[derive(Debug, Clone)] +pub enum SumcheckError { + /// The proof contains a different number of round polynomials than expected. + WrongNumberOfRounds { expected: usize, got: usize }, + /// The round polynomial's p(0) + p(1) does not match the expected sum. + RoundSumMismatch { round: usize }, +} + +impl fmt::Display for SumcheckError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SumcheckError::WrongNumberOfRounds { expected, got } => { + write!( + f, + "wrong number of rounds: expected {}, got {}", + expected, got + ) + } + SumcheckError::RoundSumMismatch { round } => { + write!(f, "round sum mismatch at round {}", round) + } + } + } +} + +/// Verify a sumcheck proof (non-interactive via Fiat-Shamir). +/// +/// Checks that the prover's round polynomials are consistent with the claimed +/// sum. The transcript operations must exactly mirror those of `sumcheck_prove` +/// so that the same challenges are derived. +/// +/// # Arguments +/// - `proof`: the sumcheck proof containing one round polynomial per variable +/// - `claimed_sum`: the asserted sum over the Boolean hypercube +/// - `num_vars`: number of Boolean variables +/// - `transcript`: Fiat-Shamir transcript (must have the same seed as the prover's) +/// +/// # Returns +/// `Ok((challenges, final_eval))` where `challenges` is the random point and +/// `final_eval` is the evaluation claim at that point (i.e., the last round +/// polynomial evaluated at the last challenge). +pub fn sumcheck_verify( + proof: &SumcheckProof, + claimed_sum: &FieldElement, + num_vars: usize, + transcript: &mut impl IsTranscript, +) -> Result<(Vec>, FieldElement), SumcheckError> { + if proof.round_polys.len() != num_vars { + return Err(SumcheckError::WrongNumberOfRounds { + expected: num_vars, + got: proof.round_polys.len(), + }); + } + + let mut current_sum = claimed_sum.clone(); + let mut challenges = Vec::with_capacity(num_vars); + + for round in 0..num_vars { + // Check that p_r(0) + p_r(1) == current_sum + if proof.round_polys[round].sum_at_binary() != current_sum { + return Err(SumcheckError::RoundSumMismatch { round }); + } + + // Append all evaluations of the round polynomial to the transcript + // (must match the prover exactly) + for eval in proof.round_polys[round].evals() { + transcript.append_field_element(eval); + } + + // Sample the challenge for this round + let challenge: FieldElement = transcript.sample_field_element(); + + // Update the running sum: next round's claimed sum is p_r(challenge) + current_sum = proof.round_polys[round].evaluate(&challenge); + + challenges.push(challenge); + } + + Ok((challenges, current_sum)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField; + + // Use the base Goldilocks field for tests. The cubic extension would work + // identically since all operations are generic over IsField. + type FE = FieldElement; + + /// Helper: evaluate a polynomial given in coefficient form at a point. + /// coeffs[i] is the coefficient of x^i. + fn eval_poly_coeffs(coeffs: &[FE], x: &FE) -> FE { + let mut result = FE::zero(); + let mut power = FE::one(); + for c in coeffs { + result = &result + &(c * &power); + power = &power * x; + } + result + } + + #[test] + fn test_sum_at_binary_linear() { + // p(x) = 3 + 5x => p(0) = 3, p(1) = 8 + // sum_at_binary = 3 + 8 = 11 + let evals = vec![FE::from(3u64), FE::from(8u64)]; + let poly = RoundPoly::new(evals); + assert_eq!(poly.sum_at_binary(), FE::from(11u64)); + } + + #[test] + fn test_sum_at_binary_quadratic() { + // p(x) = 2x^2 + 3x + 1 => p(0) = 1, p(1) = 6, p(2) = 15 + // sum_at_binary = p(0) + p(1) = 1 + 6 = 7 + let evals = vec![FE::from(1u64), FE::from(6u64), FE::from(15u64)]; + let poly = RoundPoly::new(evals); + assert_eq!(poly.sum_at_binary(), FE::from(7u64)); + } + + #[test] + fn test_evaluate_at_nodes_linear() { + // p(x) = 3 + 5x => p(0) = 3, p(1) = 8 + // Evaluating at the nodes should return the stored evaluations. + let evals = vec![FE::from(3u64), FE::from(8u64)]; + let poly = RoundPoly::new(evals); + + assert_eq!(poly.evaluate(&FE::from(0u64)), FE::from(3u64)); + assert_eq!(poly.evaluate(&FE::from(1u64)), FE::from(8u64)); + } + + #[test] + fn test_evaluate_at_nodes_quadratic() { + // p(x) = 2x^2 + 3x + 1 => p(0) = 1, p(1) = 6, p(2) = 15 + let evals = vec![FE::from(1u64), FE::from(6u64), FE::from(15u64)]; + let poly = RoundPoly::new(evals); + + assert_eq!(poly.evaluate(&FE::from(0u64)), FE::from(1u64)); + assert_eq!(poly.evaluate(&FE::from(1u64)), FE::from(6u64)); + assert_eq!(poly.evaluate(&FE::from(2u64)), FE::from(15u64)); + } + + #[test] + fn test_evaluate_at_random_point_linear() { + // p(x) = 3 + 5x + // p(0) = 3, p(1) = 8 + // p(7) = 3 + 35 = 38 + let evals = vec![FE::from(3u64), FE::from(8u64)]; + let poly = RoundPoly::new(evals); + + let point = FE::from(7u64); + let expected = FE::from(38u64); + assert_eq!(poly.evaluate(&point), expected); + } + + #[test] + fn test_evaluate_at_random_point_quadratic() { + // p(x) = 2x^2 + 3x + 1 + // p(0) = 1, p(1) = 6, p(2) = 15 + // p(5) = 2*25 + 3*5 + 1 = 50 + 15 + 1 = 66 + let coeffs = [FE::from(1u64), FE::from(3u64), FE::from(2u64)]; + let evals = vec![ + eval_poly_coeffs(&coeffs, &FE::from(0u64)), + eval_poly_coeffs(&coeffs, &FE::from(1u64)), + eval_poly_coeffs(&coeffs, &FE::from(2u64)), + ]; + let poly = RoundPoly::new(evals); + + let point = FE::from(5u64); + let expected = eval_poly_coeffs(&coeffs, &point); + assert_eq!(poly.evaluate(&point), expected); + assert_eq!(expected, FE::from(66u64)); + } + + #[test] + fn test_evaluate_at_random_point_cubic() { + // p(x) = x^3 + 2x^2 + 3x + 4 + // Need 4 evaluation points: 0, 1, 2, 3 + let coeffs = [ + FE::from(4u64), + FE::from(3u64), + FE::from(2u64), + FE::from(1u64), + ]; + let evals: Vec = (0..4) + .map(|i| eval_poly_coeffs(&coeffs, &FE::from(i as u64))) + .collect(); + let poly = RoundPoly::new(evals); + + // p(10) = 1000 + 200 + 30 + 4 = 1234 + let point = FE::from(10u64); + let expected = eval_poly_coeffs(&coeffs, &point); + assert_eq!(poly.evaluate(&point), expected); + assert_eq!(expected, FE::from(1234u64)); + } + + #[test] + fn test_evaluate_at_large_field_element() { + // Test with a large field element as the evaluation point to ensure + // field arithmetic works correctly (not just small integers). + let coeffs = [FE::from(7u64), FE::from(11u64), FE::from(13u64)]; + let evals: Vec = (0..3) + .map(|i| eval_poly_coeffs(&coeffs, &FE::from(i as u64))) + .collect(); + let poly = RoundPoly::new(evals); + + // Use a large evaluation point + let point = FE::from(1_000_000_007u64); + let expected = eval_poly_coeffs(&coeffs, &point); + assert_eq!(poly.evaluate(&point), expected); + } + + #[test] + fn test_constant_polynomial() { + // p(x) = 42 (constant) + let evals = vec![FE::from(42u64)]; + let poly = RoundPoly::new(evals); + + // A constant polynomial should evaluate to the same value everywhere + assert_eq!(poly.evaluate(&FE::from(0u64)), FE::from(42u64)); + assert_eq!(poly.evaluate(&FE::from(100u64)), FE::from(42u64)); + assert_eq!(poly.evaluate(&FE::from(999u64)), FE::from(42u64)); + } + + #[test] + fn test_num_evals() { + let evals = vec![FE::from(1u64), FE::from(2u64), FE::from(3u64)]; + let poly = RoundPoly::new(evals); + assert_eq!(poly.num_evals(), 3); + } + + #[test] + fn test_evals_accessor() { + let evals = vec![FE::from(10u64), FE::from(20u64)]; + let poly = RoundPoly::new(evals); + assert_eq!(poly.evals()[0], FE::from(10u64)); + assert_eq!(poly.evals()[1], FE::from(20u64)); + } + + #[test] + #[should_panic(expected = "RoundPoly must have at least one evaluation")] + fn test_empty_evals_panics() { + let _poly = RoundPoly::::new(vec![]); + } + + #[test] + #[should_panic(expected = "sum_at_binary requires at least 2 evaluations")] + fn test_sum_at_binary_single_eval_panics() { + let poly = RoundPoly::new(vec![FE::from(1u64)]); + let _ = poly.sum_at_binary(); + } + + #[test] + fn test_evaluate_consistency_with_coefficients() { + // Build a degree-4 polynomial from coefficients, generate evaluations + // at 0..=4, construct RoundPoly, and verify evaluate() at many points. + let coeffs = [ + FE::from(5u64), + FE::from(3u64), + FE::from(7u64), + FE::from(2u64), + FE::from(1u64), + ]; + let evals: Vec = (0..5) + .map(|i| eval_poly_coeffs(&coeffs, &FE::from(i as u64))) + .collect(); + let poly = RoundPoly::new(evals); + + // Check at several points including nodes and non-nodes + for x in [0u64, 1, 2, 3, 4, 5, 10, 100, 12345] { + let point = FE::from(x); + let expected = eval_poly_coeffs(&coeffs, &point); + assert_eq!(poly.evaluate(&point), expected, "Mismatch at x = {}", x); + } + } + + // ==================== Sumcheck prover tests ==================== + + #[test] + fn test_sumcheck_prove_2var_linear() { + // f(x1, x2) = 3*x1 + 5*x2 + 7 + // Hypercube evaluations in little-endian bit order: + // index 0 = (x1=0, x2=0): f = 7 + // index 1 = (x1=1, x2=0): f = 10 + // index 2 = (x1=0, x2=1): f = 12 + // index 3 = (x1=1, x2=1): f = 15 + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); // 7 + 10 + 12 + 15 + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, challenges) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut transcript); + + // Should have 2 round polynomials (one per variable) + assert_eq!(proof.round_polys.len(), 2); + assert_eq!(challenges.len(), 2); + + // First round poly: sum over x2 of f(x1, x2) + // p1(0) = f(0,0) + f(0,1) = 7 + 12 = 19 + // p1(1) = f(1,0) + f(1,1) = 10 + 15 = 25 + // p1(0) + p1(1) = 19 + 25 = 44 = claimed_sum + assert_eq!(proof.round_polys[0].sum_at_binary(), claimed_sum); + assert_eq!(proof.round_polys[0].evals()[0], FE::from(19u64)); + assert_eq!(proof.round_polys[0].evals()[1], FE::from(25u64)); + + // Each round poly should have max_degree+1 = 2 evaluations + assert_eq!(proof.round_polys[0].num_evals(), 2); + assert_eq!(proof.round_polys[1].num_evals(), 2); + + // Second round: the claimed sum for round 2 is p1(r1) where r1 = challenges[0]. + // Verify that p2(0) + p2(1) = p1(r1). + let r1_eval = proof.round_polys[0].evaluate(&challenges[0]); + assert_eq!(proof.round_polys[1].sum_at_binary(), r1_eval); + } + + #[test] + fn test_sumcheck_prove_1var() { + // f(x) = 2x + 3 => f(0) = 3, f(1) = 5 + // Sum = 3 + 5 = 8 + let evals = vec![FE::from(3u64), FE::from(5u64)]; + let claimed_sum = FE::from(8u64); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, challenges) = sumcheck_prove(&evals, &claimed_sum, 1, 1, &mut transcript); + + assert_eq!(proof.round_polys.len(), 1); + assert_eq!(challenges.len(), 1); + assert_eq!(proof.round_polys[0].sum_at_binary(), claimed_sum); + assert_eq!(proof.round_polys[0].evals()[0], FE::from(3u64)); + assert_eq!(proof.round_polys[0].evals()[1], FE::from(5u64)); + } + + #[test] + fn test_sumcheck_prove_3var_constant() { + // f(x1, x2, x3) = 1 for all inputs + // Sum = 8 (2^3 points, each contributing 1) + let evals = vec![FE::one(); 8]; + let claimed_sum = FE::from(8u64); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, challenges) = sumcheck_prove(&evals, &claimed_sum, 3, 1, &mut transcript); + + assert_eq!(proof.round_polys.len(), 3); + + // First round: p(0) = 4, p(1) = 4, sum = 8 + assert_eq!(proof.round_polys[0].sum_at_binary(), claimed_sum); + assert_eq!(proof.round_polys[0].evals()[0], FE::from(4u64)); + assert_eq!(proof.round_polys[0].evals()[1], FE::from(4u64)); + + // Each subsequent round's sum should match the previous round's evaluation at challenge + for r in 1..3 { + let prev_eval = proof.round_polys[r - 1].evaluate(&challenges[r - 1]); + assert_eq!(proof.round_polys[r].sum_at_binary(), prev_eval); + } + } + + #[test] + fn test_sumcheck_prove_higher_degree() { + // Test with max_degree = 3 (4 evaluation points per round poly). + // For a multilinear f, higher-degree evaluations are determined by + // the degree-1 interpolation, so they are just linear extrapolation. + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, challenges) = sumcheck_prove(&evals, &claimed_sum, 2, 3, &mut transcript); + + assert_eq!(proof.round_polys.len(), 2); + + // Each round poly should have max_degree+1 = 4 evaluations + assert_eq!(proof.round_polys[0].num_evals(), 4); + assert_eq!(proof.round_polys[1].num_evals(), 4); + + // p(0) + p(1) must still equal claimed sum + assert_eq!(proof.round_polys[0].sum_at_binary(), claimed_sum); + + // Since the underlying function is multilinear, the round poly is degree 1. + // Verify that evaluations at t=2 and t=3 are consistent with linear interpolation. + let p0 = &proof.round_polys[0].evals()[0]; + let p1 = &proof.round_polys[0].evals()[1]; + // p(t) = p0*(1-t) + p1*t = p0 + (p1 - p0)*t + let slope = p1 - p0; + let p2_expected = p0 + &(&slope * &FE::from(2u64)); + let p3_expected = p0 + &(&slope * &FE::from(3u64)); + assert_eq!(proof.round_polys[0].evals()[2], p2_expected); + assert_eq!(proof.round_polys[0].evals()[3], p3_expected); + + // Consistency between rounds + for r in 1..2 { + let prev_eval = proof.round_polys[r - 1].evaluate(&challenges[r - 1]); + assert_eq!(proof.round_polys[r].sum_at_binary(), prev_eval); + } + } + + #[test] + fn test_sumcheck_prove_deterministic() { + // Same inputs and transcript seed should produce identical proofs + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); + + let mut t1 = DefaultTranscript::::new(&[0x42]); + let (proof1, challenges1) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut t1); + + let mut t2 = DefaultTranscript::::new(&[0x42]); + let (proof2, challenges2) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut t2); + + assert_eq!(challenges1, challenges2); + for (rp1, rp2) in proof1.round_polys.iter().zip(proof2.round_polys.iter()) { + assert_eq!(rp1.evals(), rp2.evals()); + } + } + + #[test] + fn test_sumcheck_prove_final_value() { + // After all rounds, the table should contain a single element which is + // f evaluated at the random challenge point. We verify by checking that + // the last round poly evaluated at the last challenge gives this value. + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); + + let mut transcript = DefaultTranscript::::new(&[]); + let (proof, challenges) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut transcript); + + // Manually compute f(r1, r2) by multilinear evaluation + let r1 = &challenges[0]; + let r2 = &challenges[1]; + // f(r1, r2) = f(0,0)*(1-r1)*(1-r2) + f(1,0)*r1*(1-r2) + f(0,1)*(1-r1)*r2 + f(1,1)*r1*r2 + let one = FE::one(); + let one_minus_r1 = &one - r1; + let one_minus_r2 = &one - r2; + let f_at_r = &(&(&evals[0] * &one_minus_r1) * &one_minus_r2) + + &(&(&evals[1] * r1) * &one_minus_r2) + + &(&(&evals[2] * &one_minus_r1) * r2) + + &(&(&evals[3] * r1) * r2); + + // The last round poly evaluated at the last challenge should give f(r1, r2) + let last_round_eval = proof.round_polys[1].evaluate(&challenges[1]); + assert_eq!(last_round_eval, f_at_r); + } + + #[test] + #[should_panic(expected = "evals length must be 2^num_vars")] + fn test_sumcheck_prove_wrong_length_panics() { + let evals = vec![FE::from(1u64), FE::from(2u64), FE::from(3u64)]; // length 3, not a power of 2 + let mut transcript = DefaultTranscript::::new(&[]); + let _ = sumcheck_prove(&evals, &FE::from(6u64), 2, 1, &mut transcript); + } + + // ==================== Sumcheck verifier tests ==================== + + #[test] + fn test_sumcheck_verify() { + // f(x1, x2) = 3*x1 + 5*x2 + 7 + // Hypercube evaluations (little-endian bit order): + // (0,0)=7, (1,0)=10, (0,1)=12, (1,1)=15 + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); // 7 + 10 + 12 + 15 + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, prover_challenges) = + sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut prover_transcript); + + // Verify with a fresh transcript (same seed) + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = sumcheck_verify(&proof, &claimed_sum, 2, &mut verifier_transcript); + + assert!(result.is_ok(), "Verification should succeed"); + let (challenges, _final_eval) = result.unwrap(); + + // Verifier must derive the same challenges as the prover + assert_eq!(challenges, prover_challenges); + } + + #[test] + fn test_sumcheck_verify_wrong_sum() { + // f(x1, x2) = 3*x1 + 5*x2 + 7 + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); + + // Prove with the correct sum + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, _) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut prover_transcript); + + // Verify with a WRONG claimed sum + let wrong_sum = FE::from(99u64); + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let result = sumcheck_verify(&proof, &wrong_sum, 2, &mut verifier_transcript); + + assert!(result.is_err(), "Verification should fail with wrong sum"); + match result.unwrap_err() { + SumcheckError::RoundSumMismatch { round } => { + assert_eq!(round, 0, "Mismatch should be detected at round 0"); + } + other => panic!("Expected RoundSumMismatch, got {:?}", other), + } + } + + #[test] + fn test_sumcheck_verify_roundtrip_3var() { + // f(x1, x2, x3) = x1 * x2 + x3 + 2 + // Hypercube evaluations in little-endian bit order: + // index = b0 + 2*b1 + 4*b2, where b0=x1, b1=x2, b2=x3 + // (0,0,0)=2, (1,0,0)=2, (0,1,0)=2, (1,1,0)=3, (0,0,1)=3, (1,0,1)=3, (0,1,1)=3, (1,1,1)=4 + let evals = vec![ + FE::from(2u64), // f(0,0,0) = 0*0 + 0 + 2 = 2 + FE::from(2u64), // f(1,0,0) = 1*0 + 0 + 2 = 2 + FE::from(2u64), // f(0,1,0) = 0*1 + 0 + 2 = 2 + FE::from(3u64), // f(1,1,0) = 1*1 + 0 + 2 = 3 + FE::from(3u64), // f(0,0,1) = 0*0 + 1 + 2 = 3 + FE::from(3u64), // f(1,0,1) = 1*0 + 1 + 2 = 3 + FE::from(3u64), // f(0,1,1) = 0*1 + 1 + 2 = 3 + FE::from(4u64), // f(1,1,1) = 1*1 + 1 + 2 = 4 + ]; + let claimed_sum: FE = evals.iter().fold(FE::zero(), |acc, v| &acc + v); // = 22 + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[0xAB]); + let (proof, prover_challenges) = + sumcheck_prove(&evals, &claimed_sum, 3, 1, &mut prover_transcript); + + // Verify + let mut verifier_transcript = DefaultTranscript::::new(&[0xAB]); + let result = sumcheck_verify(&proof, &claimed_sum, 3, &mut verifier_transcript); + + assert!(result.is_ok(), "3-var verification should succeed"); + let (challenges, _final_eval) = result.unwrap(); + assert_eq!(challenges.len(), 3); + assert_eq!(challenges, prover_challenges); + } + + #[test] + fn test_sumcheck_verify_final_eval() { + // Verify that final_eval matches the MLE evaluated at the challenge point. + // f(x1, x2) = 3*x1 + 5*x2 + 7 + let evals = vec![ + FE::from(7u64), + FE::from(10u64), + FE::from(12u64), + FE::from(15u64), + ]; + let claimed_sum = FE::from(44u64); + + // Prove + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, _) = sumcheck_prove(&evals, &claimed_sum, 2, 1, &mut prover_transcript); + + // Verify + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let (challenges, final_eval) = + sumcheck_verify(&proof, &claimed_sum, 2, &mut verifier_transcript).unwrap(); + + // Compute MLE at the challenge point manually: + // f(r1, r2) = f(0,0)*(1-r1)*(1-r2) + f(1,0)*r1*(1-r2) + // + f(0,1)*(1-r1)*r2 + f(1,1)*r1*r2 + let r1 = &challenges[0]; + let r2 = &challenges[1]; + let one = FE::one(); + let one_minus_r1 = &one - r1; + let one_minus_r2 = &one - r2; + let mle_at_r = &(&(&evals[0] * &one_minus_r1) * &one_minus_r2) + + &(&(&evals[1] * r1) * &one_minus_r2) + + &(&(&evals[2] * &one_minus_r1) * r2) + + &(&(&evals[3] * r1) * r2); + + assert_eq!( + final_eval, mle_at_r, + "final_eval from verifier must match MLE evaluated at challenge point" + ); + } +} From 08463caa4720bbba7e8767c4cb3baa0f14128609 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:23:26 -0300 Subject: [PATCH 02/20] docs(gkr): LogUp-GKR port plan and analysis of PR #485 vs current main --- thoughts/logup-gkr/branch-analysis.md | 100 ++++++++++++++ thoughts/logup-gkr/drift-analysis.md | 68 ++++++++++ thoughts/logup-gkr/main-surface-map.md | 76 +++++++++++ thoughts/logup-gkr/port-plan.md | 176 +++++++++++++++++++++++++ 4 files changed, 420 insertions(+) create mode 100644 thoughts/logup-gkr/branch-analysis.md create mode 100644 thoughts/logup-gkr/drift-analysis.md create mode 100644 thoughts/logup-gkr/main-surface-map.md create mode 100644 thoughts/logup-gkr/port-plan.md diff --git a/thoughts/logup-gkr/branch-analysis.md b/thoughts/logup-gkr/branch-analysis.md new file mode 100644 index 000000000..875e0fa5a --- /dev/null +++ b/thoughts/logup-gkr/branch-analysis.md @@ -0,0 +1,100 @@ +# LogUp-GKR branch (`origin/feat/gkr-logup`, PR #485) — implementation analysis + +> Companion to `port-plan.md`. Produced 2026-07-21 by deep-reading the branch tip (`78be304f`). +> "SNAP" refers to full-file snapshots of the branch, "DIFF" to per-file diffs vs the merge-base `7b42e51b`. +> Regenerate them locally with: +> `git fetch origin feat/gkr-logup` +> `git show origin/feat/gkr-logup:crypto/stark/src/gkr.rs` (etc.) +> `git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/lookup.rs` (etc.) +> Line references below are into those snapshots/diffs, NOT into current main. + +## 1. Protocol architecture + +**Core idea.** Standard LogUp commits per-table aux term columns + an accumulated column and checks Σ contributions via `bus_public_inputs.table_contribution`. The GKR branch instead proves each table's total LogUp sum with a **batch GKR over fractional summation trees**, and replaces the committed aux machinery with exactly **2 aux columns** per interacting table: the **Lagrange kernel `l`** (aux col 0) and a **bridge running sum `σ`** (aux col 1). + +Pipeline per table: +1. **Leaf fractions** (`compute_logup_leaf_fractions`, DIFF lookup.rs:837-905): for each row i, all K interactions are folded into one fraction N(i)/D(i) by cross-multiplication (`n' = n·fp_k + sign_k·m_k·d`, `d' = d·fp_k`), where `fp_k = z − (bus_id·α⁰ + Σ vⱼ·αʲ)` and `m_k` from the `Multiplicity` variant. Row-parallel, no batch inversion. +2. **Summation tree** (`gen_layers`/`next_logup_layer`, SNAP gkr.rs:91-135): pairwise fraction addition halves each layer up to a 1-element root; root value n/d = the table's total contribution. `Layer::{LogUpGeneric, LogUpSingles}` (Singles = implicit numerators 1, ~50% fewer muls; production leaves are Generic). +3. **Batch GKR** (`gkr_prove_batch`, gkr.rs:1221-1851): ONE proof for all tables. Per layer (root→leaves): append each active instance's (n,d) claims; sample `sumcheck_alpha` (combines instances) and `lambda` (combines n/d claims: `claim = n + λ·d`); run one shared sumcheck over `max_parent_vars` variables with degree-3 round polys (4 evals each, `RoundPoly` in sumcheck.rs); gate = `nl·dr + nr·dl + λ·dl·dr` weighted by `eq(current_point, ·)`. Mixed-size instances: smaller tables join at later layers, with a `2^n_unused` doubling factor on their combined claim (gkr.rs:1309-1316) and constant `claim/2` round polys while inactive (1534-1550). After rounds, append 4 child claims per instance, sample `eta`, fold; `current_point = [eta] ++ challenges`. Eq handling uses Dao-Thaler halving + a scalar `eq_correction`, and **SVO** (split-value optimization, ePrint 2025/1117 Alg 5, `SVO_THRESHOLD=8`): eq split into prefix×suffix tables for √-memory (gkr.rs:522-534, batch version 1471-1494, 1614-1659). +4. **Output**: shared `random_point` + per-instance `(n_claim, d_claim)` = claimed leaf MLE evals. Per-instance point = `instance_eval_point(shared, n_vars)` = `[eta] ++ last (n_vars−1) challenges` (gkr.rs:1861-1876). +5. **Bridge back to the STARK** (the committed-trace link): + - `column_claims`: MLE of each distinct main column referenced by any interaction (sorted indices via `extract_column_indices`, DIFF lookup.rs:352-399), evaluated at the instance point via the Lagrange kernel (`finalize_logup_gkr_result`, lookup.rs diff 1336-1376). + - **Aux col 0 = Lagrange kernel** `l[i] = eq(bits(i), r)` (SNAP lagrange_kernel.rs:21-67). Bound by boundary constraint `l[0] = ∏(1−r_j)` (DIFF lookup.rs:312-330) and by an **l² self-check**: the bridge's batched value includes `γ^K·l[i]`, forcing `Σ l[i]² = ∏(r_k² + (1−r_k)²)` (BUG-004 fix; `compute_bridge_params`, lookup.rs diff 437-458). + - **Aux col 1 = σ running sum** with circular transition constraint `LookupBridgeSumConstraint` (lookup.rs diff 512-629, degree 2, `end_exemptions=0`): `σ_next − σ_curr − l_curr·batched_curr + Δ = 0`, where `batched = Σⱼ γʲ·colⱼ + γ^K·l` and `Δ = target/N`, `target = Σⱼ γʲ·cⱼ + γ^K·l_mle_claim`. Telescoping over N rows proves `⟨l, colⱼ⟩ = cⱼ` for all j (Schwartz-Zippel over γ). +6. **Bus balance**: verifier sums `root_n/root_d` over `batch_gkr_proof.root_claims` and compares to `expected_bus_balance` (DIFF verifier.rs:386-421). `bus_public_inputs` is gone for GKR tables (`build_auxiliary_trace` returns `None`, only allocates). + +**What replaces the acc columns:** aux goes from `⌈K/2⌉ term columns + 1 acc` to a **fixed 2 columns** regardless of K (`AirWithBuses::new`, DIFF lookup.rs:121-127). This is always-on for any AIR with interactions — no mode flag anywhere on the branch. + +**Memory motivation** (PR bench: peak heap −70%, 221→66 GB; prove +2.6% on fib_iterative_8M at commit `334093fc`): the committed aux trace collapses from ~⌈K/2⌉+1 extension columns at LDE size to 2 (CPU K=40 → 21→2, DVRM 34 → 18→2, MEMW 26 → 14→2). That shrinks aux LDE buffers, aux Merkle trees, OOD tables, and DEEP openings ~10× for the big tables. GKR's own layer trees (~4N ext elements/table) are transient and freed before the LDE phase. (? INFERRED from code structure — consistent with the aux-column arithmetic; the in-branch design doc doesn't cover GKR.) + +## 2. Prover integration map (`multi_prove`, DIFF prover.rs:157-510) + +Round-1 structure becomes: **Phase A** (unchanged: per-table main commits appended to main transcript, SNAP prover.rs:1583-1639) → **Phase B** (unchanged: sample z, α; `LOGUP_NUM_CHALLENGES=2`, prover.rs:1648-1654) → **NEW Phase B′**: collect `gkr_table_indices` where `air.has_trace_interaction()`; compute layer trees in parallel (`compute_logup_layers` from `air.bus_interactions()` + `trace.columns_main()`); run `gkr_prove_batch` on the **main** transcript; distribute per-table `LogUpGkrResult` in parallel (`finalize_logup_gkr_result`; `table_contrib = root_n·root_d⁻¹`) → **NEW Phase B″**: append every table's `column_claims` values to main transcript, then sample **γ** (BUG-012) → **Phase C pass 1 changed**: for GKR tables, the aux trace is built **inline in prover.rs** (not in `build_auxiliary_trace`): write kernel to aux col 0; compute `l_mle_claim`, `compute_bridge_params`, precompute `batched[i]` (row-parallel), then sequentially fill `σ` forward with `σ[0]=0` (prover diff 326-401). Non-GKR aux path retained for AIRs without interactions. → **metadata build**: per-table `rap_challenges = [z, α] ++ bridge params` via `extend_rap_challenges_with_bridge` (layout: `[2]=γ, [3]=Δ, [4..4+K+1]=γ⁰..γ^K, then random_point`; lookup.rs diff 460-495); `bus_public_inputs: None`; new `Round1Metadata.logup_gkr_result` field → **Phase C pass 2 / Rounds 2-4 unchanged** (fork per table, aux root appended in fork), except each proof gets `logup_gkr_proof = {gkr_proof: stub-with-claimed_sum, random_point, column_claims}` attached post-hoc (prover diff 483-496); `MultiProof.batch_gkr_proof` carries the real batch proof. Single-table `prove` copies `batch_gkr_proof` into the `StarkProof` (BUG-014, prover diff 513-528). New bound `Field: IsPrimeField` on `multi_prove`/`prove`. + +## 3. Verifier integration map (`multi_verify`, DIFF verifier.rs:171-429) + +The old monolithic `step_1_replay_rounds_and_recover_challenges` is **deleted** (verifier diff 21-166; rounds 2-4 replay lives in `verify_rounds_2_to_4`). After Phase B challenge sampling: **Phase B′**: require `multi_proof.batch_gkr_proof` if any table has interactions; `gkr_verify_batch(batch_proof, n_layers_by_instance, transcript)` where `n_vars = trace_length.trailing_zeros()` per table; on success, per table: require `proof.logup_gkr_proof`; check `column_claims.len() == extract_column_indices(air.bus_interactions()).len()`; run `reconstruct_and_verify_gkr_claims(n_claim, d_claim, column_claims, interactions, challenges, n_layers)` (see §7); store `gkr_bridge_claims[idx]` and transcript-derived `gkr_random_points[idx] = instance_eval_point(shared_random_point, n_vars)`. **Phase B″**: append column_claims in the same order, sample γ. **Per-table**: fork transcript, append aux root, build `table_rap_challenges` via the same `extend_rap_challenges_with_bridge`, then `verify_rounds_2_to_4` — the bridge is checked as an ordinary transition constraint at the OOD point, the kernel by the boundary constraint. **Bus balance**: Σ over `batch_proof.root_claims` of `root_n·root_d⁻¹` (zero denominator → `return false`) `== expected_bus_balance`. `verify` wraps into a `MultiProof` carrying `proof.batch_gkr_proof`. + +`gkr_verify_batch` (gkr.rs:1911-2170) replays exactly: activation → append active claims → sample α, λ → per-round `sum_at_binary` check + append evals + sample challenge → per-instance gate check with `eq(instance_point, challenges[n_unused..])` → append child claims → sample eta → fold claims. + +## 4. Fiat-Shamir ordering (main transcript, branch-tip behavior) + +1. Per table in index order: [preprocessed root] main root (`append_bytes`). +2. If any interactions: sample z, α (2 × `sample_field_element`). +3. `append_bytes(b"gkr_batch")`; `append_bytes(n_instances as u64 LE)`. Then per batch layer: for each instance active this layer (ascending index): append n, d; sample `sumcheck_alpha`; sample `lambda`; if non-trivial layer, per round: append 4 round-poly evals, sample challenge; then per active instance: append 4 child claims; sample `eta`. +4. Per GKR table (ascending index): append each `column_claims` value (sorted column order); then sample γ (ONE shared γ on the main transcript). +5. Per-table fork: clone; if multi-table append idx domain separator (pre-existing); append aux root; rounds 2-4 unchanged (β, OOD z, trace OOD evals, composition OOD, deep-γ, FRI ζ/roots, last value, nonce, iotas). + +Note the (pre-existing) fork-time `append_field_element(bpi.table_contribution)` is dead for GKR tables since `bus_public_inputs` is `None`. + +## 5. Proof format changes (DIFF proof_stark.rs) + +- New `LogUpGkrProof { gkr_proof: GkrProof, random_point: Vec, column_claims: Vec<(usize, E)> }`. +- `StarkProof` += `logup_gkr_proof: Option>` and `#[serde(default)] batch_gkr_proof: Option>`. +- `MultiProof` += `#[serde(default)] batch_gkr_proof: Option>`. +- `BatchGkrProof { root_claims: Vec<(E,E)>, layer_proofs: Vec }`; `BatchGkrLayerProof { sumcheck_proof: SumcheckProof, child_claims_by_instance: Vec<[E;4]> }`. All serde-derived (branch predates rkyv-authoritative wire format). +- **Dead weight**: in batch mode `LogUpGkrProof.gkr_proof` is a stub (empty `layer_proofs`) and `random_point` is written by the prover but **never read by the verifier** (✓ VERIFIED by grep: verifier only reads `column_claims`). The port should drop both fields. +- Proof size: replaces per-table aux OOD/opening data with the batch GKR proof (≈ Σ_layers rounds × 4 evals + 4 claims/instance/layer, extension elements). No in-branch size numbers. + +## 6. API/trait changes needed from the (now-rewritten) lookup layer + +- `AIR::bus_interactions(&self) -> &[BusInteraction]` (new trait method, default `&[]`; DIFF traits.rs:18-24) — the only trait addition. +- The GKR path consumes: `BusInteraction{bus_id, is_sender, values, multiplicity}`, all `Multiplicity` variants, `BusValue::{Packed, Linear}` + `Packing::combine`/`num_columns` + `accumulate_fingerprint` + `column_indices()`, `LinearTerm`, `num_bus_elements()`, `PackingShifts`, `compute_alpha_powers`. It reads **raw main columns** (`trace.columns_main()`) — it never touches the constraint system. +- New lookup.rs exports: `LOGUP_CHALLENGE_Z/GAMMA`, `LOGUP_BRIDGE_OFFSET_IDX`, `LOGUP_GAMMA_POWERS_START`, `logup_random_point_start`, `extract_column_indices`, `compute_bridge_params`, `extend_rap_challenges_with_bridge`, `compute_logup_leaf_fractions`, `compute_logup_layers`, `finalize_logup_gkr_result`, `reconstruct_and_verify_gkr_claims`, `LogUpGkrResult`, `LookupBridgeSumConstraint` (a boxed `TransitionConstraint` — on current main this must become an emit-style constraint body). Old machinery (`split_interactions`, term-column builders, `LookupBatchedTermConstraint`, `LookupAccumulatedConstraint`, debug bus sums) is `#[allow(dead_code)]`-retired on the branch, not deleted. +- `AirWithBuses::new`: aux layout = 2 columns; appends the bridge constraint; `transition_offsets` stays `[0,1]`. + +## 7. Soundness history + review-item status at branch tip + +Branch bug-fix commits: `d04f642a` (BUG-004 Lagrange-kernel binding via γ^K·l² self-check; BUG-011 0-layer forgery → `n_layers == 0` direct rational check; BUG-012 column_claims FS gap → Phase B″ binding before γ), `54dfa975` + `95deda49` (verifier DoS/panic bounds), `dfd4c53b` (zero root denominator → reject not panic), `e7761ed9` (trivial-layer gate check in batch verifier + BUG-014 single-proof API), `72ad8cf7` (Result-returning cleanups + tests). + +**The Codex review's three items, verified at branch tip:** + +1. **Payload `random_point` trust — FIXED.** ✓ VERIFIED: the verifier uses only the transcript-derived point: `gkr_verify_batch(...) → shared_random_point` (SNAP verifier.rs:828), `instance_eval_point(&shared_random_point, n_vars)` → `gkr_random_points[table_idx]` (889-890), fed to `extend_rap_challenges_with_bridge` (956). `proof.logup_gkr_proof.random_point` is never read anywhere in verifier.rs. Covered by the batch-verify rework + `d04f642a`. + +2. **`reconstruct_and_verify_gkr_claims` fail-open — STILL FAIL-OPEN for multi-interaction tables; a real open soundness gap.** ✓ VERIFIED at tip (DIFF lookup.rs:1108-1118): `if interactions.len() == 1 || n_layers == 0 { rational cross-check n·d̂ == n̂·d } else { true }` — the `n_layers == 0` arm is BUG-011's fix, but K>1 tables with N>1 rows still pass after structural checks only. The in-code justification ("bridge ensures column_claims consistency") is true but insufficient: the bridge binds **column_claims ↔ trace** (⟨l,colⱼ⟩=cⱼ), while **nothing binds the GKR leaf claims (n_claim,d_claim) to the columns** — the leaf fraction is a nonlinear (product) combination, MLE doesn't commute with products, and no other check touches `per_instance_claims` for these tables. Consequence: a prover can run an honest GKR over **fabricated leaf vectors** (arbitrary root contribution, e.g. to fake bus balance) while supplying honest column_claims so the bridge, kernel checks, and FS all pass. Since every production table is multi-interaction, this needs a protocol fix before production use — e.g. per-interaction leaves (K separate GKR instances, or a wider input layer where numerator/denominator are **linear** in columns, as in Winterfell/Stwo LogUp-GKR) or an extra input-layer sumcheck reducing the leaf claim to column MLE claims. + +3. **`gkr_verify_batch` malformed-proof panics — LARGELY FIXED, one residual vector.** ✓ VERIFIED: `layer_proofs.len() == max_layers` (SNAP gkr.rs:1939-1947), `child_claims_by_instance.len() >= active_instances.len()` (1999-2008), `num_rounds >= parent_num_vars_i` (2099-2106), zero root denominator → `return false` (DIFF verifier.rs:397-411). **Residual**: `RoundPoly` derives `Deserialize` with no length invariant; `sum_at_binary` **asserts** `evals.len() >= 2` (SNAP sumcheck.rs:38-44) and `evaluate` computes `evals.len()-1` (57) — a proof containing an empty/1-eval round poly still panics the verifier (called at gkr.rs:2072 before any length validation). Minor: `1u64 << n_unused` (gkr.rs:1989) can overflow-panic in debug if instance size spans ≥64 (needs absurd `trace_length`; `trace_length=0` gives `trailing_zeros()=64` unless validated upstream). The port must add a `num_evals` check (or make `sum_at_binary`/`evaluate` fallible). + +## 8. Tests + +- **Portable as-is** (self-contained modules): gkr.rs `mod tests` (fraction add, tree building, single + batch prove/verify roundtrips, tamper rejection), sumcheck.rs tests, lagrange_kernel.rs tests (kernel values, partition of unity, MLE evals). +- **lookup.rs tests** (DIFF 1378-1623): leaf-fraction unit tests (single sender / receiver-with-multiplicity / two-interaction cross-multiply / consistency vs `compute_logup_term_column`) — portable modulo the old term-column reference (the consistency test needs a local reimplementation or deletion). +- **bus_tests/soundness_tests.rs** (DIFF, −4/+5 tests): DELETED obsolete `test_tampered_table_contribution`, `test_missing_bus_public_inputs_rejected`, `test_injected_bus_public_inputs_on_non_logup_air_rejected`, `test_zeroed_table_contribution_rejected`; ADDED `generate_valid_multi_proof` helper + `test_tampered_gkr_column_claims_rejected`, `test_tampered_gkr_claimed_sum_rejected` (tampers `batch_proof.root_claims`), `test_missing_gkr_proof_rejected`, `test_tampered_sigma_ood_rejected`, `test_tampered_lagrange_kernel_random_point_rejected` (tampers `child_claims`); kept/adapted `test_tampered_acc_ood_evaluation`. All portable in spirit; they use `AirWithBuses` test AIRs which exist (rewritten) on main. +- **completeness_tests.rs**: +`test_single_table_prove_verify_with_gkr` (BUG-014 regression). packing_tests.rs: 8-line layout-constant tweak. +- No test covers the §7-item-2 gap (a fabricated-leaves forgery test would FAIL-to-reject on this branch — worth adding as an expected-fail/#[ignore] marker test during the port). + +## 9. Perf-commit contamination (exclude from the GKR port) + +Unrelated hunks (commits `bbe4dbc4`, `d5666021`, `c3f5719b`, `3aa03e6c`, `ccb75483`): +- `crypto/stark/src/fri/fri_functions.rs` — entire diff (parallel `fold_evaluations_in_place`), commit `3aa03e6c`. NOT on main; separable follow-up PR if still wanted. +- `crypto/stark/src/prover.rs` — `commit_columns_bit_reversed` `map_init` row-buffer hunk (DIFF prover.rs:35-79), commit `ccb74483` — function deleted on main (#735); `Arc` dedup hunks (82-105, 117-147, 152-153, 410-411, 465-466), commit `c3f5719b` — superseded on main. +- `prover/src/tables/mod.rs` — entire diff (max_rows caps), commit `d5666021` — already on main as #499. +- `crypto/stark/src/constraints/evaluator.rs` — entire diff (timing instrumentation only). +- `crypto/stark/src/debug.rs` — 2 comment lines. +- `docs/superpowers/specs/2026-04-09-scaling-improvements-design.md` — commit `bbe4dbc4` (see §10). + +GKR-core also includes its own perf commits (`7ec0b566`/`fb65d0f2` SVO, `dd6a2147` O(1) claim updates, `36150b81` parallel fold_table) — these live inside gkr.rs/sumcheck.rs and port with the files. + +## 10. The in-branch design doc + +`docs/superpowers/specs/2026-04-09-scaling-improvements-design.md` is **not** about GKR — it's the spec for the unrelated Item-9 perf commits (unified 2^20 sizing, twiddle dedup, FRI-fold parallelization, commit-alloc fix) plus future MMCS batched commitments and shared FRI; it explicitly lists "GKR-based LogUp" as a **non-goal**. GKR's actual motivation is the memory profile (peak heap −70% at +2.6% prove time, per PR bench). The disk-spill follow-up (PR #489) is not described in any in-branch doc. diff --git a/thoughts/logup-gkr/drift-analysis.md b/thoughts/logup-gkr/drift-analysis.md new file mode 100644 index 000000000..a3cd51df7 --- /dev/null +++ b/thoughts/logup-gkr/drift-analysis.md @@ -0,0 +1,68 @@ +# Drift analysis: `feat/gkr-logup` (PR #485) vs `origin/main` + +> Companion to `port-plan.md`. Produced 2026-07-21. Merge-base `7b42e51b` (2026-04-10); main (`a8648320`) has ~148 commits since, including full rewrites of every GKR integration point. + +## 1. Per-commit classification (22 commits on the branch) + +| Commit | Subject | Class | Verdict | +|---|---|---|---| +| `292fce80` | port logup gkr (the v2 port: gkr.rs +2663, sumcheck.rs +828, lagrange_kernel.rs +316, lookup/prover/verifier wiring) | GKR-core | **port** (via branch-tip state, not cherry-pick) | +| `bbe4dbc4` | docs: scaling improvements spec | docs | optional (unrelated to GKR — see branch-analysis §10) | +| `d5666021` | cap LT/LOAD/BRANCH/MEMW_R max_rows at 2^20 | perf | **already-in-main** — landed as PR #499 (`902e1027`); main's `prover/src/tables/mod.rs` has a `max_rows` module with all four at `1 << 20` | +| `c3f5719b` | dedup LdeTwiddles by domain size in multi_prove | perf | **already-in-main (superseded)** — main has a full shared-cache design (`LdeTwiddles` + `OnceLock` composition twiddles, `twiddle_caches: &[Arc]` threaded through multi_prove) | +| `3aa03e6c` | rayon-parallelize `fold_evaluations_in_place` | perf | **still-absent on main** — main's fold is sequential (verified). Independent of GKR; resubmit separately if still wanted (signature changed via #591/#598, needs re-port) | +| `ccb75483` | eliminate per-row Vec alloc in `commit_columns_bit_reversed` | perf | **superseded-by-rewrite** — function gone; #735 replaced it with `commitment.rs::commit_bit_reversed` (`rows_per_leaf`) | +| `b2821bf8` | Merge feat/logup_gkr_v2 | merge | n/a | +| `36150b81` | parallelize GKR `fold_table` inner loop | GKR-core (perf) | port (already in branch-tip gkr.rs) | +| `9b038bfb` | wire batch GKR into STARK prover+verifier | GKR-core (wiring) | **re-implement by hand** — targets pre-#764/#823 prover/verifier | +| `7ec0b566` | SVO for eq polynomial in sumcheck | GKR-core (perf) | port | +| `dd6a2147` | eliminate O(N) combined_claims recompute | GKR-core (perf) | port | +| `a9aa266d` | save work (prover.rs wiring tweaks) | GKR-core (wiring) | re-implement | +| `fb65d0f2` | SVO port to batch GKR inner loop | GKR-core (perf) | port | +| `95deda49` | "Fix three verifier panics in batch GKR" | GKR-core | port (note: diffstat is only +1 line in fri_functions.rs — message/content mismatch; real fixes appear in later commits) | +| `54dfa975` | verifier DoS guards (layer_proofs len, child_claims bounds) | GKR-core (soundness) | port — in branch-tip gkr.rs | +| `d04f642a` | Lagrange kernel soundness, 0-layer bus forgery, column_claims FS gap | GKR-core (soundness) | port — spans gkr/lookup/prover/verifier; the lookup/prover/verifier parts need re-implementation | +| `334093fc` / `6ee8c779` | merges of main (spec v0.2 etc.) | merge | n/a | +| `dfd4c53b` | return false on zero GKR root denominator | GKR-core (soundness) | port | +| `e7761ed9` | gate check for trivial layers + single-proof API fix | GKR-core | port | +| `72ad8cf7` | dead-pub cleanup, Result instead of panic, tests | GKR-core | port | +| `78be304f` | fix fmt | chore | n/a | + +Key point: **all GKR algorithm work (incl. all soundness fixes) is contained in the branch-tip versions of the three new files**; the wiring commits are the only part that must be redone. + +## 2. Sibling branches + +- **`feat/logup_gkr_v2`**: fully contained in `feat/gkr-logup` (`git log gkr-logup..logup_gkr_v2` is empty). Ignore. +- **`feat/gkr-logup_opt`**: stale precursor. Unique commits: `a64da708` "add instruments" (+31 lines of timing instrumentation in prover.rs) and a 4-line save-work. It **lacks** the later soundness fixes. Nothing worth salvaging; ignore. +- **PR #489 disk-spill** (branch deleted; fetchable as `refs/pull/489/head`, tip `f2b2bbf2`): stacked on gkr-logup, adds mmap disk-spill of trace tables + Merkle trees (`crypto/stark/src/{table,trace}.rs` +555, `prover/src/tables/trace_builder.rs` +199, `prover/src/tests/disk_spill_tests.rs`, memmap deps). Orthogonal scaling feature, also missing the final soundness fixes. Treat as separate future work, not part of this port. + +## 3. Main-side rewrites of touched files (all still exist on main; none renamed/deleted) + +| File | Main commits since base | Dominant reshapers | +|---|---|---| +| `lookup.rs` | 10 | **#764 single-source constraints** (rewrote LogUp emission wholesale), **#823 forward accumulation** (acc = sole next-row OOD read), #762 GPU aux-trace build, #696, #769. Branch is 3099 lines vs main 2755; **zero** GKR identifiers survive on main | +| `prover.rs` | 28 | **#799 full GPU trace residency**, #748 GPU-resident LDE/Merkle, #735 commitment-layer unification, #715 row-major LDE, #729 FRI early termination, #762, #823. Branch 2457 vs main 3325 lines | +| `verifier.rs` | 14 | #823, #826 deep-composition fusion, #815 OOD guards, **#769 rkyv in-place verify** (verifier reads archived proofs), #764, #729, #735 | +| `traits.rs` | 4 | #764 (AIR trait: `compute_transition_prover`/`constraint_program()`), #823 | +| `proof/stark.rs` | 4 | #769 (rkyv derive on proof types!), #729, #735, #823 — any new GKR proof fields must be rkyv-compatible | +| `constraints/evaluator.rs` | 6 | #764, **#798 GPU constraint/composition eval**, #799 | +| `fri/fri_functions.rs` | 2 | #591, #598 (signature changes only; branch's rayon fold not present) | +| `debug.rs` | 3 | #764 etc. (branch delta trivial: +2 lines) | +| `lib.rs` | 9 | module-registration churn (trivial to redo: add `pub mod gkr; pub mod sumcheck; pub mod lagrange_kernel;`) | +| `prover/src/tables/mod.rs` | 7 | #499 already covers the branch's change; #685 continuations, #657/#753 ECSM, #644 | +| `tests/bus_tests/*` | many | all three touched files exist; heavily churned by #764/#823/#688 — branch test edits must be re-authored against the new harness | + +## 4. Mechanical conflict probe (`git merge-tree --write-tree origin/main origin/feat/gkr-logup`) + +**Conflicted (9):** `constraints/evaluator.rs`, `lib.rs`, `lookup.rs`, `proof/stark.rs`, `prover.rs`, `bus_tests/soundness_tests.rs`, `traits.rs`, `verifier.rs`, `prover/src/tables/mod.rs` — i.e., every wiring file. +**Clean:** `gkr.rs`, `sumcheck.rs`, `lagrange_kernel.rs`, the docs spec (pure adds), plus `debug.rs`, `fri_functions.rs`, `completeness_tests.rs`, `packing_tests.rs` — though "auto-merged" here still means semantically wrong (they auto-merge against code #764 deleted). + +## 5. Verdict + +**Rebase/cherry-pick is not viable; port-plus-rewire is.** Concretely: + +1. **Port wholesale (clean adds, branch-tip versions):** `crypto/stark/src/gkr.rs` (3188 lines), `sumcheck.rs` (830), `lagrange_kernel.rs` (311). These carry every algorithmic commit including all five soundness-fix commits. Expect only compile-level fixups (field-trait/API drift from #598 etc.). +2. **Re-implement by hand on main:** all wiring — `lookup.rs` (against #764's single-source LogUp emission + #823's forward accumulation), `prover.rs` (against #748/#799's GPU-resident round structure), `verifier.rs` (against #769's rkyv archived-proof access + #826), `proof/stark.rs` (new GKR fields need rkyv derives), `traits.rs`, `lib.rs` module decls, and re-authored bus_tests. Use the branch's diffs (`git diff 7b42e51b..origin/feat/gkr-logup -- `) as the semantic spec of *what* to wire, not *how*. +3. **Drop:** `d5666021` (in main as #499), `c3f5719b` (superseded by shared twiddle caches), `ccb75483` (function deleted by #735), everything from `_opt` and `logup_gkr_v2`. +4. **Optional separate PR:** re-port `3aa03e6c` (rayon FRI fold) — genuinely absent on main's CPU path. +5. **Defer:** PR #489 disk-spill (orthogonal, stacked, stale). diff --git a/thoughts/logup-gkr/main-surface-map.md b/thoughts/logup-gkr/main-surface-map.md new file mode 100644 index 000000000..78e2ac3f0 --- /dev/null +++ b/thoughts/logup-gkr/main-surface-map.md @@ -0,0 +1,76 @@ +# Current main — LogUp-GKR integration surfaces + +> Companion to `port-plan.md`. Produced 2026-07-21 against main at `a8648320` (+1 trivial math-test commit). All file:line references verified by reading the working tree at that commit; re-verify anchors if main has moved. + +## 1. `crypto/stark/src/lookup.rs` (2755 lines) — the LogUp core + +**`AirWithBuses`** (`lookup.rs:813`) is the ONE production AIR type. Fields: `constraint_set: CS` (table's base constraints), `logup: LogUpLayout`, derived `meta: Vec` (base prefix + LogUp ext appended), `num_base`, lazy `constraint_program: OnceLock`, `auxiliary_trace_build_data` (the `Vec`), `preprocessed_commitment`/`num_precomputed_cols`, `max_bus_elements`. Constructor `new()` at `lookup.rs:870`: derives LogUp meta by running `emit_logup_constraints` through a `MetaBuilder` (`lookup.rs:890-892`). + +**Aux column layout** (from `split_interactions` `lookup.rs:117`): N interactions → `⌈N/2⌉` aux (extension-field) columns = `num_committed_pairs` batched **term columns** (each holds `Σ ±m/fp` for a PAIR of interactions) + **1 accumulated column** (last; index `logup.acc_column_idx = num_term_columns`). The last 1–2 interactions are "absorbed" into the accumulated constraint (no committed column). **Multiplicities and bus values live in MAIN trace columns** — aux columns are derived. `LogUpLayout` struct at `lookup.rs:1872` (`interactions`, `num_committed_pairs`, `num_term_columns`, `acc_column_idx`; `absorbed()` at 1902, `num_constraints()` at 1914). + +**Inputs a GKR prover needs — exactly these functions:** +- `BusInteraction` (`lookup.rs:1428`): `{bus_id: u64, multiplicity: Multiplicity, values: Vec, is_sender: bool}`. +- Per-row fingerprint = `z − (bus_id + Σ αⁱ·elementᵢ)` computed inside `compute_logup_term_column` (`lookup.rs:1610`, phase 1 at 1645-1666) via `BusValue::accumulate_fingerprint` (`lookup.rs:626`) / `Packing::accumulate_fingerprint_with` (`lookup.rs:274`, the canonical packing arithmetic). Chunked (1024 rows), one batch-inverse per chunk. +- Per-row multiplicity: `Multiplicity::evaluate_at_row` (`lookup.rs:1403`), variants `One/Column/Sum/Negated/Diff/Sum3/Linear` (`lookup.rs:1328`). +- Standalone per-row eval helpers already exist for the GPU parity path: `eval_fingerprint` (`logup_gpu.rs:498`) and `eval_term` (`logup_gpu.rs:529`). +- `build_auxiliary_trace` (`lookup.rs:1103`) orchestrates: GPU-resident build (`logup_gpu::try_build_aux_resident_gpu`, 1149-1164) → GPU term-column build (1168-1174) → CPU build → writes committed columns, then `build_accumulated_column_from_terms` (`lookup.rs:1740`) which returns **L = table_contribution** (Σ all terms) and builds the circular acc column with per-row offset `L/N`. + +**Constraint emission (single-source):** `emit_logup_constraints` (`lookup.rs:2256`) emits, at indices `idx_base..`: one degree-3 **batched-term constraint per committed pair** (`emit_logup_batched_term`, 2136: `c·fp_a·fp_b − ±m_a·fp_b − ±m_b·fp_a`) then ONE **accumulated constraint** (`emit_logup_accumulated`, 2184: `(acc_next − acc_curr − Σterms + L/N)·f₁[·f₂] − …`, degree `1+absorbed`). Helpers: `emit_multiplicity` (1925), `emit_linear_terms` (1945), `emit_packing_fingerprint` (1982), `emit_busvalue_fingerprint` (2066), `emit_fingerprint` (2114 — note base-operand-LEFT: `z − bus` written `−(bus − z)`). `logup_max_degree` (2239). These run through the generic `ConstraintBuilder` (builder hooks: `main/aux/challenge/alpha_pow/table_offset/fold_fingerprint_term` at `constraints/builder.rs:104-154`), so ONE body serves ProverEvalFolder, VerifierEvalFolder, CaptureBuilder (IR), MetaBuilder. + +**Challenges:** `LOGUP_NUM_CHALLENGES = 2` (`z` at index 0, `alpha` at index `LOGUP_CHALLENGE_ALPHA = 1`), `lookup.rs:102-105`. `build_rap_challenges` (`lookup.rs:1268`) samples them. + +## 2. #823 forward accumulation (already on main) + +`acc[0] = 0` (boundary constraint pinned in `AirWithBuses::boundary_constraints`, `lookup.rs:1286-1295`); acc column built by writing `acc[row]` BEFORE folding the row's terms (`lookup.rs:1768-1779`): `acc[i+1] − acc[i] = row_terms[i] − L/N`. Consequence: **`acc_next` is the ONLY next-row operand in the entire constraint system** — `AIR::trace_ood_next_row_columns` (`lookup.rs:1006-1017`) returns just `[main_width + acc_column_idx]`, which drives g·z OOD pruning (proof carries a full current-row block + a 1-column next-row block; `ood.rs::OodLayout`, built at `prover.rs:1447-1456`, mirrored by the verifier at `verifier.rs:1438-1453`). A GKR port that reshapes the aux trace changes `trace_ood_next_row_columns`, the OOD blocks, DEEP coefficients, and proof shape simultaneously. + +## 3. `crypto/stark/src/prover.rs` (3325 lines) — round structure & Fiat-Shamir order + +`multi_prove` (`prover.rs:2319`) drives everything (`prove` at 3020 is the 1-table wrapper): + +1. **Pre-pass**: Domain+LdeTwiddles dedup cache (2356-2393); VRAM-budgeted chunking `plan_table_chunks` (2395-2422). +2. **R1 Phase A — main commits** (2444-2509): per-table `commit_main_trace` (chunked parallel; under cuda takes `device_only` from `device_only_for` → `gpu_lde::device_only_gate`, `prover.rs:2477-2491`), then **sequential** transcript absorb: `precomputed_root?` then main `root` per table (2500-2503). +3. **R1 Phase B — shared LogUp challenges** (2521-2530): sample `z, α` from the SHARED transcript (only if any AIR `has_aux_trace`). +4. **R1 Phase C pass 1 — aux build** (2543-2601): parallel `air.build_auxiliary_trace(trace, &lookup_challenges)` per table → `Vec>`; device-resident main trace handles threaded in for the GPU aux build (2564-2571), dropped after (2592-2601). +5. **Transcript forking** (2632-2641): per-table fork = clone shared transcript + `append_bytes(idx as u64 LE)` when `num_airs > 1`. +6. **R1 Phase C pass 2 — aux commits** (2643-2838): per table, resident-GPU aux LDE (2704-2734, hard-abort if resident build fired but LDE fails) / row-major GPU LDE (2738-2773) / CPU row-major LDE (2775-2819); **aux root absorbed into the table's FORKED transcript** (2828-2836). +7. **Rounds 2–4 per table** (2898-2965): `build_round1` from cached LDE, then **absorb `bus_public_inputs.table_contribution` (L) into the forked transcript** (2954-2956), then `prove_rounds_2_to_4` (3048): + - **R2**: sample β (3068); β-powers split into transition then boundary coefficients (3082-3089); `round_2_compute_composition_polynomial` (1224) → `ConstraintEvaluator::evaluate` (CPU or GPU-fused, see §10) → quotient decompose (d2 fast path 1265-1271) → composition Merkle (GPU tree kept device-resident 1332-1351); absorb composition root (3102). + - **R3**: sample `z` via `sample_z_ood` (3109); barycentric OOD evals (1382-1439); absorb **pruned OOD blocks** (current-row block then next-row block, column-major; 3130-3138) then `Hᵢ(z^N)` (3141-3143). + - **R4** (1459): sample γ (1475); DEEP coefficients = γ-powers, trace terms first then composition parts (1484-1498); DEEP eval; FRI commit phase `fri::commit_phase_from_evaluations` (1529, samples ζₖ per layer, absorbs layer roots, final-poly coeffs; early-termination-aware `FriFoldLayout` in `fri/terminal.rs`); grinding nonce (1543-1550); query iotas (1553); openings (1562). + - Proof assembled at 3183-3214. + +**GPU gating**: `device_only_gate` (`gpu_lde.rs:187-211`) = Goldilocks-ext3 tower (TypeId check, `gpu_lde.rs:216`) ∧ not disabled by env ∧ lde_size power-of-2 ≥ threshold ∧ n ≥ bary threshold ∧ **not preprocessed** ∧ **contiguous transition offsets `[0,1,…]`** ∧ **uniform zerofier** ∧ not debug-checks. When false, everything falls back to CPU per-feature (each `try_*` returns `None` → CPU path). + +## 4. `crypto/stark/src/verifier.rs` (1676 lines) + +`multi_verify_views` (`verifier.rs:1126`): per-table pre-checks — composition part count derived from AIR not proof (1156-1167, #699), `ood_blocks_well_formed` (#815 width/shape guard, 1176), preprocessed commitment equality (1179-1204); absorb main roots (1203-1207) → sample shared LogUp challenges (1216-1222) → **bus_public_inputs presence must match `has_trace_interaction`** (1232-1246) → per-table fork (+idx domain sep) → absorb aux root → **absorb L** (1266-1273) → `verify_rounds_2_to_4` (1517) which calls `replay_rounds_after_round_1` (1356: β → composition root → z → OOD blocks → Hᵢ(z^N) → γ → per-layer ζ+root → final-fold ζ → final-poly coeffs → nonce → iotas) then `step_2_verify_claimed_composition_polynomial` (199), `step_3_verify_fri` (387), `step_4_verify_trace_and_composition_openings` (602). **Global bus balance** at 1310-1330: `Σ table_contribution == expected_bus_balance` (nonzero target supported for externally-computed contributions, e.g. COMMIT bus). This is THE cross-table soundness check GKR replaces. + +OOD constraint check (step 2): rebuilds frame from pruned blocks, `TransitionEvaluationContext::new_verifier` with `logup_alpha_powers` + `logup_table_offset = L/N` (296-327), `air.compute_transition` (VerifierEvalFolder), zerofiers from `constraints_meta` (331-340), compares against Horner-combined `Hᵢ(z^N)` (351-362). + +## 5. Proof struct — `crypto/stark/src/proof/stark.rs` + +`StarkProof` (`proof/stark.rs:72`): `trace_length`, main/aux/precomputed roots, `trace_ood_evaluations` + `trace_ood_next_evaluations` (pruned block), `composition_poly_root`, `composition_poly_parts_ood_evaluation`, `fri_layers_merkle_roots`, `fri_final_poly_coeffs`, `query_list`, `deep_poly_openings`, `nonce`, **`bus_public_inputs: Option`** (`lookup.rs:1524` — production field is just `table_contribution: FieldElement`; per-bus maps are debug-checks-only and rkyv-Skipped), `public_inputs`. **rkyv is the authoritative wire format** (verifier reads archives **in place** via `StarkProofView`, #769 — `proof/` module has the view layer); serde kept only for the cross-version examples CLI. Any new GKR proof fields must get rkyv Archive types AND zero-copy view accessors. + +## 6. `crypto/stark/src/traits.rs` — AIR trait surface (338 lines) + +Key methods (all at `traits.rs:122-338`): `step_size`, `name`, `build_auxiliary_trace` (returns `Option`), `build_rap_challenges`, `trace_layout() -> (main, aux)`, `has_aux_trace`, `has_trace_interaction`, `max_bus_elements`, `is_preprocessed`/`num_precomputed_columns`/`precomputed_commitment`, `trace_ood_next_row_columns` (g·z pruning set; conservative default = all), `composition_poly_degree_bound`, **`compute_transition`** (verifier/OOD; this monomorphization IS the recursion-guest path), `num_base_transition_constraints`, **`compute_transition_prover`** (base/ext split slices), `constraints_meta() -> &[ConstraintMeta]`, **`constraint_program()`** (default panics; guest must never call), `boundary_constraints(pub_inputs, rap_challenges, bus_public_inputs, trace_length)`, `context`, `transition_zerofier_evaluations_grouped`. `TransitionEvaluationContext` (traits.rs:66) carries `logup_alpha_powers` + `logup_table_offset` in BOTH variants. + +## 7. `prover/src/tables/mod.rs` + +The branch's diff there (commit `d5666021`, max_rows caps) is **superseded**: main has a full `max_rows` module (`prover/src/tables/mod.rs:82-98`) with LT/LOAD/BRANCH/MEMW_R = 2^20 plus `MaxRowsConfig` (105). Note main now has ~24 table modules (cpu32, keccak×3, dvrm, ecdas, ecsm, eq, bytewise, store, memw_aligned, memw_register, global_memory, local_to_global, commit… — `mod.rs:24-51`), vs ~11 at the GKR merge-base: **the interaction/table count a batch-GKR must handle roughly doubled** (26 production AIRs per `constraint_program_tests`). + +## 8. Multi-table orchestration + +Cross-table loop = `multi_prove` (`prover.rs:2319`) / `multi_verify_views` (`verifier.rs:1126`). Challenge sharing: ONE shared transcript absorbs all main roots, then z/α sampled once; everything after runs on per-table forks (domain-separated by index). **A batch GKR across all tables' interactions must run between Phase A (all main commits absorbed) and the per-table forking** — exactly where z/α sampling + aux build sit today — and its transcript interactions must happen on the SHARED transcript. VM-level assembly: `VmAirs` (`prover/src/lib.rs:478`), `multi_prove` call sites `prover/src/lib.rs:~1147,~1377`. + +Per-file reshaping history since merge-base `7b42e51b`: lookup.rs ← #514/#548/#591/#699/#650/#696/#764/#762/#769/#823; prover.rs ← ~29 PRs (heaviest: #591/#637/#648/#582/#699/#650/#700/#715/#735/#748/#762/#729/#799/#823); verifier.rs ← #592/#591/#606/#699/#735/#726/#764/#783/#729/#769/#815/#826/#823. + +## 9. Namespace collisions & the recursion guest + +**No `gkr.rs`, `sumcheck.rs`, or `lagrange_kernel.rs` anywhere on main** — the branch's new files land without collision. **Recursion guest**: `bench_vs/lambda/recursion` (main.rs) compiles `lambda_vm_prover::recursion::verify_and_attest_blob` (`prover/src/recursion.rs`) → `verify_recursion_blob` (`prover/src/lib.rs:307`) → the in-place rkyv archived verifier. It is **std via build-std** (not no_std; prove-side code DCE'd), reads the archive zero-copy, and runs `crypto/stark`'s verifier directly — so any GKR verifier code becomes **in-circuit guest code**: no HashMap (hashing in-circuit is prohibitively expensive — project rule), avoid allocation-heavy layer verification, and `constraint_program()`/CaptureBuilder must stay unreachable from the verify path. PR #768 (batched-FRI recursive verifier, active) measures guest cycles on this path — a GKR verifier adds sumcheck-round transcript work to exactly this cost profile. + +## 10. GPU residency vs. an aux-trace reshape + +The aux trace is GPU-built and GPU-resident on main: `logup_gpu.rs` lowers interactions to a `FingerprintDescriptor` (`logup_gpu.rs:108`, `build_fingerprint_descriptor`:273 → `math_cuda::logup::LogupDescriptor`:322); `try_build_aux_resident_gpu` (418) builds term + acc columns on device, returns L, and `trace.set_aux_resident(ra)` keeps them for the aux LDE (`try_expand_leaf_and_tree_ext3_row_major_keep_dev`, `prover.rs:2704-2734`); under `device_only` the host never sees main or aux LDEs (#799). The **fused GPU composition** (`constraints/evaluator.rs:304-393` → `constraint_ir/gpu_interp.rs::try_eval_composition_gpu`:197) consumes `air.constraint_program()` — **which includes the LogUp constraints** (captured in `AirWithBuses::constraint_program`, `lookup.rs:1088-1101`) — plus `logup_alpha_powers`, `logup_table_offset`, uniform `z_inv`, and the boundary spec, producing `H` on device. `DeviceProgram::lower` (`constraint_ir/device.rs:125`) is Goldilocks-hardcoded; CPU parity oracle `eval_device_program` (device.rs:247); `constraint_program_tests` pins interpreter==folders for all 26 AIRs (the GPU contract). + +Consequences for GKR: (a) if the aux trace becomes {Lagrange kernel + σ}, `logup_gpu.rs`'s builder and the resident-aux LDE path must either be taught the new shape or gated off; (b) whatever replaces the LogUp transition constraints must still be emitted through `ConstraintBuilder` or the captured `ConstraintProgram` (and hence GPU composition + the parity tests) silently diverges from the folders — the emission path is the contract; (c) any new committed aux column is ext3 and flows through the existing aux commit machinery unchanged as long as it lives in the aux segment; (d) the GKR sumcheck itself is a NEW workload with no existing kernel — CPU-first with the tower/TypeId gating pattern (`is_goldilocks_ext3_tower`, `gpu_lde.rs:216`) is the established fallback idiom. diff --git a/thoughts/logup-gkr/port-plan.md b/thoughts/logup-gkr/port-plan.md new file mode 100644 index 000000000..7a41edb91 --- /dev/null +++ b/thoughts/logup-gkr/port-plan.md @@ -0,0 +1,176 @@ +# Plan: bring LogUp-GKR (PR #485) up to date with the current prover + +> Written 2026-07-21 against main `a8648320`. Execution target: an implementation agent working phase-by-phase. +> Companion docs in this directory — READ ALL THREE BEFORE STARTING: +> - `branch-analysis.md` — what the GKR branch implements (protocol, wiring semantics, FS order, soundness status). This is the **semantic spec** of what to port. +> - `main-surface-map.md` — the current shape of every integration point on main (verified file:line anchors). +> - `drift-analysis.md` — per-commit classification, what to drop, why rebase is not viable. +> +> Also read the project memory doc `single-source-constraints-architecture.md` (auto-memory) if available — it explains the constraint-emission contract this port must respect. + +## 0. Executive summary + +PR #485 (`feat/gkr-logup`, tip `78be304f`, closed draft) replaces committed LogUp aux columns (⌈K/2⌉ term columns + 1 acc per table) with a **batch GKR proof over fractional summation trees** plus exactly **2 committed aux columns** per table (Lagrange kernel `l`, bridge running sum `σ`). Measured on the PR: **peak heap −70 % (221 GB → 66 GB), prove time +2.6 %**. That memory win is the point of resurrecting it (it composes with continuations: bigger epochs per box). + +Main has moved ~148 commits since the merge-base: the entire constraint system was replaced (single-source constraints #764/#772), the acc column became the sole next-row OOD read (#823), the proof wire format became rkyv-in-place (#769), and the prover is now GPU-resident end-to-end (#748/#798/#799). **A git rebase/cherry-pick is not viable** (all 9 wiring files conflict against rewritten code). The strategy is: + +1. **Port the three new modules wholesale** from the branch tip (`gkr.rs`, `sumcheck.rs`, `lagrange_kernel.rs`) — they are self-contained, conflict-free, and already contain all five soundness-fix commits. +2. **Re-implement the wiring by hand** against current main, as an **opt-in mode** (`LogUpMode::Gkr`), keeping standard LogUp the default and byte-identical. +3. Fix the known small defects during the port (dead proof fields, a verifier deserialization panic); **document but do not yet fix** the one open protocol-level soundness gap (multi-interaction leaf binding — see §6), which becomes the immediate follow-up. + +## 1. Ground rules (non-negotiable, from project conventions) + +- **Opt-in mode, default OFF.** With GKR disabled, the build must be **wire-identical to main**: same transcript, same proof bytes accepted. Gate: cross-version verification (old binary ↔ new proofs, both directions) — this is the king gate for constraint-system changes; proof determinism is NOT a goal (proofs are nondeterministic by design). +- **All new/changed constraints go through the single-source `ConstraintBuilder` emission path.** The bridge constraint must be ONE `emit`-style body serving all four interpretations (ProverEvalFolder / VerifierEvalFolder / CaptureBuilder→IR / MetaBuilder-derived meta). Never a hand-written parallel evaluator. Constraint bodies are declarative — no micro-opts without a measured bench win. +- **Verifier code is recursion-guest code.** `crypto/stark/src/verifier.rs`, `gkr.rs` (verify paths), `sumcheck.rs`, `lagrange_kernel.rs` all compile into the RV64 guest (`prover/src/recursion.rs` → `bench_vs/lambda/recursion`). No `HashMap`/`HashSet` on any verify-reachable path (audit the ported files — replace with sorted `Vec` + linear scan/binary search if found). Keep allocations in layer verification minimal. +- **Never trust proof-carried data for control flow or randomness.** Mode comes from AIR/verifier configuration, never from proof-field presence. All GKR randomness comes from the transcript (the branch already fixed this — preserve it). +- **Hygiene:** `cargo fmt` + `make lint` (workspace-level, not per-package clippy) before every push. No AI attribution anywhere. Bench runs happen on the bench server via PR comment (`/bench`, `/bench-abba`), never locally. + +## 2. Source material + +```bash +git fetch origin feat/gkr-logup # tip 78be304f, merge-base 7b42e51b +git show origin/feat/gkr-logup:crypto/stark/src/gkr.rs > /tmp/gkr.rs +git show origin/feat/gkr-logup:crypto/stark/src/sumcheck.rs > /tmp/sumcheck.rs +git show origin/feat/gkr-logup:crypto/stark/src/lagrange_kernel.rs > /tmp/lagrange_kernel.rs +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/lookup.rs # wiring spec +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/prover.rs +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/verifier.rs +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/proof/stark.rs +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/traits.rs +git diff 7b42e51b..origin/feat/gkr-logup -- crypto/stark/src/tests/ +``` + +Exclusions (do NOT port — see drift-analysis §1/§5): everything touching `fri/fri_functions.rs`, `constraints/evaluator.rs` (timers), `prover/src/tables/mod.rs`, the `commit_columns_bit_reversed`/`LdeTwiddles` hunks in prover.rs, the scaling-improvements docs spec. Ignore branches `feat/gkr-logup_opt` and `feat/logup_gkr_v2` entirely. Defer PR #489 (disk spill). + +## 3. Target design on current main + +### 3.1 Mode selection + +Add `LogUpMode { Standard, Gkr }`: + +- Carried by `AirWithBuses` (constructor parameter or a `with_logup_mode` builder), stored alongside `LogUpLayout`. In `Standard` mode every code path is exactly today's. +- VM assembly: thread a mode flag through `VmAirs` construction (`prover/src/lib.rs:478`) — one switch for all tables. For experiments, wire an env var (e.g. `LAMBDA_LOGUP_GKR=1`) at the `VmAirs` call sites; library API takes the enum explicitly. +- **Prover and verifier must agree out-of-band** (same `VmAirs` config). The verifier decides expectations from ITS mode: in `Gkr` mode `batch_gkr_proof` is required and `bus_public_inputs` must be absent; in `Standard` mode a proof carrying GKR fields is rejected (fail-closed both ways). + +### 3.2 Aux layout in GKR mode + +Per interacting table: aux = `[l (kernel), σ (bridge sum)]`, ext3 columns, aux width 2. Constants: `GKR_AUX_KERNEL_COL = 0`, `GKR_AUX_SIGMA_COL = 1`. + +- `trace_layout()` → `(main_width, 2)`. +- `trace_ood_next_row_columns()` → `[main_width + GKR_AUX_SIGMA_COL]` (σ is the only next-row read; the kernel is current-row only). This slots directly into the #823 OOD-pruning machinery (`ood.rs::OodLayout`) — the pruned next-row block stays width 1. +- Boundary constraints (mode-aware in `AirWithBuses::boundary_constraints`, replacing the `acc[0]=0` pin): `l[0] = ∏ⱼ(1−rⱼ)` and `σ[0] = 0`. +- `composition_poly_degree_bound`: `logup_max_degree` becomes mode-aware — bridge constraint is degree 2 (vs 3 for standard batched terms), so the bound is `max(CS::max_degree(), 2)` in GKR mode. The verifier derives composition part count from the AIR (#699), so this stays consistent automatically — but note proof shape differs between modes by design. + +### 3.3 Constraint emission (the part that MUST be re-designed, not transliterated) + +The branch's `LookupBridgeSumConstraint` is a boxed `TransitionConstraint` — that world is gone. Write `emit_logup_gkr_constraints(builder, layout, num_base)` next to `emit_logup_constraints` (`lookup.rs:2256`), emitting exactly ONE ext constraint via `b.emit_ext(idx, expr)`: + +``` +σ_next − σ_curr − l_curr · batched_curr + Δ = 0 (degree 2, RowDomain::ALL) +batched_curr = Σⱼ γʲ·colⱼ(curr) + γᴷ·l_curr +``` + +where `colⱼ` ranges over the K distinct main columns referenced by any interaction (sorted order from `extract_column_indices`), and `Δ`, `γ`-powers, and the random point come from the extended rap-challenges vector (§3.4). Follow the invariants: base-operand-LEFT for base×ext products (`colⱼ · γʲ`, not `γʲ · colⱼ`), `const_base`/`const_signed` as the only constant path, exact-once emission per index. `MetaBuilder` then derives meta for free; `AirWithBuses::new` selects which emit body to run based on mode (both for meta derivation at construction and for `compute_transition{,_prover}` at runtime). The captured `ConstraintProgram` (lazy OnceLock) automatically reflects the mode — which keeps `constraint_program_tests`' folder==interpreter contract meaningful in both modes. + +Access to challenges inside the body goes through the existing builder hooks (`challenge(idx)`, `alpha_pow`, `table_offset` — see `constraints/builder.rs:104-154`). Reuse `table_offset` for `Δ` (it plays exactly the acc-offset role `L/N` plays today, so `TransitionEvaluationContext.logup_table_offset` carries it unchanged) and map γ-powers onto `logup_alpha_powers` OR add a parallel `gkr_gamma_powers` slot in `TransitionEvaluationContext` — decide by whichever avoids overloading semantics the recursion guest also compiles; prefer a new explicit field. + +### 3.4 Extended rap-challenges layout (GKR mode) + +Adopt the branch's layout (branch lookup.rs diff 460-495): `[0]=z, [1]=α, [2]=γ, [3]=Δ, [4..4+K+1]=γ⁰..γᴷ, [4+K+1..]=r (instance random point, n_vars elements)`. Keep the branch's named constants (`LOGUP_CHALLENGE_Z/GAMMA`, `LOGUP_BRIDGE_OFFSET_IDX`, `LOGUP_GAMMA_POWERS_START`, `logup_random_point_start`). Note z/α/γ are shared across tables; Δ, the γ-power count (K varies per table), and r (length = log₂ trace_length, varies per table) are per-table. + +### 3.5 Prover flow (into `multi_prove`, `prover.rs:2319`) + +Insert between Phase B (z/α sampling, `prover.rs:2521-2530`) and Phase C pass 1 (aux build, 2543): + +- **Phase B′ (GKR batch prove, shared transcript):** collect interacting tables; build leaf fractions + layer trees in parallel (`compute_logup_leaf_fractions`/`compute_logup_layers`, ported from the branch — they read `air.bus_interactions()` + host main-trace columns); `gkr_prove_batch(instances, transcript)`; `finalize_logup_gkr_result` per table (kernel-weighted `column_claims`, `l_mle_claim`). Free the layer trees before the LDE phase (this is where the memory win lives). +- **Phase B″ (binding + γ):** append every table's `column_claims` to the shared transcript (ascending table index, sorted column order), then sample the single shared `γ`. +- **Aux build stays in `build_auxiliary_trace`** (do NOT copy the branch's inline-in-prover.rs hack — it only existed because the old code structure forced it). Pass the full extended challenge vector (§3.4) into the existing Phase-C-pass-1 parallel call; in GKR mode the impl writes `l` (from `lagrange_kernel.rs`) and fills `σ` forward (`σ[0]=0`; per-row `batched` precomputed row-parallel, σ filled sequentially), and returns `bus_public_inputs = None`. +- Fork-time `append L` (`prover.rs:2954-2956`) is skipped in GKR mode (bpi is None; make the skip explicit and mode-checked, not accidental-on-None). +- Rounds 2-4 unchanged. `MultiProof` gains `batch_gkr_proof`; per-table `StarkProof` gains `column_claims` (drop the branch's dead `random_point` field and the stub `gkr_proof` — see branch-analysis §5). Single-table `prove` wrapper carries the batch proof (branch BUG-014). +- **Main-trace host access:** leaf-fraction building reads host main-trace columns — these exist regardless of GPU residency (#799 keeps LDE on device, not the raw trace). Verify this assumption at implementation time (`trace.columns_main()` availability in `multi_prove` before Phase C). + +### 3.6 Verifier flow (into `multi_verify_views`, `verifier.rs:1126`) + +Mirror insertion after shared z/α sampling (1216-1222), replacing the bpi-presence check (1232-1246) with the mode-aware rule from §3.1: + +- Require `batch_gkr_proof`; compute per-table `n_vars = trace_length.trailing_zeros()` (validate `trace_length` is a power of two ≥ 1 first); `gkr_verify_batch(proof, n_layers_by_instance, transcript)` → shared random point + per-instance `(n̂, d̂)` claims. +- Per table: check `column_claims.len() == extract_column_indices(...).len()`; run `reconstruct_and_verify_gkr_claims` (ported as-is, gap documented — §6); derive `r = instance_eval_point(shared_point, n_vars)` **from the transcript-derived point only**. +- Append column_claims + sample γ exactly as the prover (Phase B″ mirror). +- Per-table: fork, absorb aux root, skip the L absorb in GKR mode, build extended rap challenges, `verify_rounds_2_to_4` unchanged — bridge checked as an ordinary OOD transition constraint (VerifierEvalFolder runs the same emit body), kernel bound by the `l[0]` boundary constraint + the γᴷ·l self-check inside the bridge (`Σl² = ∏(rₖ²+(1−rₖ)²)` folded into Δ/target — see branch-analysis §1.5). +- **Global bus balance replacement:** `Σ_instances root_nᵢ·root_dᵢ⁻¹ == expected_bus_balance` (reject on any zero root denominator). Must support the nonzero `expected_bus_balance` target exactly like today's check at `verifier.rs:1310-1330` (continuations/COMMIT-bus depend on it). +- Keep #815-style shape guards; the `ood_blocks_well_formed` check derives from the AIR so aux-width-2 flows through. + +### 3.7 Proof format (`proof/stark.rs` + the rkyv view layer) + +New types (from branch, minus dead fields): `BatchGkrProof { root_claims: Vec<(E,E)>, layer_proofs: Vec }`, `BatchGkrLayerProof { sumcheck_proof, child_claims_by_instance: Vec<[E;4]> }`, `SumcheckProof`/`RoundPoly` from sumcheck.rs. All must derive **rkyv Archive/Serialize/Deserialize** (wire format, #769) AND get zero-copy accessors on the `StarkProofView`/multi-proof view layer, plus serde for the examples CLI. `MultiProof.batch_gkr_proof: Option<...>`, `StarkProof.gkr_column_claims: Option>` (or similar; rkyv-friendly index type). In-place validation rules for the new archived types must enforce the §5 length invariants at access time (the view layer is how the verifier reads untrusted bytes). + +### 3.8 GPU path in GKR mode: v1 = CPU-only, explicitly gated + +The GPU stack encodes the standard-LogUp aux shape and constraint set: +- `logup_gpu.rs::try_build_aux_resident_gpu` + term-column builder: term/acc layout → **must not fire** in GKR mode. +- Fused GPU composition (`constraint_ir/gpu_interp.rs::try_eval_composition_gpu`) consumes the captured program + `logup_alpha_powers`/`logup_table_offset` plumbing; the bridge constraint changes that challenge surface. + +v1 decision: in GKR mode, force `device_only = false` and make every GPU `try_*` return None (cleanest: add `logup_mode == Standard` to `device_only_gate` and to the aux-GPU/composition-GPU entry conditions, each with a comment naming the assumption it protects). Main-trace GPU LDE/Merkle MAY be left on where it is shape-agnostic, but only if the `device_only` interactions are verified — when in doubt, CPU everything under GKR mode. The original −70 % heap / +2.6 % time bench was CPU-vs-CPU, so the A/B story survives. GPU support for GKR (aux build of l/σ on device, bridge in the fused kernel, GKR sumcheck kernels) is follow-up work (§8). + +## 4. Execution phases (each ends compiling + green on its gate) + +**Phase 0 — Branch + module port.** New branch off main (e.g. `feat/logup-gkr-v3`). Add `gkr.rs`, `sumcheck.rs`, `lagrange_kernel.rs` from the branch tip; register in `lib.rs`; fix compile drift (field-trait APIs, transcript API, rand/test-helper churn since April). Port their in-module tests. Audit all three files for HashMap/HashSet on verify paths (§1). Gate: `cargo test -p lambda-vm-crypto-stark gkr sumcheck lagrange` green (adjust package/filter names to the workspace layout). + +**Phase 1 — Lookup-layer adapter (`lookup.rs`).** `LogUpMode`; `extract_column_indices`, `compute_logup_leaf_fractions`, `compute_logup_layers`, `finalize_logup_gkr_result`, `compute_bridge_params`, `extend_rap_challenges_with_bridge`, `reconstruct_and_verify_gkr_claims` (ported from branch diffs, adapted to current `BusInteraction`/`Multiplicity`/`Packing` — reuse `logup_gpu.rs::eval_fingerprint`-style per-row helpers where they fit); `emit_logup_gkr_constraints` per §3.3; mode-aware `AirWithBuses` (aux layout, meta derivation, boundary constraints, `trace_ood_next_row_columns`, degree bound, `build_auxiliary_trace`). Gate: unit tests — leaf fractions vs a direct per-row `m/fp` sum; `Σ leaf fractions == standard-mode table_contribution` on a small AIR (the two modes MUST agree on the total, this is the cross-mode oracle); meta/emit parity (MetaBuilder-derived meta matches expectations; `EmitTracker`-style exact-once holds). + +**Phase 2 — Proof format + views.** §3.7 types with rkyv + serde + view accessors + length-validating access. Gate: roundtrip serialization tests (rkyv archive → view → values; serde for CLI), plus a malformed-bytes rejection test through the view layer. + +**Phase 3 — Prover wiring.** §3.5 into `multi_prove` + single-table `prove` + `Round1Metadata` extension + GPU gating per §3.8. Gate: `Standard` mode untouched (full existing prover test suite green); `Gkr` mode single-table and multi-table prove runs produce proofs (verified in Phase 4). + +**Phase 4 — Verifier wiring.** §3.6 into `multi_verify_views` + rounds replay + balance check. Add the branch's DoS guards AND the residual fix: bound `RoundPoly.evals` length (exactly 4 for degree-3 gates; reject otherwise — never assert), guard `1u64 << n_unused`, validate `trace_length`. Gate: GKR-mode prove→verify roundtrip green on unit AIRs + fibonacci examples; tampering any transcript-bound value (root_claims, column_claims, child_claims, σ OOD, aux root) → reject. + +**Phase 5 — Test port + new tests.** Re-author the branch's `bus_tests` additions against the current harness (list in branch-analysis §8): the 5 GKR soundness tests, the single-table completeness test; KEEP the standard-mode tests they replaced (both modes coexist now — do not delete `test_tampered_table_contribution` etc.). Add: (a) a mode-mismatch test (Standard verifier rejects GKR proof and vice versa), (b) the **fabricated-leaves forgery test** for the §6 gap, `#[ignore]`d with a comment naming the gap (it documents the known hole and becomes the acceptance test for the follow-up fix). Gate: full `cargo test` workspace green; `make lint` clean. + +**Phase 6 — System validation.** +1. **Standard-mode wire-identity (the king gate):** cross-version verification against an origin/main binary, both directions, examples + full VM programs (6/6). Any mismatch with GKR disabled = a bug in the mode plumbing. +2. **GKR-mode e2e:** full VM prove+verify on the 6 test programs + `executor/tests/ethrex_{5,20}_transfers.bin`, GKR on, CPU path. +3. **Memory + time:** the PR-comment bench on the bench server (`/bench` for the cheap tier; `/bench-abba` for the paired run) — GKR mode on vs main. Expect the heap win to have shifted since April (table count ~doubled; aux-heavy tables like CPU K=40: 21→2 aux cols still apply). Report peak-heap and prove-time deltas in the PR description. Do not run benches locally. +4. **Recursion-guest impact (measure, don't fix):** build the recursion guest with a GKR proof as input and record guest-cycle delta vs standard (the GKR verifier adds sumcheck transcript work in-circuit). This number decides whether recursion needs its own follow-up. + +**Phase 7 — PR.** Draft PR referencing #485, with: motivation (memory numbers old + new), mode semantics, the documented §6 gap + follow-up plan, exclusions (drift-analysis §5). No AI attribution. + +## 5. Defect fixes folded into the port (small, mandatory) + +| Defect (branch-analysis ref) | Fix | +|---|---| +| Dead `LogUpGkrProof.random_point` + stub `gkr_proof` fields (§5) | Drop; keep only `column_claims` per table + the batch proof | +| `RoundPoly` deserialization → `sum_at_binary`/`evaluate` panic on <2 evals (§7.3) | Enforce eval-count == 4 at view/deserialization boundary; make the sumcheck helpers fallible or pre-checked | +| `1u64 << n_unused` debug overflow (§7.3) | Validate instance size spread / trace_length upstream; checked shift | +| Fork-time L absorb silently dead in GKR mode (§4 note) | Explicit mode-checked skip on both sides | + +## 6. KNOWN OPEN SOUNDNESS GAP — documented, follow-up, not this PR + +`reconstruct_and_verify_gkr_claims` is **fail-open for multi-interaction tables** (every production table): nothing binds the batch-GKR leaf claims `(n̂, d̂)` to the committed columns — the bridge binds only `column_claims ↔ trace`, and the leaf fraction is a nonlinear (cross-multiplied) combination of columns, so MLE evaluation doesn't factor through `column_claims`. A malicious prover can run an honest GKR over fabricated leaves (arbitrary root contribution → fake bus balance) while all present checks pass. Full analysis: branch-analysis §7 item 2. + +Why not fix it in this PR: it is a protocol change to the GKR **input layer**, and mixing it into the port destroys the only available correctness reference (the branch's own behavior). Port faithfully first, keep the experiment measurable, land the fix as the immediate next PR before any production consideration. + +Fix directions to evaluate in the follow-up (in rough preference order): +1. **Input-layer sumcheck**: one extra sumcheck reducing each instance's leaf claims `(n̂, d̂)` at point r to claims about the K column MLEs at r (the leaf is a low-degree polynomial in the columns per row — this is the standard LogUp-GKR construction; see Winterfell's and Stwo's logup-gkr input layers for reference shapes). +2. **Per-interaction instances**: leaves = `(±m_k, fp_k)` per interaction (numerator/denominator LINEAR in columns) → K GKR instances per table; MLE then commutes with the linear map, so `column_claims` directly verify the leaf claims. More instances, but the batch already handles mixed sizes. +The `#[ignore]`d forgery test from Phase 5 is the acceptance criterion. + +## 7. Risk register + +| Risk | Mitigation | +|---|---| +| FS-order mistake in the B′/B″ insertion | The order is specified twice (branch-analysis §4 = branch behavior; §3.5/3.6 here = target). Prover and verifier are written from the SAME spec; every tamper test doubles as an FS check | +| Breaking standard mode subtly | Mode plumbed as data, no behavioral change when `Standard`; Phase 6.1 cross-version gate is decisive | +| GPU path fires on a GKR-shaped table | Explicit mode condition in `device_only_gate` + aux/composition `try_*` entries; add a debug assert that resident-aux never coexists with `Gkr` mode | +| Guest bloat / cycle blowup in recursion | Phase 6.4 measures; no HashMap audit in Phase 0; follow-up decides | +| rkyv view layer holes (panics on malformed archives) | Phase 2 malformed-bytes tests; length checks at accessor level (§5) | +| Memory win regressed by 26-table batch (more instances, deeper trees) | Layer trees are transient; free before LDE (§3.5). Phase 6.3 measures the real number | +| `trailing_zeros` on non-power-of-two/zero trace length | Validated upfront in verifier (§3.6) | + +## 8. Follow-up queue (explicitly out of scope here) + +1. **Leaf-binding soundness fix** (§6) — required before any production use. +2. **GPU support for GKR mode**: device build of `l`/`σ`, bridge constraint through the fused composition kernel (it's already just an emitted constraint → captured IR; the challenge plumbing is the work), GKR layer/sumcheck kernels. +3. **PR #489 disk-spill** re-evaluation on top (orthogonal memory scaling). +4. **Rayon FRI fold** (`3aa03e6c`) as a separate tiny perf PR — genuinely absent on main. +5. Recursion-guest optimization of the GKR verifier if Phase 6.4 numbers demand it (batched-FRI work #768 is the sibling effort). From b1e5c88f72a49782965feb1aa9afd8ae2a8edbf0 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:39:17 -0300 Subject: [PATCH 03/20] style(gkr): clippy op-ref and clone-on-copy fixes in ported modules --- crypto/stark/src/gkr.rs | 40 ++++++++++++++--------------- crypto/stark/src/lagrange_kernel.rs | 4 +-- crypto/stark/src/sumcheck.rs | 34 ++++++++++++------------ 3 files changed, 39 insertions(+), 39 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index e00d64543..3eda95df0 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -2252,7 +2252,7 @@ mod tests { // Check: 1136/384 = 71/24 (both sides: 1136*24 = 27264, 71*384 = 27264) let root_n = &tree[2].numerators[0]; let root_d = &tree[2].denominators[0]; - assert_eq!(root_n * &FE::from(24u64), &FE::from(71u64) * root_d); + assert_eq!(root_n * FE::from(24u64), FE::from(71u64) * root_d); } #[test] @@ -2281,8 +2281,8 @@ mod tests { let root_n = &tree[3].numerators[0]; let root_d = &tree[3].denominators[0]; assert_eq!( - root_n * &acc.denominator, - &acc.numerator * root_d, + root_n * acc.denominator, + acc.numerator * root_d, "Tree root must equal sequential sum as a fraction" ); } @@ -2348,17 +2348,17 @@ mod tests { assert_eq!(evals.len(), 4); let one = FE::one(); - let one_minus_r0 = &one - &r0; - let one_minus_r1 = &one - &r1; + let one_minus_r0 = one - r0; + let one_minus_r1 = one - r1; // Index 0 = (b0=0, b1=0): (1-r0)*(1-r1) - assert_eq!(evals[0], &one_minus_r0 * &one_minus_r1); + assert_eq!(evals[0], one_minus_r0 * one_minus_r1); // Index 1 = (b0=1, b1=0): r0*(1-r1) - assert_eq!(evals[1], &r0 * &one_minus_r1); + assert_eq!(evals[1], r0 * one_minus_r1); // Index 2 = (b0=0, b1=1): (1-r0)*r1 - assert_eq!(evals[2], &one_minus_r0 * &r1); + assert_eq!(evals[2], one_minus_r0 * r1); // Index 3 = (b0=1, b1=1): r0*r1 - assert_eq!(evals[3], &r0 * &r1); + assert_eq!(evals[3], r0 * r1); } #[test] @@ -2419,7 +2419,7 @@ mod tests { gkr_prove(&tree, &mut transcript).unwrap(); // claimed_sum = 68/55 - let expected_sum = &FE::from(68u64) * &FE::from(55u64).inv().unwrap(); + let expected_sum = FE::from(68u64) * FE::from(55u64).inv().unwrap(); assert_eq!(proof.claimed_sum, expected_sum); // Should have 1 layer proof (root -> leaves) @@ -2441,8 +2441,8 @@ mod tests { // n_MLE(eta) = 3*(1-eta) + 7*eta = 3 + 4*eta // d_MLE(eta) = 5*(1-eta) + 11*eta = 5 + 6*eta let eta = &final_point[0]; - let expected_n = &FE::from(3u64) * &(&FE::one() - eta) + &(&FE::from(7u64) * eta); - let expected_d = &FE::from(5u64) * &(&FE::one() - eta) + &(&FE::from(11u64) * eta); + let expected_n = FE::from(3u64) * (FE::one() - eta) + (FE::from(7u64) * eta); + let expected_d = FE::from(5u64) * (FE::one() - eta) + (FE::from(11u64) * eta); assert_eq!(final_n_claim, expected_n); assert_eq!(final_d_claim, expected_d); } @@ -2471,7 +2471,7 @@ mod tests { gkr_prove(&tree, &mut transcript).unwrap(); // claimed_sum = root_n / root_d = 1136 / 384 - let expected_sum = &FE::from(1136u64) * &FE::from(384u64).inv().unwrap(); + let expected_sum = FE::from(1136u64) * FE::from(384u64).inv().unwrap(); assert_eq!(proof.claimed_sum, expected_sum); // Should have 2 layer proofs @@ -2519,7 +2519,7 @@ mod tests { let root_n = &tree[2].numerators[0]; let root_d = &tree[2].denominators[0]; - let expected_sum = root_n * &root_d.inv().unwrap(); + let expected_sum = root_n * root_d.inv().unwrap(); let mut transcript = DefaultTranscript::::new(&[]); let (proof, _, _, _) = gkr_prove(&tree, &mut transcript).unwrap(); @@ -2665,7 +2665,7 @@ mod tests { for (layer_idx, lp) in proof.layer_proofs.iter().enumerate() { let lambda: FE = replay.sample_field_element(); - let combined_claim = &n_claim + &(&lambda * &d_claim); + let combined_claim = n_claim + (lambda * d_claim); if lp.sumcheck_proof.round_polys.is_empty() { // Trivial layer: verify gate equation directly @@ -2673,7 +2673,7 @@ mod tests { let nr = &lp.child_claims[1]; let dl = &lp.child_claims[2]; let dr = &lp.child_claims[3]; - let gate_val = &(nl * dr) + &(nr * dl) + &(&lambda * &(dl * dr)); + let gate_val = (nl * dr) + (nr * dl) + (lambda * (dl * dr)); assert_eq!( combined_claim, gate_val, "Gate equation failed at trivial layer {}", @@ -2705,9 +2705,9 @@ mod tests { let eta: FE = replay.sample_field_element(); // Update claims for next layer - let one_minus_eta = &FE::one() - η - n_claim = &(&lp.child_claims[0] * &one_minus_eta) + &(&lp.child_claims[1] * &eta); - d_claim = &(&lp.child_claims[2] * &one_minus_eta) + &(&lp.child_claims[3] * &eta); + let one_minus_eta = FE::one() - eta; + n_claim = (lp.child_claims[0] * one_minus_eta) + (lp.child_claims[1] * eta); + d_claim = (lp.child_claims[2] * one_minus_eta) + (lp.child_claims[3] * eta); } } @@ -2843,7 +2843,7 @@ mod tests { let (mut proof, _, _, _) = gkr_prove(&tree, &mut prover_transcript).unwrap(); // Tamper with the claimed_sum - proof.claimed_sum = &proof.claimed_sum + &FE::one(); + proof.claimed_sum += FE::one(); // Verify with the tampered proof let mut verifier_transcript = DefaultTranscript::::new(&[0xAA]); diff --git a/crypto/stark/src/lagrange_kernel.rs b/crypto/stark/src/lagrange_kernel.rs index 59d5f9d5e..bd5a832fa 100644 --- a/crypto/stark/src/lagrange_kernel.rs +++ b/crypto/stark/src/lagrange_kernel.rs @@ -209,7 +209,7 @@ mod tests { assert_eq!(s[3], FE::from(21u64)); // Verify partition of unity: sum = 12 + (-18) + (-14) + 21 = 1 - let sum = &(&(&s[0] + &s[1]) + &s[2]) + &s[3]; + let sum = ((s[0] + s[1]) + s[2]) + s[3]; assert_eq!(sum, FE::one()); } @@ -229,7 +229,7 @@ mod tests { let mut sum = FE::zero(); for val in &s { - sum = &sum + val; + sum = sum + val; } assert_eq!( sum, diff --git a/crypto/stark/src/sumcheck.rs b/crypto/stark/src/sumcheck.rs index 83eb0a386..1510a5a5a 100644 --- a/crypto/stark/src/sumcheck.rs +++ b/crypto/stark/src/sumcheck.rs @@ -324,8 +324,8 @@ mod tests { let mut result = FE::zero(); let mut power = FE::one(); for c in coeffs { - result = &result + &(c * &power); - power = &power * x; + result += c * power; + power *= x; } result } @@ -619,8 +619,8 @@ mod tests { let p1 = &proof.round_polys[0].evals()[1]; // p(t) = p0*(1-t) + p1*t = p0 + (p1 - p0)*t let slope = p1 - p0; - let p2_expected = p0 + &(&slope * &FE::from(2u64)); - let p3_expected = p0 + &(&slope * &FE::from(3u64)); + let p2_expected = p0 + (slope * FE::from(2u64)); + let p3_expected = p0 + (slope * FE::from(3u64)); assert_eq!(proof.round_polys[0].evals()[2], p2_expected); assert_eq!(proof.round_polys[0].evals()[3], p3_expected); @@ -675,12 +675,12 @@ mod tests { let r2 = &challenges[1]; // f(r1, r2) = f(0,0)*(1-r1)*(1-r2) + f(1,0)*r1*(1-r2) + f(0,1)*(1-r1)*r2 + f(1,1)*r1*r2 let one = FE::one(); - let one_minus_r1 = &one - r1; - let one_minus_r2 = &one - r2; - let f_at_r = &(&(&evals[0] * &one_minus_r1) * &one_minus_r2) - + &(&(&evals[1] * r1) * &one_minus_r2) - + &(&(&evals[2] * &one_minus_r1) * r2) - + &(&(&evals[3] * r1) * r2); + let one_minus_r1 = one - r1; + let one_minus_r2 = one - r2; + let f_at_r = ((evals[0] * one_minus_r1) * one_minus_r2) + + ((evals[1] * r1) * one_minus_r2) + + ((evals[2] * one_minus_r1) * r2) + + ((evals[3] * r1) * r2); // The last round poly evaluated at the last challenge should give f(r1, r2) let last_round_eval = proof.round_polys[1].evaluate(&challenges[1]); @@ -771,7 +771,7 @@ mod tests { FE::from(3u64), // f(0,1,1) = 0*1 + 1 + 2 = 3 FE::from(4u64), // f(1,1,1) = 1*1 + 1 + 2 = 4 ]; - let claimed_sum: FE = evals.iter().fold(FE::zero(), |acc, v| &acc + v); // = 22 + let claimed_sum: FE = evals.iter().fold(FE::zero(), |acc, v| acc + v); // = 22 // Prove let mut prover_transcript = DefaultTranscript::::new(&[0xAB]); @@ -815,12 +815,12 @@ mod tests { let r1 = &challenges[0]; let r2 = &challenges[1]; let one = FE::one(); - let one_minus_r1 = &one - r1; - let one_minus_r2 = &one - r2; - let mle_at_r = &(&(&evals[0] * &one_minus_r1) * &one_minus_r2) - + &(&(&evals[1] * r1) * &one_minus_r2) - + &(&(&evals[2] * &one_minus_r1) * r2) - + &(&(&evals[3] * r1) * r2); + let one_minus_r1 = one - r1; + let one_minus_r2 = one - r2; + let mle_at_r = ((evals[0] * one_minus_r1) * one_minus_r2) + + ((evals[1] * r1) * one_minus_r2) + + ((evals[2] * one_minus_r1) * r2) + + ((evals[3] * r1) * r2); assert_eq!( final_eval, mle_at_r, From 2a3c13b33fc6cbe7cb0b476e6072006e68a7104d Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:39:31 -0300 Subject: [PATCH 04/20] =?UTF-8?q?feat(stark):=20LogUpMode::Gkr=20=E2=80=94?= =?UTF-8?q?=20lookup-layer=20adapter=20and=20single-source=20bridge=20cons?= =?UTF-8?q?traint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Opt-in per-AIR mode (default Standard, wire-identical). GKR mode: - aux trace = 2 columns (Lagrange kernel l, bridge running sum sigma) regardless of interaction count - one degree-2 bridge constraint sigma_next - sigma_curr - l*batched + delta, emitted through the generic ConstraintBuilder (all four interpretations: prover/verifier folders, IR capture, derived meta) with gamma powers and delta read as challenge leaves from the extended rap_challenges layout - boundary constraints l[0] = prod(1 - r_j) and sigma[0] = 0 - sigma is the sole next-row OOD read (preserves g*z pruning) - logup_gkr module: leaf fractions, layer trees, column-claim finalize, bridge params, aux-column build, and verifier-side claim reconstruction (sorted-slice lookups, no hashing on verify paths; the multi-interaction leaf-binding gap is documented fail-open pending the input-layer fix) - AIR trait: bus_interactions() and logup_mode() accessors Gates: bridge folder==IR parity (prover+verifier, 1000 random frames), GKR tree root == standard-mode table contribution (cross-mode oracle), aux-column recurrence/boundary tests; 275 stark tests green. --- crypto/stark/src/lib.rs | 1 + crypto/stark/src/logup_gkr.rs | 852 ++++++++++++++++++++++++++++++++++ crypto/stark/src/lookup.rs | 406 ++++++++++++++-- crypto/stark/src/traits.rs | 14 + 4 files changed, 1247 insertions(+), 26 deletions(-) create mode 100644 crypto/stark/src/logup_gkr.rs diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index c2bd817c4..0e5e2a99e 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -22,6 +22,7 @@ pub mod grinding; #[cfg(feature = "instruments")] pub mod instruments; pub mod lagrange_kernel; +pub mod logup_gkr; #[cfg(feature = "cuda")] pub mod logup_gpu; pub mod lookup; diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs new file mode 100644 index 000000000..321068995 --- /dev/null +++ b/crypto/stark/src/logup_gkr.rs @@ -0,0 +1,852 @@ +//! LogUp-GKR adapter: the glue between a table's bus interactions and the +//! batch GKR protocol in [`crate::gkr`]. +//! +//! In [`crate::lookup::LogUpMode::Gkr`] the per-table LogUp sums are proven by +//! a batch GKR over fraction-summation trees instead of committed term/acc +//! columns. Each interacting table commits exactly two auxiliary columns: +//! +//! - column [`GKR_AUX_KERNEL_COL`]: the Lagrange kernel `l[i] = eq(bits(i), r)` +//! at the GKR random point `r`, bound by the boundary constraint +//! `l[0] = ∏(1 − r_j)` and the `γ^K·l²` self-check folded into the bridge. +//! - column [`GKR_AUX_SIGMA_COL`]: the bridge running sum `σ`, whose circular +//! transition constraint telescopes to `⟨l, col_j⟩ = c_j` for every main +//! column `j` referenced by an interaction (Schwartz-Zippel over `γ`). +//! +//! This module holds the prover-side leaf/tree/claim computation, the shared +//! challenge-vector layout, and the verifier-side claim reconstruction. The +//! constraint *emission* for the bridge lives in [`crate::lookup`], beside the +//! standard LogUp emission, so both modes share the single-source discipline. + +use math::field::{ + element::FieldElement, + traits::{IsFFTField, IsField, IsPrimeField, IsSubFieldOf}, +}; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::gkr::{Layer, gen_layers}; +use crate::lagrange_kernel::{compute_lagrange_kernel, eval_mle_base_with_kernel}; +use crate::lookup::{BusInteraction, LOGUP_CHALLENGE_ALPHA, PackingShifts, compute_alpha_powers}; +use crate::trace::TraceTable; + +// ============================================================================= +// Challenge-vector layout (GKR mode) +// ============================================================================= +// The per-table rap_challenges vector in GKR mode is the shared `[z, α]` prefix +// extended with the bridge parameters: +// [0] = z, [1] = α (shared, sampled in Phase B) +// [2] = γ (shared, sampled after column claims) +// [3] = bridge offset Δ = target/N (derived, per table) +// [4 .. 4+K+1] = γ⁰, γ¹, …, γᴷ (derived; γᴷ backs the l² self-check) +// [4+K+1 ..] = r₀, …, r_{n−1} (GKR random point, per table length) +// K = number of distinct main columns referenced by the table's interactions. + +/// Index of the `z` challenge in the LogUp challenges vector. +pub const LOGUP_CHALLENGE_Z: usize = 0; + +/// Index of the `γ` challenge (column-claim batching) in the per-table +/// rap_challenges vector. Sampled on the main transcript after every table's +/// column claims are absorbed (binding them before γ exists). +pub const LOGUP_CHALLENGE_GAMMA: usize = 2; + +/// Index of the bridge offset `Δ = target/N` in the per-table rap_challenges +/// vector. A derived value, not a random challenge. +pub const LOGUP_BRIDGE_OFFSET_IDX: usize = 3; + +/// Start index of the precomputed γ powers in the per-table rap_challenges +/// vector: `rap_challenges[LOGUP_GAMMA_POWERS_START + j] = γ^j` for +/// `j = 0..=K`. The extra `γ^K` power backs the Lagrange-kernel `l²` +/// self-check. +pub const LOGUP_GAMMA_POWERS_START: usize = 4; + +/// Auxiliary column index of the Lagrange kernel `l`. +pub const GKR_AUX_KERNEL_COL: usize = 0; + +/// Auxiliary column index of the bridge running sum `σ`. +pub const GKR_AUX_SIGMA_COL: usize = 1; + +/// Number of auxiliary columns an interacting table commits in GKR mode. +pub const GKR_NUM_AUX_COLUMNS: usize = 2; + +/// Start index of the GKR random-point coordinates in the per-table +/// rap_challenges vector, given the table's distinct-column count `K`. +pub const fn logup_random_point_start(num_columns: usize) -> usize { + // +1 for the extra γ^K power (the l² self-check). + LOGUP_GAMMA_POWERS_START + num_columns + 1 +} + +// ============================================================================= +// Column extraction +// ============================================================================= + +/// Extract the sorted distinct main-column indices referenced by any +/// interaction (bus values and multiplicities). This canonical order is shared +/// by the prover's `column_claims`, the bridge constraint's batched sum, and +/// the verifier's claim checks. +pub fn extract_column_indices(interactions: &[BusInteraction]) -> Vec { + let mut cols = Vec::new(); + for inter in interactions { + for val in &inter.values { + cols.extend(val.column_indices()); + } + inter.multiplicity.collect_columns(&mut cols); + } + cols.sort_unstable(); + cols.dedup(); + cols +} + +// ============================================================================= +// Leaf fractions and layer trees (prover side) +// ============================================================================= + +/// The fingerprint `z − (bus_id + Σ αⁱ·elementᵢ)` for one interaction at one +/// row. `α⁰ = 1`: the bus-id term is embedded directly, no multiply. +fn compute_fingerprint_at_row( + interaction: &BusInteraction, + main_segment_cols: &[Vec>], + row: usize, + z: &FieldElement, + alpha_powers: &[FieldElement], + shifts: &PackingShifts, +) -> FieldElement +where + F: IsField + IsSubFieldOf, + E: IsField, +{ + let mut lc = FieldElement::::from(interaction.bus_id); + let mut alpha_offset = 1; + for bv in &interaction.values { + alpha_offset += bv.accumulate_fingerprint( + main_segment_cols, + row, + alpha_powers, + alpha_offset, + &mut lc, + shifts, + ); + } + z - &lc +} + +/// Computes the leaf fractions of the GKR summation tree from a table's bus +/// interactions and main trace. +/// +/// For each row `i`, all K interactions fold into a single fraction +/// `N(i)/D(i)` by cross-multiplication, starting from `0/1`: +/// +/// - `n' = n·fp_k ± m_k·d` (`+` for senders, `−` for receivers) +/// - `d' = d·fp_k` +/// +/// so `D(i) = Π_k fp_k(i)` and `N(i) = Σ_k ±m_k(i)·Π_{j≠k} fp_j(i)`, i.e. the +/// row's total LogUp contribution `Σ_k ±m_k/fp_k` as one fraction. +/// +/// Returns `(numerators, denominators)`, each of length `trace_len`. +pub fn compute_logup_leaf_fractions( + interactions: &[BusInteraction], + main_segment_cols: &[Vec>], + trace_len: usize, + challenges: &[FieldElement], +) -> (Vec>, Vec>) +where + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + E: IsField + Send + Sync, +{ + assert!( + !interactions.is_empty(), + "leaf fractions require at least one interaction" + ); + + let z = &challenges[LOGUP_CHALLENGE_Z]; + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let max_bus_elements = interactions + .iter() + .map(|inter| inter.num_bus_elements()) + .max() + .unwrap(); + let alpha_powers = compute_alpha_powers(alpha, max_bus_elements); + let shifts = PackingShifts::::new(); + + let leaf_at_row = |row: usize| { + let mut running_n = FieldElement::::zero(); + let mut running_d = FieldElement::::one(); + for inter in interactions { + let fp = compute_fingerprint_at_row( + inter, + main_segment_cols, + row, + z, + &alpha_powers, + &shifts, + ); + let m = inter.multiplicity.evaluate_at_row(main_segment_cols, row); + // m·d is F×E (base operand LEFT); the sign resolves as add vs neg. + let cross = &m * &running_d; + let cross = if inter.is_sender { cross } else { -cross }; + running_n = &running_n * &fp + cross; + running_d = &running_d * &fp; + } + (running_n, running_d) + }; + + #[cfg(feature = "parallel")] + let (numerators, denominators): (Vec<_>, Vec<_>) = + (0..trace_len).into_par_iter().map(leaf_at_row).unzip(); + #[cfg(not(feature = "parallel"))] + let (numerators, denominators): (Vec<_>, Vec<_>) = (0..trace_len).map(leaf_at_row).unzip(); + + (numerators, denominators) +} + +/// Compute the full GKR layer tree for one table's interactions: leaf +/// fractions, then pairwise fraction-summation layers up to the root. +pub fn compute_logup_layers( + interactions: &[BusInteraction], + main_segment_cols: &[Vec>], + trace_len: usize, + challenges: &[FieldElement], +) -> Vec> +where + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + E: IsField + Send + Sync, +{ + let (numerators, denominators) = + compute_logup_leaf_fractions(interactions, main_segment_cols, trace_len, challenges); + gen_layers(Layer::LogUpGeneric { + numerators, + denominators, + }) +} + +// ============================================================================= +// Per-table GKR result (prover side) +// ============================================================================= + +/// One table's outcome of the batch GKR run: the transcript-bound claims that +/// feed the bridge parameters and the proof. +#[derive(Debug, Clone)] +pub struct LogUpGkrResult { + /// The table's total LogUp contribution (`root_n / root_d`). + pub table_contribution: FieldElement, + /// The GKR random evaluation point for this instance + /// (length = log2(trace_len)). + pub random_point: Vec>, + /// Claimed MLE evaluation of the leaf numerator at `random_point`. + pub n_claim: FieldElement, + /// Claimed MLE evaluation of the leaf denominator at `random_point`. + pub d_claim: FieldElement, + /// MLE claim per distinct referenced main column, in + /// [`extract_column_indices`] order: `(column_index, ⟨l, col⟩)`. + pub column_claims: Vec<(usize, FieldElement)>, +} + +/// Finalize a table's GKR result: evaluate every referenced main column's MLE +/// at the instance random point (via one shared Lagrange kernel, dropped after +/// use — the aux-trace build recomputes it to keep this result small). +pub fn finalize_logup_gkr_result( + interactions: &[BusInteraction], + main_segment_cols: &[Vec>], + random_point: Vec>, + n_claim: FieldElement, + d_claim: FieldElement, + table_contribution: FieldElement, +) -> LogUpGkrResult +where + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + E: IsField + Send + Sync, +{ + let col_indices = extract_column_indices(interactions); + let kernel = compute_lagrange_kernel(&random_point); + + #[cfg(feature = "parallel")] + let col_iter = col_indices.into_par_iter(); + #[cfg(not(feature = "parallel"))] + let col_iter = col_indices.into_iter(); + + let column_claims: Vec<(usize, FieldElement)> = col_iter + .map(|col_idx| { + let claim = eval_mle_base_with_kernel(&main_segment_cols[col_idx], &kernel); + (col_idx, claim) + }) + .collect(); + + LogUpGkrResult { + table_contribution, + random_point, + n_claim, + d_claim, + column_claims, + } +} + +// ============================================================================= +// Bridge parameters (shared prover/verifier derivation) +// ============================================================================= + +/// The expected boundary value of the Lagrange kernel column: +/// `l[0] = eq(bits(0), r) = ∏_j (1 − r_j)`. +pub fn lagrange_l0(random_point: &[FieldElement]) -> FieldElement { + let one = FieldElement::::one(); + random_point + .iter() + .fold(one.clone(), |acc, r_j| acc * (&one - r_j)) +} + +/// The expected squared ℓ₂ norm of the true Lagrange kernel: +/// `Σ_i l[i]² = ∏_k (r_k² + (1 − r_k)²)`. Folding this into the bridge target +/// (via `γ^K`) forces the committed `l` column to be `eq(·, r)` — a forged +/// kernel would have to preserve this norm AND every `⟨l, col_j⟩` claim. +pub fn lagrange_kernel_norm_claim(random_point: &[FieldElement]) -> FieldElement { + let one = FieldElement::::one(); + random_point.iter().fold(one.clone(), |acc, r_k| { + let one_minus_r = &one - r_k; + acc * (r_k.square() + one_minus_r.square()) + }) +} + +/// Compute the bridge parameters from column claims: +/// `Δ = (Σ_j γ^j·c_j + γ^K·norm_claim) / N` and the γ powers `γ^0..γ^K`. +/// +/// Both prover and verifier derive these from transcript-bound values. +pub fn compute_bridge_params( + column_claims: &[(usize, FieldElement)], + gamma: &FieldElement, + trace_len: usize, + kernel_norm_claim: &FieldElement, +) -> (FieldElement, Vec>) { + let k = column_claims.len(); + let gamma_powers = compute_alpha_powers(gamma, k + 1); + + let mut target = FieldElement::::zero(); + for ((_, c_j), gp) in column_claims.iter().zip(gamma_powers.iter()) { + target += c_j * gp; + } + target += kernel_norm_claim * &gamma_powers[k]; + + let n_inv = FieldElement::::from(trace_len as u64) + .inv() + .expect("trace length is nonzero"); + let bridge_offset = &target * &n_inv; + + (bridge_offset, gamma_powers) +} + +/// Extend the shared `[z, α]` challenge prefix into a table's full GKR-mode +/// rap_challenges vector (see the module-top layout). +pub fn extend_rap_challenges_with_bridge( + rap_challenges: &mut Vec>, + column_claims: &[(usize, FieldElement)], + gamma: &FieldElement, + trace_len: usize, + random_point: &[FieldElement], +) { + debug_assert_eq!(rap_challenges.len(), LOGUP_CHALLENGE_GAMMA); + let norm_claim = lagrange_kernel_norm_claim(random_point); + let (bridge_offset, gamma_powers) = + compute_bridge_params(column_claims, gamma, trace_len, &norm_claim); + rap_challenges.push(gamma.clone()); + rap_challenges.push(bridge_offset); + rap_challenges.extend(gamma_powers); + rap_challenges.extend_from_slice(random_point); +} + +// ============================================================================= +// Auxiliary-trace build (GKR mode) +// ============================================================================= + +/// Fill the two GKR auxiliary columns (Lagrange kernel `l`, bridge running sum +/// `σ`) from the extended challenge vector. The columns must already be +/// allocated. `challenges` must be the full extended vector — the bridge +/// parameters and random point are read from their layout positions. +pub(crate) fn build_gkr_aux_columns( + trace: &mut TraceTable, + column_indices: &[usize], + challenges: &[FieldElement], +) where + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + E: IsField + Send + Sync, +{ + let trace_len = trace.num_rows(); + assert!( + trace_len.is_power_of_two(), + "GKR aux build requires a power-of-two trace length, got {trace_len}" + ); + let n_vars = trace_len.trailing_zeros() as usize; + let k = column_indices.len(); + let rp_start = logup_random_point_start(k); + assert_eq!( + challenges.len(), + rp_start + n_vars, + "GKR aux build requires the extended challenge vector \ + (got {} challenges, expected {} for K={k}, n_vars={n_vars})", + challenges.len(), + rp_start + n_vars, + ); + + let random_point = &challenges[rp_start..]; + let gamma_powers = &challenges[LOGUP_GAMMA_POWERS_START..LOGUP_GAMMA_POWERS_START + k + 1]; + let bridge_offset = &challenges[LOGUP_BRIDGE_OFFSET_IDX]; + + let kernel = compute_lagrange_kernel(random_point); + let main_segment_cols = trace.columns_main(); + + // batched[i] = Σ_j γ^j·col_j[i] + γ^K·l[i] — row-parallel, matches the + // bridge constraint's batched sum exactly. + let batched_at_row = |row: usize| { + let mut acc = &gamma_powers[k] * &kernel[row]; + for (j, &col_idx) in column_indices.iter().enumerate() { + // col·γ is F×E (base operand LEFT). + acc += &main_segment_cols[col_idx][row] * &gamma_powers[j]; + } + acc + }; + #[cfg(feature = "parallel")] + let batched: Vec> = + (0..trace_len).into_par_iter().map(batched_at_row).collect(); + #[cfg(not(feature = "parallel"))] + let batched: Vec> = (0..trace_len).map(batched_at_row).collect(); + + // σ[0] = 0; σ[i+1] = σ[i] + l[i]·batched[i] − Δ. The circular constraint + // σ_next − σ_curr − l·batched + Δ = 0 then holds on every row, including + // the wraparound (Σ l·batched = target = N·Δ closes the cycle). + let mut sigma = FieldElement::::zero(); + for row in 0..trace_len { + trace.set_aux(row, GKR_AUX_KERNEL_COL, kernel[row].clone()); + trace.set_aux(row, GKR_AUX_SIGMA_COL, sigma.clone()); + sigma = sigma + &kernel[row] * &batched[row] - bridge_offset; + } + debug_assert_eq!( + sigma, + FieldElement::::zero(), + "bridge running sum must wrap to zero (Σ l·batched = N·Δ)" + ); +} + +// ============================================================================= +// Verifier-side claim reconstruction +// ============================================================================= + +/// Sorted-slice column-claim lookup (guest-friendly: no hashing). +struct ClaimLookup<'a, E: IsField>(&'a [(usize, FieldElement)]); + +impl ClaimLookup<'_, E> { + /// The claimed MLE value for `col`. The caller has already checked the + /// claim index set equals [`extract_column_indices`], so every referenced + /// column is present; a miss returns zero (unreachable after that check). + fn get(&self, col: usize) -> FieldElement { + match self.0.binary_search_by_key(&col, |(c, _)| *c) { + Ok(i) => self.0[i].1.clone(), + Err(_) => FieldElement::::zero(), + } + } +} + +/// Verify a table's `column_claims` against the batch-GKR leaf claims +/// `(n_claim, d_claim)`. +/// +/// Always enforced: the claim index set must EQUAL the canonical sorted +/// distinct column set of the interactions (same indices, same order). +/// +/// For single-interaction tables and 0-layer (single-row) instances the leaf +/// fraction is reconstructible from the column MLEs — single-interaction +/// leaves are linear in the columns, and at the empty point the MLE is the row +/// value itself — so the rational cross-check +/// `n_recon·d_claim == n_claim·d_recon` is exact and enforced. +/// +/// # KNOWN SOUNDNESS GAP (multi-interaction tables) +/// +/// For `interactions.len() > 1` with `n_layers > 0` this check is FAIL-OPEN: +/// the leaf fraction is a nonlinear (cross-multiplied) function of the +/// columns, MLE does not commute with products, and nothing else binds +/// `(n_claim, d_claim)` to the committed trace — the bridge constraint binds +/// only `column_claims`. A malicious prover can therefore run an honest GKR +/// over fabricated leaves. Every production table is multi-interaction. The +/// fix (a linear input layer or an input-layer sumcheck) is designed as the +/// immediate follow-up — see `thoughts/logup-gkr/port-plan.md` §6. Do NOT +/// treat GKR mode as production-sound until it lands. +pub fn reconstruct_and_verify_gkr_claims( + n_claim: &FieldElement, + d_claim: &FieldElement, + column_claims: &[(usize, FieldElement)], + interactions: &[BusInteraction], + challenges: &[FieldElement], + n_layers: usize, +) -> bool { + // Structural binding: exact index-set (and order) equality with the + // canonical column list. Subsumes presence checks and pins the transcript + // absorption order of the claims. + let expected = extract_column_indices(interactions); + if column_claims.len() != expected.len() { + return false; + } + if !column_claims + .iter() + .zip(expected.iter()) + .all(|((c, _), e)| c == e) + { + return false; + } + + let claims = ClaimLookup(column_claims); + let z = &challenges[LOGUP_CHALLENGE_Z]; + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let max_bus_elements = interactions + .iter() + .map(|inter| inter.num_bus_elements()) + .max() + .unwrap_or(0); + let alpha_powers = compute_alpha_powers(alpha, max_bus_elements); + + // Reconstruct the leaf fraction from the column claims with the SAME + // cross-multiplication recurrence as `compute_logup_leaf_fractions`. + let mut running_n = FieldElement::::zero(); + let mut running_d = FieldElement::::one(); + for inter in interactions { + let mut lc = FieldElement::::from(inter.bus_id); + let mut alpha_offset = 1; + for bv in &inter.values { + for elem in bv.combine_from(|col| claims.get(col)) { + lc += &elem * &alpha_powers[alpha_offset]; + alpha_offset += 1; + } + } + let fp = z - &lc; + let m = inter.multiplicity.evaluate_from(|col| claims.get(col)); + let cross = &m * &running_d; + let cross = if inter.is_sender { cross } else { -cross }; + running_n = &running_n * &fp + cross; + running_d = &running_d * &fp; + } + + if interactions.len() == 1 || n_layers == 0 { + // Rational cross-check: n_recon/d_recon == n_claim/d_claim. + &running_n * d_claim == n_claim * &running_d + } else { + // FAIL-OPEN — see the doc comment's KNOWN SOUNDNESS GAP. + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::lookup::{Multiplicity, Packing}; + use math::field::goldilocks::GoldilocksField; + + type F = GoldilocksField; + type FE = FieldElement; + + /// Single sender, multiplicity 1: N(i) = 1, D(i) = fp(i). + #[test] + fn leaf_fractions_single_sender() { + let trace_len = 4; + let col0: Vec = [10u64, 20, 30, 40].iter().map(|&v| FE::from(v)).collect(); + let main_segment_cols = vec![col0.clone()]; + let interactions = vec![BusInteraction::sender( + 1u64, + Multiplicity::One, + Packing::Direct.columns(&[0]), + )]; + let z = FE::from(100u64); + let alpha = FE::from(3u64); + let challenges = vec![z, alpha]; + + let (numerators, denominators) = compute_logup_leaf_fractions::( + &interactions, + &main_segment_cols, + trace_len, + &challenges, + ); + + let alpha_powers = compute_alpha_powers(&alpha, 2); + for row in 0..trace_len { + let lc = FE::from(1u64) + col0[row] * alpha_powers[1]; + let expected_fp = z - lc; + assert_eq!(numerators[row], FE::one()); + assert_eq!(denominators[row], expected_fp); + } + } + + /// Single receiver with a column multiplicity: N(i) = −m(i), D(i) = fp(i). + #[test] + fn leaf_fractions_single_receiver_column_multiplicity() { + let trace_len = 4; + let col0: Vec = [5u64, 6, 7, 8].iter().map(|&v| FE::from(v)).collect(); + let col1: Vec = [2u64, 0, 1, 3].iter().map(|&v| FE::from(v)).collect(); + let main_segment_cols = vec![col0.clone(), col1.clone()]; + let interactions = vec![BusInteraction::receiver( + 0u64, + Multiplicity::Column(1), + Packing::Direct.columns(&[0]), + )]; + let z = FE::from(50u64); + let alpha = FE::from(7u64); + let challenges = vec![z, alpha]; + + let (numerators, denominators) = compute_logup_leaf_fractions::( + &interactions, + &main_segment_cols, + trace_len, + &challenges, + ); + + let alpha_powers = compute_alpha_powers(&alpha, 2); + for row in 0..trace_len { + let lc = col0[row] * alpha_powers[1]; + let expected_fp = z - lc; + assert_eq!(numerators[row], -col1[row]); + assert_eq!(denominators[row], expected_fp); + } + } + + /// Two interactions: n = fp₁ − fp₀ (sender then receiver, both m=1), + /// d = fp₀·fp₁. + #[test] + fn leaf_fractions_two_interactions_cross_multiply() { + let trace_len = 2; + let col0: Vec = [10u64, 20].iter().map(|&v| FE::from(v)).collect(); + let col1: Vec = [30u64, 40].iter().map(|&v| FE::from(v)).collect(); + let main_segment_cols = vec![col0.clone(), col1.clone()]; + let interactions = vec![ + BusInteraction::sender(0u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver(1u64, Multiplicity::One, Packing::Direct.columns(&[1])), + ]; + let z = FE::from(200u64); + let alpha = FE::from(5u64); + let challenges = vec![z, alpha]; + + let (numerators, denominators) = compute_logup_leaf_fractions::( + &interactions, + &main_segment_cols, + trace_len, + &challenges, + ); + + let alpha_powers = compute_alpha_powers(&alpha, 2); + for row in 0..trace_len { + let fp_0 = z - (col0[row] * alpha_powers[1]); + let fp_1 = z - (FE::from(1u64) + col1[row] * alpha_powers[1]); + assert_eq!(numerators[row], fp_1 - fp_0); + assert_eq!(denominators[row], fp_0 * fp_1); + } + } + + /// Column extraction covers bus values AND multiplicity columns, sorted + /// and deduplicated. + #[test] + fn extract_column_indices_covers_values_and_multiplicities() { + let interactions = vec![ + BusInteraction::sender(0u64, Multiplicity::Sum(7, 2), Packing::Word2L.columns(&[4])), + BusInteraction::receiver(1u64, Multiplicity::Column(2), Packing::Direct.columns(&[5])), + ]; + assert_eq!(extract_column_indices(&interactions), vec![2, 4, 5, 7]); + } + + /// The verifier-side reconstruction accepts honest single-interaction + /// claims and rejects a tampered one; the exact index-set check rejects + /// missing or extra claims. + #[test] + fn reconstruct_single_interaction_roundtrip_and_tamper() { + // Single-row table (n_vars = 0): the MLE at the empty point IS the row + // value, so honest claims are the row values themselves. + let col0 = FE::from(10u64); + let col1 = FE::from(3u64); + let interactions = vec![BusInteraction::sender( + 2u64, + Multiplicity::Column(1), + Packing::Direct.columns(&[0]), + )]; + let z = FE::from(1000u64); + let alpha = FE::from(11u64); + let challenges = vec![z, alpha]; + let alpha_powers = compute_alpha_powers(&alpha, 2); + + let fp = z - (FE::from(2u64) + col0 * alpha_powers[1]); + let n_claim = col1; + let d_claim = fp; + let column_claims = vec![(0usize, col0), (1usize, col1)]; + + assert!(reconstruct_and_verify_gkr_claims( + &n_claim, + &d_claim, + &column_claims, + &interactions, + &challenges, + 0, + )); + + let tampered = FE::from(999u64); + assert!(!reconstruct_and_verify_gkr_claims( + &tampered, + &d_claim, + &column_claims, + &interactions, + &challenges, + 0, + )); + + // Missing claim → reject. + assert!(!reconstruct_and_verify_gkr_claims( + &n_claim, + &d_claim, + &column_claims[..1], + &interactions, + &challenges, + 0, + )); + // Extra claim → reject. + let mut extra = column_claims.clone(); + extra.push((9, FE::from(1u64))); + assert!(!reconstruct_and_verify_gkr_claims( + &n_claim, + &d_claim, + &extra, + &interactions, + &challenges, + 0, + )); + } + + /// The cross-mode oracle: the GKR summation-tree ROOT must equal the + /// standard-mode table contribution `L = Σ_rows Σ_k ±m_k/fp_k` (computed + /// via the standard aux path's per-interaction term columns). This ties + /// leaf fractions AND `gen_layers` to the standard LogUp semantics — the + /// two modes must claim the same bus balance. + #[test] + fn gkr_root_matches_standard_table_contribution() { + let trace_len = 8usize; + let col0: Vec = (1..=8).map(|v| FE::from(v as u64)).collect(); + let col1: Vec = (21..=28).map(|v| FE::from(v as u64)).collect(); + let col2: Vec = [1u64, 0, 2, 1, 0, 3, 1, 1] + .iter() + .map(|&v| FE::from(v)) + .collect(); + let main_segment_cols = vec![col0, col1, col2]; + + let interactions = vec![ + BusInteraction::sender(3u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver(3u64, Multiplicity::Column(2), Packing::Direct.columns(&[1])), + BusInteraction::sender(5u64, Multiplicity::Column(2), Packing::Word2L.columns(&[0])), + ]; + let challenges = vec![FE::from(0xDEAD_BEEFu64), FE::from(0x1234_5678u64)]; + + // Standard-mode L: sum every interaction's term column. + let mut expected_l = FE::zero(); + for inter in &interactions { + let terms = crate::lookup::compute_logup_term_column::( + &[inter], + &main_segment_cols, + trace_len, + &challenges, + "oracle", + ); + for t in &terms { + expected_l = expected_l + t; + } + } + + // GKR: leaf fractions → summation tree → root claim n/d. + let layers = + compute_logup_layers::(&interactions, &main_segment_cols, trace_len, &challenges); + let (root_n, root_d) = match layers.last().expect("root layer") { + Layer::LogUpGeneric { + numerators, + denominators, + } => { + assert_eq!(numerators.len(), 1, "root layer has one element"); + (numerators[0], denominators[0]) + } + other => panic!("unexpected root layer variant: {other:?}"), + }; + let root_value = root_n * root_d.inv().expect("nonzero root denominator"); + assert_eq!(root_value, expected_l, "GKR root != standard-mode L"); + } + + /// Bridge parameters: Δ·N == Σ γʲ·cⱼ + γᴷ·norm_claim, and the kernel norm + /// claim matches the actual kernel's Σ l². + #[test] + fn bridge_params_and_kernel_norm() { + let r = vec![FE::from(3u64), FE::from(7u64)]; + let kernel = compute_lagrange_kernel(&r); + let norm: FE = kernel.iter().fold(FE::zero(), |acc, l| acc + l * l); + assert_eq!(norm, lagrange_kernel_norm_claim(&r)); + + let column_claims = vec![(0usize, FE::from(5u64)), (3usize, FE::from(8u64))]; + let gamma = FE::from(13u64); + let trace_len = 4usize; + let (delta, gamma_powers) = compute_bridge_params(&column_claims, &gamma, trace_len, &norm); + assert_eq!(gamma_powers.len(), 3); + let target = FE::from(5u64) * gamma_powers[0] + + FE::from(8u64) * gamma_powers[1] + + norm * gamma_powers[2]; + assert_eq!(delta * FE::from(trace_len as u64), target); + } + + /// The GKR aux columns satisfy the bridge recurrence and boundary values, + /// and honest column claims telescope: Σ l·batched = N·Δ. + #[test] + fn gkr_aux_columns_satisfy_bridge_recurrence() { + use crate::trace::TraceTable; + + let trace_len = 8usize; + let n_vars = 3usize; + let col0: Vec = (1..=8).map(|v| FE::from(v as u64)).collect(); + let col1: Vec = (11..=18).map(|v| FE::from(v as u64)).collect(); + let column_indices = vec![0usize, 1usize]; + let k = column_indices.len(); + + let mut trace = TraceTable::::from_columns_main(vec![col0.clone(), col1.clone()], 1); + trace.allocate_aux_table(GKR_NUM_AUX_COLUMNS); + + // Honest bridge parameters from honest column claims. + let random_point: Vec = [3u64, 7, 5].iter().map(|&v| FE::from(v)).collect(); + assert_eq!(random_point.len(), n_vars); + let kernel = compute_lagrange_kernel(&random_point); + let column_claims: Vec<(usize, FE)> = column_indices + .iter() + .map(|&c| { + let cols = [&col0, &col1]; + let claim = cols[c] + .iter() + .zip(kernel.iter()) + .fold(FE::zero(), |acc, (v, l)| acc + v * l); + (c, claim) + }) + .collect(); + let gamma = FE::from(17u64); + + let mut challenges = vec![FE::from(1u64), FE::from(2u64)]; // z, α (unused here) + extend_rap_challenges_with_bridge( + &mut challenges, + &column_claims, + &gamma, + trace_len, + &random_point, + ); + assert_eq!(challenges.len(), logup_random_point_start(k) + n_vars); + + build_gkr_aux_columns(&mut trace, &column_indices, &challenges); + + // Boundary values. + let l0 = trace.get_aux(0, GKR_AUX_KERNEL_COL); + assert_eq!(*l0, lagrange_l0(&random_point)); + let sigma0 = trace.get_aux(0, GKR_AUX_SIGMA_COL); + assert_eq!(*sigma0, FE::zero()); + + // The circular recurrence on every row (including wraparound). + let delta = &challenges[LOGUP_BRIDGE_OFFSET_IDX]; + let gp = &challenges[LOGUP_GAMMA_POWERS_START..LOGUP_GAMMA_POWERS_START + k + 1]; + for row in 0..trace_len { + let next = (row + 1) % trace_len; + let l = *trace.get_aux(row, GKR_AUX_KERNEL_COL); + let sigma_curr = *trace.get_aux(row, GKR_AUX_SIGMA_COL); + let sigma_next = *trace.get_aux(next, GKR_AUX_SIGMA_COL); + let batched = col0[row] * gp[0] + col1[row] * gp[1] + gp[k] * l; + assert_eq!( + sigma_next - sigma_curr - l * batched + delta, + FE::zero(), + "bridge recurrence failed at row {row}" + ); + } + } +} diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 8a89ea727..949eed1aa 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -104,6 +104,21 @@ pub const LOGUP_CHALLENGE_ALPHA: usize = 1; /// Number of challenges required by the LogUp protocol. pub const LOGUP_NUM_CHALLENGES: usize = 2; +/// How a table's LogUp bus sums are proven. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum LogUpMode { + /// Committed batched term columns + a circular accumulated column; the + /// per-table contribution `L` travels as a public input and the verifier + /// checks `Σ L == expected_bus_balance`. + #[default] + Standard, + /// Batch GKR over fraction-summation trees: two committed aux columns + /// (Lagrange kernel + bridge running sum) bind the GKR claims to the + /// trace; the bus balance is checked on the GKR root claims. See + /// [`crate::logup_gkr`]. + Gkr, +} + /// Chunk size for fused chunk-local LogUp processing. /// Each chunk processes all interactions for CHUNK_SIZE rows, fitting in L2 cache. #[cfg(feature = "parallel")] @@ -825,6 +840,13 @@ pub struct AirWithBuses< /// The LogUp layout: the framework generates the LogUp (extension) /// constraints from this and appends them after the `constraint_set` ones. logup: LogUpLayout, + /// How this table's LogUp sums are proven (default [`LogUpMode::Standard`]; + /// switch with [`Self::with_logup_mode`]). + mode: LogUpMode, + /// GKR mode only: the sorted distinct main columns referenced by the + /// interactions ([`crate::logup_gkr::extract_column_indices`]); empty in + /// standard mode. + gkr_column_indices: Vec, /// Idx-ordered metadata for all transition constraints, DERIVED at /// construction: `constraint_set.meta()` (base prefix) followed by the /// LogUp emission's derived metadata (ext). @@ -921,6 +943,8 @@ impl< trace_layout, constraint_set, logup, + mode: LogUpMode::Standard, + gkr_column_indices: Vec::new(), meta, num_base, constraint_program: std::sync::OnceLock::new(), @@ -933,6 +957,44 @@ impl< } } + /// Switch this table's LogUp mode, re-deriving every mode-dependent piece: + /// the auxiliary layout (⌈N/2⌉ term+acc columns vs the 2 GKR columns), the + /// appended LogUp constraint metadata, the context sizes, and the lazily + /// captured constraint program. + pub fn with_logup_mode(mut self, mode: LogUpMode) -> Self { + self.mode = mode; + let num_interactions = self.logup.interactions.len(); + self.gkr_column_indices = match mode { + LogUpMode::Standard => Vec::new(), + LogUpMode::Gkr => crate::logup_gkr::extract_column_indices(&self.logup.interactions), + }; + self.trace_layout.1 = if num_interactions == 0 { + 0 + } else { + match mode { + LogUpMode::Standard => self.logup.num_term_columns + 1, + LogUpMode::Gkr => crate::logup_gkr::GKR_NUM_AUX_COLUMNS, + } + }; + + let mut meta = self.constraint_set.meta(); + let mut logup_mb = crate::constraints::builder::MetaBuilder::new(); + emit_logup_for_mode::( + &mut logup_mb, + mode, + &self.logup, + &self.gkr_column_indices, + self.num_base, + ); + meta.extend(logup_mb.into_meta()); + self.meta = meta; + + self.context.trace_columns = self.trace_layout.0 + self.trace_layout.1; + self.context.num_transition_constraints = self.meta.len(); + self.constraint_program = std::sync::OnceLock::new(); + self + } + /// Marks this AIR as a preprocessed table with a hardcoded commitment. /// /// Preprocessed tables have columns that are fully deterministic and known @@ -1004,15 +1066,20 @@ where } fn trace_ood_next_row_columns(&self) -> Vec { - // The only transition constraint that reads the next row is the circular - // LogUp accumulator, and after forward accumulation it reads only the - // accumulated column there (all committed terms and absorbed operands - // read the current row). Its full-width index is the main width plus the - // accumulated column's aux index. No interactions => no next-row reads. + // The only transition constraint that reads the next row is the + // circular running-sum column — the LogUp accumulator in standard mode + // (forward accumulation reads only `acc` there), the bridge sum σ in + // GKR mode (the kernel and every main column read the current row). + // Its full-width index is the main width plus its aux index. No + // interactions => no next-row reads. if self.auxiliary_trace_build_data.interactions.is_empty() { Vec::new() } else { - vec![self.trace_layout.0 + self.logup.acc_column_idx] + let aux_idx = match self.mode { + LogUpMode::Standard => self.logup.acc_column_idx, + LogUpMode::Gkr => crate::logup_gkr::GKR_AUX_SIGMA_COL, + }; + vec![self.trace_layout.0 + aux_idx] } } @@ -1020,6 +1087,14 @@ where !self.auxiliary_trace_build_data.interactions.is_empty() } + fn bus_interactions(&self) -> &[BusInteraction] { + &self.auxiliary_trace_build_data.interactions + } + + fn logup_mode(&self) -> LogUpMode { + self.mode + } + fn max_bus_elements(&self) -> usize { self.max_bus_elements } @@ -1029,10 +1104,17 @@ where // once via `ConstraintSet::max_degree()`; the framework's LogUp // constraints contribute their own known max (batched terms degree 3, // accumulator `1 + absorbed`). - let max_degree = self - .constraint_set - .max_degree() - .max(logup_max_degree(&self.logup)); + let logup_degree = if self.auxiliary_trace_build_data.interactions.is_empty() { + 0 + } else { + match self.mode { + LogUpMode::Standard => logup_max_degree(&self.logup), + // The bridge constraint `σ_next − σ_curr − l·batched + Δ` is + // degree 2 (`l` × trace columns). + LogUpMode::Gkr => 2, + } + }; + let max_degree = self.constraint_set.max_degree().max(logup_degree); // The composition polynomial is the constraint QUOTIENT H = Σ βᵢ·Cᵢ/Zᵢ. Its degree is // deg(Cᵢ) − deg(Zᵢ) = (max_degree−1)·N − max_degree + eᵢ, so with the end-exemptions // eᵢ < max_degree (the max-degree LogUp constraints have eᵢ = 0) it fits in @@ -1065,7 +1147,9 @@ where // by the base-constraint count). run_air_transition_prover( &self.constraint_set, + self.mode, &self.logup, + &self.gkr_column_indices, ctx, base_evals, ext_evals, @@ -1078,7 +1162,9 @@ where ) -> Vec> { run_air_transition_verifier( &self.constraint_set, + self.mode, &self.logup, + &self.gkr_column_indices, self.num_base, self.meta.len(), ctx, @@ -1094,7 +1180,13 @@ where self.constraint_program.get_or_init(|| { let mut cb = crate::constraints::builder::CaptureBuilder::::new(); self.constraint_set.eval(&mut cb); - emit_logup_constraints(&mut cb, &self.logup, self.num_base); + emit_logup_for_mode( + &mut cb, + self.mode, + &self.logup, + &self.gkr_column_indices, + self.num_base, + ); let (prog, _degrees) = cb.finish(self.num_base); prog }) @@ -1117,6 +1209,15 @@ where return None; } + // GKR mode: fill the Lagrange-kernel and bridge-sum columns from the + // extended challenge vector (the batch GKR has already run — the + // bridge parameters and random point are in `challenges`). No bus + // public inputs travel: the balance check runs on the GKR root claims. + if self.mode == LogUpMode::Gkr { + crate::logup_gkr::build_gkr_aux_columns(trace, &self.gkr_column_indices, challenges); + return None; + } + // Clone main columns once (shared across all interactions) let main_segment_cols = trace.columns_main(); let trace_len = trace.num_rows(); @@ -1283,15 +1384,44 @@ where ) -> BoundaryConstraints { let mut boundary_constraints = B::boundary_constraints(pub_inputs, rap_challenges); - // Pin acc[0] = 0 to remove the constant-shift degree of freedom in the - // circular transition constraint (forward accumulation starts at 0). if !self.auxiliary_trace_build_data.interactions.is_empty() { - let acc_col_idx = self.trace_layout.1 - 1; // last aux column = accumulated - boundary_constraints.push(BoundaryConstraint::new_aux( - acc_col_idx, - 0, - FieldElement::zero(), - )); + match self.mode { + // Pin acc[0] = 0 to remove the constant-shift degree of freedom + // in the circular transition constraint (forward accumulation + // starts at 0). + LogUpMode::Standard => { + let acc_col_idx = self.trace_layout.1 - 1; // last aux column = accumulated + boundary_constraints.push(BoundaryConstraint::new_aux( + acc_col_idx, + 0, + FieldElement::zero(), + )); + } + // Bind the Lagrange kernel column to the GKR random point via + // `l[0] = ∏(1 − r_j)` (its Σl² norm is checked by the bridge's + // γ^K term), and pin σ[0] = 0 as in standard mode. + LogUpMode::Gkr => { + let rp_start = + crate::logup_gkr::logup_random_point_start(self.gkr_column_indices.len()); + assert!( + rap_challenges.len() > rp_start, + "GKR-mode boundary constraints require the extended \ + challenge vector (got {} challenges, random point starts at {rp_start})", + rap_challenges.len(), + ); + let random_point = &rap_challenges[rp_start..]; + boundary_constraints.push(BoundaryConstraint::new_aux( + crate::logup_gkr::GKR_AUX_KERNEL_COL, + 0, + crate::logup_gkr::lagrange_l0(random_point), + )); + boundary_constraints.push(BoundaryConstraint::new_aux( + crate::logup_gkr::GKR_AUX_SIGMA_COL, + 0, + FieldElement::zero(), + )); + } + } } BoundaryConstraints::from_constraints(boundary_constraints) @@ -1407,6 +1537,36 @@ impl Multiplicity { ) -> FieldElement { self.evaluate_with(|col| main_segment_cols[col][row].clone()) } + + /// Evaluate the multiplicity from an arbitrary column-value source (e.g. + /// the GKR column claims — MLE evaluations instead of trace rows). + #[inline] + pub fn evaluate_from(&self, get_col: G) -> FieldElement + where + F: IsField, + G: Fn(usize) -> FieldElement, + { + self.evaluate_with(get_col) + } + + /// Append every main-column index this multiplicity reads to `out`. + pub(crate) fn collect_columns(&self, out: &mut Vec) { + match self { + Multiplicity::One => {} + Multiplicity::Column(c) | Multiplicity::Negated(c) => out.push(*c), + Multiplicity::Sum(a, b) | Multiplicity::Diff(a, b) => out.extend([*a, *b]), + Multiplicity::Sum3(a, b, c) => out.extend([*a, *b, *c]), + Multiplicity::Linear(terms) => { + for term in terms { + match term { + LinearTerm::Column { column, .. } + | LinearTerm::ColumnUnsigned { column, .. } => out.push(*column), + LinearTerm::Constant(_) => {} + } + } + } + } + } } /// Struct representing a lookup interaction for a given table. @@ -1607,7 +1767,7 @@ where /// /// With `parallel`: chunked over rows via `par_chunks_mut`. /// Without `parallel`: processed as a single chunk. -fn compute_logup_term_column( +pub(crate) fn compute_logup_term_column( interactions: &[&BusInteraction], main_segment_cols: &[Vec>], trace_len: usize, @@ -2270,6 +2430,70 @@ where emit_logup_accumulated::(b, layout, idx); } +/// Emit the GKR-mode bridge running-sum constraint (the ONLY LogUp transition +/// constraint in [`LogUpMode::Gkr`]): +/// +/// `σ_next − σ_curr − l_curr·batched_curr + Δ = 0` (degree 2, every row) +/// +/// with `batched = γ^K·l + Σ_j col_j·γ^j` over the table's distinct referenced +/// main columns (canonical [`crate::logup_gkr::extract_column_indices`] order) +/// and `Δ = target/N` from the challenge vector. Telescoped over the circular +/// domain this proves `⟨l, col_j⟩ = c_j` for every claim (Schwartz-Zippel over +/// γ) plus the kernel's `Σ l²` norm (the `γ^K` term). `σ_next` is the sole +/// next-row read; `l` and every main column read the current row, preserving +/// the 1-column pruned OOD next-row block. +/// +/// γ powers and Δ are read as [`ConstraintBuilder::challenge`] leaves from the +/// extended rap_challenges vector (layout in [`crate::logup_gkr`]) — no new +/// context plumbing. +pub fn emit_logup_gkr_constraints(b: &mut B, column_indices: &[usize], idx_base: usize) +where + F: IsField, + E: IsField, + B: ConstraintBuilder, +{ + use crate::logup_gkr::{ + GKR_AUX_KERNEL_COL, GKR_AUX_SIGMA_COL, LOGUP_BRIDGE_OFFSET_IDX, LOGUP_GAMMA_POWERS_START, + }; + + let k = column_indices.len(); + let l = b.aux(0, GKR_AUX_KERNEL_COL); + let sigma_curr = b.aux(0, GKR_AUX_SIGMA_COL); + let sigma_next = b.aux(1, GKR_AUX_SIGMA_COL); + + // γ^K·l first (an ExprE seed even for K = 0), then the column terms; + // col·γ is base×ext (base operand LEFT). + let mut batched = b.challenge(LOGUP_GAMMA_POWERS_START + k) * l.clone(); + for (j, &col_idx) in column_indices.iter().enumerate() { + batched = batched + b.main(0, col_idx) * b.challenge(LOGUP_GAMMA_POWERS_START + j); + } + let delta = b.challenge(LOGUP_BRIDGE_OFFSET_IDX); + b.emit_ext(idx_base, sigma_next - sigma_curr - l * batched + delta); +} + +/// Emit the LogUp transition constraints for the given mode: the standard +/// batched-term/accumulated set, or the single GKR bridge constraint. Emits +/// nothing when there are no interactions. +fn emit_logup_for_mode( + b: &mut B, + mode: LogUpMode, + layout: &LogUpLayout, + gkr_column_indices: &[usize], + idx_base: usize, +) where + F: IsField, + E: IsField, + B: ConstraintBuilder, +{ + if layout.interactions.is_empty() { + return; + } + match mode { + LogUpMode::Standard => emit_logup_constraints(b, layout, idx_base), + LogUpMode::Gkr => emit_logup_gkr_constraints(b, gkr_column_indices, idx_base), + } +} + /// Run an [`AirWithBuses`] table's transition constraints through the /// [`ProverEvalFolder`] in ONE pass: the constraint set's base-field body /// followed by the appended LogUp constraints (idx offset by `num_base`, the @@ -2277,7 +2501,9 @@ where /// constraint count. fn run_air_transition_prover( constraint_set: &CS, + mode: LogUpMode, logup: &LogUpLayout, + gkr_column_indices: &[usize], ctx: &TransitionEvaluationContext<'_, F, E>, base_evals: &mut [FieldElement], ext_evals: &mut [FieldElement], @@ -2289,7 +2515,7 @@ fn run_air_transition_prover( let num_base = base_evals.len(); let mut folder = ProverEvalFolder::new(ctx, base_evals, ext_evals); constraint_set.eval(&mut folder); - emit_logup_constraints(&mut folder, logup, num_base); + emit_logup_for_mode(&mut folder, mode, logup, gkr_column_indices, num_base); folder.assert_all_emitted(); } @@ -2303,7 +2529,9 @@ fn run_air_transition_prover( /// base-prefix results. fn run_air_transition_verifier( constraint_set: &CS, + mode: LogUpMode, logup: &LogUpLayout, + gkr_column_indices: &[usize], num_base: usize, num_constraints: usize, ctx: &TransitionEvaluationContext<'_, F, E>, @@ -2318,14 +2546,14 @@ where TransitionEvaluationContext::Verifier { .. } => { let mut folder = VerifierEvalFolder::new(ctx, &mut ext_evals); constraint_set.eval(&mut folder); - emit_logup_constraints(&mut folder, logup, num_base); + emit_logup_for_mode(&mut folder, mode, logup, gkr_column_indices, num_base); folder.assert_all_emitted(); } TransitionEvaluationContext::Prover { .. } => { let mut base_evals = vec![FieldElement::::zero(); num_base]; let mut folder = ProverEvalFolder::new(ctx, &mut base_evals, &mut ext_evals); constraint_set.eval(&mut folder); - emit_logup_constraints(&mut folder, logup, num_base); + emit_logup_for_mode(&mut folder, mode, logup, gkr_column_indices, num_base); folder.assert_all_emitted(); // Promote the base-prefix results into the extension slots. for (slot, base) in ext_evals.iter_mut().zip(base_evals) { @@ -2433,12 +2661,12 @@ mod logup_single_source_tests { for i in 0..n_rows { let mut row_sum = Fp3::zero(); for col in &term_columns { - row_sum = row_sum + &col[i]; + row_sum += col[i]; } let acc_i = *trace.get_aux(i, acc_col_idx); let acc_next = *trace.get_aux((i + 1) % n_rows, acc_col_idx); - let lhs = (acc_next - acc_i) * &n_fe; - let rhs = row_sum * &n_fe - &l; + let lhs = (acc_next - acc_i) * n_fe; + let rhs = row_sum * n_fe - l; assert_eq!(lhs, rhs, "forward circular recurrence broken at row {i}"); } } @@ -2708,6 +2936,132 @@ mod logup_single_source_tests { check_layout("two_committed_pairs", &layout, 8); } + /// The GKR bridge constraint run three ways from ONE definition + /// ([`emit_logup_gkr_constraints`]): derived meta is a single ext + /// constraint, capture measures degree 2, and on random frames + /// [`ProverEvalFolder`] == capture→[`eval_program`] == + /// [`VerifierEvalFolder`] == capture→[`eval_program_verifier`] bit-for-bit + /// (the GPU-contract parity the standard emission gets from + /// [`check_layout`]). + fn check_gkr_bridge(label: &str, interactions: Vec, num_main_cols: usize) { + use crate::logup_gkr::{ + GKR_NUM_AUX_COLUMNS, extract_column_indices, logup_random_point_start, + }; + + let column_indices = extract_column_indices(&interactions); + let k = column_indices.len(); + let n_base = 0usize; + let n = 1usize; // the bridge constraint alone + + let meta = { + let mut mb = crate::constraints::builder::MetaBuilder::new(); + emit_logup_gkr_constraints::(&mut mb, &column_indices, n_base); + mb.into_meta() + }; + assert_eq!(meta.len(), n, "[{label}] one bridge constraint"); + assert_eq!(meta[0].kind, RootKind::Ext, "[{label}] bridge is ext"); + assert_eq!(meta[0].end_exemptions, 0, "[{label}] bridge is circular"); + + let mut cb = CaptureBuilder::::new(); + emit_logup_gkr_constraints(&mut cb, &column_indices, n_base); + let (prog, degrees) = cb.finish(0); + assert_eq!(degrees.len(), n, "[{label}] one emit"); + assert_eq!(degrees[0], (0, 2), "[{label}] bridge measures degree 2"); + + // Random rap_challenges covering the full extended layout (γ powers, + // Δ, and a fake 3-var random point — parity only reads the leaves). + let rp_len = 3usize; + let n_challenges = logup_random_point_start(k) + rp_len; + + for trial in 0..TRIALS { + let mut rng = SplitMix64::new(0x6B12_0000_u64 ^ (label.len() as u64) ^ trial as u64); + let mk_step = |rng: &mut SplitMix64| { + let main: Vec = (0..num_main_cols) + .map(|_| Fp::from(rng.next_u64())) + .collect(); + let aux: Vec = (0..GKR_NUM_AUX_COLUMNS).map(|_| rand_fp3(rng)).collect(); + TableView::new(vec![main], vec![aux]) + }; + let frame = Frame::::new(vec![mk_step(&mut rng), mk_step(&mut rng)]); + let rap_challenges: Vec = (0..n_challenges).map(|_| rand_fp3(&mut rng)).collect(); + let alpha_powers: Vec = (0..4).map(|_| rand_fp3(&mut rng)).collect(); + let table_offset = rand_fp3(&mut rng); + + let prover_ctx = TransitionEvaluationContext::new_prover( + frame.as_row_frame(), + &rap_challenges, + &alpha_powers, + &table_offset, + ); + + let mut base_out = vec![Fp::zero(); n_base]; + let mut ext_out = vec![Fp3::zero(); n]; + let mut folder = ProverEvalFolder::new(&prover_ctx, &mut base_out, &mut ext_out); + emit_logup_gkr_constraints(&mut folder, &column_indices, n_base); + folder.assert_all_emitted(); + + let mut ir_base = vec![Fp::zero(); n_base]; + let mut ir_ext = vec![Fp3::zero(); n]; + eval_program(&prog, &prover_ctx, &mut ir_base, &mut ir_ext); + assert_eq!( + ext_out[0], ir_ext[0], + "[{label}] prover folder vs interpreter mismatch, trial {trial}" + ); + + let embed_step = |step: &TableView| -> TableView { + let main: Vec = (0..num_main_cols) + .map(|c| step.get_main_evaluation_element(0, c).to_extension()) + .collect(); + let aux: Vec = (0..GKR_NUM_AUX_COLUMNS) + .map(|c| *step.get_aux_evaluation_element(0, c)) + .collect(); + TableView::new(vec![main], vec![aux]) + }; + let vframe: Frame = Frame::new(vec![ + embed_step(frame.get_evaluation_step(0)), + embed_step(frame.get_evaluation_step(1)), + ]); + let vctx = TransitionEvaluationContext::::new_verifier( + &vframe, + &rap_challenges, + &alpha_powers, + &table_offset, + ); + + let mut vext_out = vec![Fp3::zero(); n]; + let mut vfolder = VerifierEvalFolder::new(&vctx, &mut vext_out); + emit_logup_gkr_constraints(&mut vfolder, &column_indices, n_base); + vfolder.assert_all_emitted(); + + let mut ir_vext = vec![Fp3::zero(); n]; + eval_program_verifier(&prog, &vctx, &mut ir_vext); + assert_eq!( + vext_out[0], ir_vext[0], + "[{label}] verifier folder vs interpreter mismatch, trial {trial}" + ); + assert_eq!( + ext_out[0], vext_out[0], + "[{label}] prover vs verifier folder mismatch, trial {trial}" + ); + } + } + + #[test] + fn logup_gkr_bridge_single_interaction() { + check_gkr_bridge("gkr_single", vec![direct_sender(7)], 4); + } + + #[test] + fn logup_gkr_bridge_multi_interaction() { + // Shared and distinct columns across interactions + a multiplicity + // column: exercises dedup and the multi-term batched sum. + check_gkr_bridge( + "gkr_multi", + vec![direct_sender(7), column_receiver(11), direct_sender(13)], + 6, + ); + } + #[test] fn logup_linear_zero_skip() { // The prover folder zero-skips the F×E multiply for Linear bus diff --git a/crypto/stark/src/traits.rs b/crypto/stark/src/traits.rs index 0aec97a2a..225c2f52d 100644 --- a/crypto/stark/src/traits.rs +++ b/crypto/stark/src/traits.rs @@ -171,6 +171,20 @@ pub trait AIR: Send + Sync { 0 } + /// The table's bus interactions (empty when the AIR has none). The GKR + /// LogUp path reads these to build leaf fractions (prover) and reconstruct + /// claims (verifier). + fn bus_interactions(&self) -> &[crate::lookup::BusInteraction] { + &[] + } + + /// How this AIR's LogUp sums are proven. Non-bus AIRs report the default + /// [`crate::lookup::LogUpMode::Standard`]; the prover and verifier require + /// every interacting table in a multi-proof to agree on the mode. + fn logup_mode(&self) -> crate::lookup::LogUpMode { + crate::lookup::LogUpMode::Standard + } + /// Returns true if this AIR has preprocessed (precomputed) columns. /// /// Preprocessed tables have columns that are fully deterministic and known From c249c1646a24a7855f78e2aefe0a30dec9b88558 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:43:52 -0300 Subject: [PATCH 05/20] =?UTF-8?q?feat(stark):=20GKR=20wire=20format=20?= =?UTF-8?q?=E2=80=94=20rkyv=20proof=20types=20and=20GkrMultiProof=20wrappe?= =?UTF-8?q?r?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - rkyv Archive/Serialize/Deserialize on RoundPoly, SumcheckProof, BatchGkrLayerProof, BatchGkrProof (rkyv is the authoritative format) - GkrMultiProof = MultiProof + batch GKR proof + per-table column claims, a SEPARATE top-level type so standard-mode proofs stay byte-identical (cross-version verification remains meaningful with GKR off) - untrusted-proof shape guards in gkr_verify_batch: round polynomials must have exactly 4 evals (deserialization bypasses RoundPoly::new's assert), max_layers bounded before the 2^n_unused shift - tests: rkyv + cbor roundtrip verifies; malformed round poly and truncated layer list rejected as errors, not panics --- crypto/stark/src/gkr.rs | 93 ++++++++++++++++++++++++++++++++- crypto/stark/src/proof/stark.rs | 35 ++++++++++++- crypto/stark/src/sumcheck.rs | 25 ++++++++- 3 files changed, 148 insertions(+), 5 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index 3eda95df0..0838630ed 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -141,7 +141,15 @@ pub fn gen_layers(input_layer: Layer) -> Vec> { /// Proof for a single layer in a batch GKR reduction. /// /// Contains the shared sumcheck proof and per-instance child claims (masks). -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] #[serde(bound = "")] pub struct BatchGkrLayerProof { /// Shared sumcheck proof for this layer (combined across all active instances). @@ -154,7 +162,15 @@ pub struct BatchGkrLayerProof { /// Complete batch GKR proof for multiple fractional summation trees. /// /// All instances share one sumcheck per layer via random linear combination. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] #[serde(bound = "")] pub struct BatchGkrProof { /// Per-instance root claims as (numerator, denominator) pairs. @@ -1936,6 +1952,14 @@ pub fn gkr_verify_batch( let max_layers = *n_layers_by_instance.iter().max().unwrap(); + // The 2^n_unused doubling factor below is computed as a u64 shift; layer + // counts are log2(trace_length) so this bound is never near in practice. + if max_layers >= 64 { + return Err(GkrError::InvalidTree { + reason: format!("max_layers {max_layers} out of range"), + }); + } + if proof.layer_proofs.len() != max_layers { return Err(GkrError::InvalidTree { reason: format!( @@ -1946,6 +1970,22 @@ pub fn gkr_verify_batch( }); } + // Deserialization bypasses `RoundPoly::new`; reject wrong-shaped round + // polynomials before any `sum_at_binary`/`evaluate` call can assert. The + // batch gate is degree 3, so every round polynomial has exactly 4 evals. + for (layer_idx, layer_proof) in proof.layer_proofs.iter().enumerate() { + for rp in &layer_proof.sumcheck_proof.round_polys { + if rp.num_evals() != 4 { + return Err(GkrError::InvalidTree { + reason: format!( + "layer {layer_idx}: round polynomial has {} evals, expected 4", + rp.num_evals(), + ), + }); + } + } + } + // Track per-instance state let mut n_claims: Vec>> = vec![None; n_instances]; let mut d_claims: Vec>> = vec![None; n_instances]; @@ -3062,6 +3102,55 @@ mod tests { assert_eq!(v_claims, final_claims); } + /// The wire format: a batch proof survives an rkyv roundtrip and still + /// verifies; a malformed round polynomial smuggled in through + /// deserialization (bypassing `RoundPoly::new`) is REJECTED, not a panic. + #[test] + fn test_batch_gkr_rkyv_roundtrip_and_malformed_round_poly_rejected() { + let instances: Vec>> = + (0..2).map(|_| gen_layers(make_generic_leaf(2))).collect(); + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (proof, _, _) = gkr_prove_batch(instances, &mut prover_transcript); + let n_layers = vec![2usize, 2]; + + // rkyv roundtrip → verifies. + let bytes = rkyv::to_bytes::(&proof).unwrap(); + let deserialized: BatchGkrProof = + rkyv::from_bytes::<_, rkyv::rancor::Error>(&bytes).unwrap(); + let mut vt = DefaultTranscript::::new(&[]); + assert!( + gkr_verify_batch(&deserialized, &n_layers, &mut vt).is_ok(), + "rkyv-roundtripped batch proof must verify" + ); + + // serde roundtrip → verifies (examples-CLI compatibility path). + let cbor = serde_cbor::to_vec(&proof).unwrap(); + let from_cbor: BatchGkrProof = serde_cbor::from_slice(&cbor).unwrap(); + let mut vt = DefaultTranscript::::new(&[]); + assert!(gkr_verify_batch(&from_cbor, &n_layers, &mut vt).is_ok()); + + // A wrong-shaped round polynomial (2 evals instead of 4) must be + // rejected with an error before any evaluation helper can assert. + let mut malformed = proof.clone(); + let layer = malformed + .layer_proofs + .iter_mut() + .find(|lp| !lp.sumcheck_proof.round_polys.is_empty()) + .expect("some layer has sumcheck rounds"); + layer.sumcheck_proof.round_polys[0] = RoundPoly::new(vec![FE::from(1u64), FE::from(2u64)]); + let mut vt = DefaultTranscript::::new(&[]); + assert!( + gkr_verify_batch(&malformed, &n_layers, &mut vt).is_err(), + "wrong-shaped round polynomial must be rejected" + ); + + // Wrong layer count must be rejected. + let mut truncated = proof.clone(); + truncated.layer_proofs.pop(); + let mut vt = DefaultTranscript::::new(&[]); + assert!(gkr_verify_batch(&truncated, &n_layers, &mut vt).is_err()); + } + #[test] fn test_batch_gkr_mixed_size_instances() { // This is the key test: 3 instances with DIFFERENT sizes. diff --git a/crypto/stark/src/proof/stark.rs b/crypto/stark/src/proof/stark.rs index ba4aca2dc..4edc5afda 100644 --- a/crypto/stark/src/proof/stark.rs +++ b/crypto/stark/src/proof/stark.rs @@ -5,7 +5,8 @@ use math::field::{ }; use crate::{ - config::Commitment, fri::fri_decommit::FriDecommitment, lookup::BusPublicInputs, table::Table, + config::Commitment, fri::fri_decommit::FriDecommitment, gkr::BatchGkrProof, + lookup::BusPublicInputs, table::Table, }; // The proof types below intentionally derive both serde and rkyv. rkyv is the @@ -127,3 +128,35 @@ pub struct StarkProof, E: IsField, PI> { pub struct MultiProof, E: IsField, PI> { pub proofs: Vec>, } + +/// A multi-table proof under [`crate::lookup::LogUpMode::Gkr`]: the per-table +/// STARK proofs plus the LogUp-GKR artifacts. +/// +/// This is a SEPARATE top-level wire type on purpose: standard-mode proofs +/// keep the exact [`MultiProof`] rkyv/serde layout (byte-identical to a +/// GKR-unaware build), and GKR mode is purely additive on the wire. Everything +/// else GKR-related is transcript-derived by the verifier (random points, +/// instance claims) or recomputed (bridge parameters) — only the batch proof +/// and the per-table column claims travel. +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "PI: serde::Serialize + serde::de::DeserializeOwned")] +pub struct GkrMultiProof, E: IsField, PI> { + /// The per-table STARK proofs (same layout as standard mode; their + /// `bus_public_inputs` are `None` — the balance check runs on the GKR + /// root claims instead). + pub multi: MultiProof, + /// The batch GKR proof across every interacting table's summation tree. + pub batch_gkr_proof: BatchGkrProof, + /// Per-table column claims, aligned with `multi.proofs`: `None` for + /// non-interacting tables, else `(column_index, ⟨l, col⟩ MLE claim)` in + /// canonical [`crate::logup_gkr::extract_column_indices`] order. + pub column_claims_by_table: Vec)>>>, +} diff --git a/crypto/stark/src/sumcheck.rs b/crypto/stark/src/sumcheck.rs index 1510a5a5a..f5ce6f84d 100644 --- a/crypto/stark/src/sumcheck.rs +++ b/crypto/stark/src/sumcheck.rs @@ -11,10 +11,23 @@ use math::field::{element::FieldElement, traits::IsField}; /// - The prover constructs the polynomial by evaluating it at small integer points. /// - The verifier only needs to check p(0) + p(1) and evaluate at a random challenge. /// - Lagrange interpolation over integer nodes is cheap (small denominators). -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] #[serde(bound = "")] pub struct RoundPoly { /// Evaluations at x = 0, 1, ..., d where d = evals.len() - 1. + /// + /// Deserialization bypasses [`RoundPoly::new`]'s non-empty assert, so + /// consumers of untrusted proofs MUST length-check via + /// [`RoundPoly::num_evals`] before calling [`RoundPoly::sum_at_binary`] or + /// [`RoundPoly::evaluate`] (the batch GKR verifier rejects wrong shapes). evals: Vec>, } @@ -124,7 +137,15 @@ impl RoundPoly { } /// Proof produced by the sumcheck prover: one round polynomial per variable. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] #[serde(bound = "")] pub struct SumcheckProof { pub round_polys: Vec>, From 6c53036c73d609a4545e965561af22f54753d0c3 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 12:58:19 -0300 Subject: [PATCH 06/20] feat(stark): wire LogUp-GKR into multi_prove / multi_verify MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prover (multi_prove_gkr -> GkrMultiProof): after the shared z/alpha sampling, a new Phase B'/B'' on the SHARED transcript — parallel leaf fractions + summation trees per interacting table, one gkr_prove_batch, parallel per-table column-claim finalize, claims absorbed then gamma sampled, per-table challenge vectors extended with the bridge params. Aux build receives the extended vector (kernel + sigma columns); layer trees are freed before the LDE phases (the memory win). The fork-time L absorb vanishes with bus_public_inputs. Standard mode is untouched: multi_prove fails closed on GKR-mode AIRs and vice versa. Verifier (multi_verify_gkr): mirrors the phase order exactly — batch GKR replay (random point and instance claims derived from the transcript, never the proof), exact column-claim index-set check + reconstruction, claims absorbed then gamma sampled, per-table extended challenges into the unchanged rounds 2-4, and the bus balance checked on the batch proof's root claims (zero denominators rejected). GKR proofs must not carry bus_public_inputs (fail-closed). GPU paths are explicitly gated to standard mode (device_only_for and the fused composition): GKR v1 runs the CPU composition/aux paths. Tests: 12 e2e GKR-mode tests — multi-table and single-table roundtrips (mixed instance sizes), rkyv wrapper roundtrip, tampered root claim / column claim / child claims / sigma OOD / kernel OOD rejected, smuggled bus_public_inputs rejected, mode mismatches fail closed both ways, wrong expected balance rejected. Plus an #[ignore]d test documenting the known multi-interaction leaf-binding gap (port-plan §6). --- crypto/stark/src/constraints/evaluator.rs | 6 + crypto/stark/src/logup_gkr.rs | 43 ++- crypto/stark/src/proof/stark.rs | 9 +- crypto/stark/src/prover.rs | 229 +++++++++++- crypto/stark/src/test_utils.rs | 24 +- .../src/tests/bus_tests/gkr_mode_tests.rs | 344 ++++++++++++++++++ crypto/stark/src/tests/bus_tests/mod.rs | 1 + crypto/stark/src/verifier.rs | 221 ++++++++++- 8 files changed, 847 insertions(+), 30 deletions(-) create mode 100644 crypto/stark/src/tests/bus_tests/gkr_mode_tests.rs diff --git a/crypto/stark/src/constraints/evaluator.rs b/crypto/stark/src/constraints/evaluator.rs index 4e82a7a1b..7b2ff806d 100644 --- a/crypto/stark/src/constraints/evaluator.rs +++ b/crypto/stark/src/constraints/evaluator.rs @@ -321,6 +321,12 @@ where if crate::gpu_lde::gpu_composition_disabled() { return None; } + // GKR mode: the bridge constraint reads challenge leaves beyond the + // [z, α] pair the device challenge plumbing was built for — CPU path + // until the kernel grows GKR support. + if air.logup_mode() != crate::lookup::LogUpMode::Standard { + return None; + } if !zerofier_data.is_uniform() { return None; } diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index 321068995..f88fda228 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -19,7 +19,7 @@ use math::field::{ element::FieldElement, - traits::{IsFFTField, IsField, IsPrimeField, IsSubFieldOf}, + traits::{IsFFTField, IsField, IsSubFieldOf}, }; #[cfg(feature = "parallel")] use rayon::prelude::*; @@ -149,7 +149,7 @@ pub fn compute_logup_leaf_fractions( challenges: &[FieldElement], ) -> (Vec>, Vec>) where - F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { assert!( @@ -207,7 +207,7 @@ pub fn compute_logup_layers( challenges: &[FieldElement], ) -> Vec> where - F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { let (numerators, denominators) = @@ -252,7 +252,7 @@ pub fn finalize_logup_gkr_result( table_contribution: FieldElement, ) -> LogUpGkrResult where - F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { let col_indices = extract_column_indices(interactions); @@ -363,7 +363,7 @@ pub(crate) fn build_gkr_aux_columns( column_indices: &[usize], challenges: &[FieldElement], ) where - F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, + F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { let trace_len = trace.num_rows(); @@ -762,6 +762,39 @@ mod tests { assert_eq!(root_value, expected_l, "GKR root != standard-mode L"); } + /// KNOWN SOUNDNESS GAP (`thoughts/logup-gkr/port-plan.md` §6): for + /// multi-interaction tables nothing binds the batch-GKR leaf claims + /// `(n̂, d̂)` to the committed columns — the leaf fraction is nonlinear in + /// the columns, so the reconstruction check is fail-open there. This test + /// asserts the DESIRED fail-closed behavior with fabricated leaf claims + /// and honest column claims; it is `#[ignore]`d because it FAILS today. + /// The input-layer binding fix un-ignores it. + #[test] + #[ignore = "documents the multi-interaction leaf-binding gap (fail-open); the input-layer fix makes this pass"] + fn reconstruct_multi_interaction_rejects_fabricated_leaf_claims() { + let interactions = vec![ + BusInteraction::sender(1u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver(2u64, Multiplicity::One, Packing::Direct.columns(&[1])), + ]; + let challenges = vec![FE::from(1000u64), FE::from(7u64)]; + // Honest-looking column claims, fabricated leaf claims that no leaf + // vector consistent with the columns could produce. + let column_claims = vec![(0usize, FE::from(5u64)), (1usize, FE::from(9u64))]; + let fabricated_n = FE::from(0xBADu64); + let fabricated_d = FE::from(0xC0DEu64); + assert!( + !reconstruct_and_verify_gkr_claims( + &fabricated_n, + &fabricated_d, + &column_claims, + &interactions, + &challenges, + 3, + ), + "fabricated multi-interaction leaf claims must be rejected" + ); + } + /// Bridge parameters: Δ·N == Σ γʲ·cⱼ + γᴷ·norm_claim, and the kernel norm /// claim matches the actual kernel's Σ l². #[test] diff --git a/crypto/stark/src/proof/stark.rs b/crypto/stark/src/proof/stark.rs index 4edc5afda..23bac2875 100644 --- a/crypto/stark/src/proof/stark.rs +++ b/crypto/stark/src/proof/stark.rs @@ -129,6 +129,10 @@ pub struct MultiProof, E: IsField, PI> { pub proofs: Vec>, } +/// One table's GKR column claims: `(column_index, ⟨l, col⟩ MLE claim)` pairs +/// in canonical [`crate::logup_gkr::extract_column_indices`] order. +pub type GkrColumnClaims = Vec<(usize, FieldElement)>; + /// A multi-table proof under [`crate::lookup::LogUpMode::Gkr`]: the per-table /// STARK proofs plus the LogUp-GKR artifacts. /// @@ -156,7 +160,6 @@ pub struct GkrMultiProof, E: IsField, PI> { /// The batch GKR proof across every interacting table's summation tree. pub batch_gkr_proof: BatchGkrProof, /// Per-table column claims, aligned with `multi.proofs`: `None` for - /// non-interacting tables, else `(column_index, ⟨l, col⟩ MLE claim)` in - /// canonical [`crate::logup_gkr::extract_column_indices`] order. - pub column_claims_by_table: Vec)>>>, + /// non-interacting tables, else the table's [`GkrColumnClaims`]. + pub column_claims_by_table: Vec>>, } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8c44c42a9..333956895 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -27,8 +27,13 @@ use rayon::prelude::{ #[cfg(feature = "debug-checks")] use crate::debug::validate_trace; use crate::fri; -use crate::lookup::LOGUP_NUM_CHALLENGES; -use crate::proof::stark::{DeepPolynomialOpenings, PolynomialOpenings}; +use crate::gkr::{BatchGkrProof, gkr_prove_batch, instance_eval_point}; +use crate::logup_gkr::{ + LogUpGkrResult, compute_logup_layers, extend_rap_challenges_with_bridge, + finalize_logup_gkr_result, +}; +use crate::lookup::{LOGUP_NUM_CHALLENGES, LogUpMode}; +use crate::proof::stark::{DeepPolynomialOpenings, GkrMultiProof, PolynomialOpenings}; #[cfg(feature = "disk-spill")] use crate::storage_mode::StorageMode; use crate::table::Table; @@ -55,6 +60,13 @@ type AirTracePair<'a, Field, FieldExtension, PI> = ( &'a PI, ); +/// GKR-mode artifacts produced by round 1's batch GKR phase, destined for the +/// [`GkrMultiProof`] wrapper. +struct GkrProverArtifacts { + batch_gkr_proof: BatchGkrProof, + column_claims_by_table: Vec>>, +} + /// A default STARK prover implementing `IsStarkProver`. pub struct Prover< Field: IsSubFieldOf + IsFFTField + Send + Sync, @@ -809,6 +821,12 @@ pub trait IsStarkProver< if !air.has_aux_trace() || air.constraints_meta().is_empty() { return false; } + // GKR mode runs the CPU composition/aux paths (the GPU LogUp builders + // and the fused composition's challenge plumbing encode the standard + // term/acc layout); device-only would strand round 2 without host LDEs. + if air.logup_mode() != LogUpMode::Standard { + return false; + } let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let n = domain.interpolation_domain_size; let offsets_contiguous = @@ -2317,10 +2335,84 @@ pub trait IsStarkProver< /// /// The transcript must be safely initialized before passing it to this method. fn multi_prove( - mut air_trace_pairs: Vec>, + air_trace_pairs: Vec>, transcript: &mut (impl IsStarkTranscript + Clone + Send), #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + PI: Send + Sync + Clone, + Field: Copy + 'static, + FieldExtension: Copy + 'static, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, + { + let (multi, gkr_artifacts) = Self::multi_prove_impl( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + if gkr_artifacts.is_some() { + return Err(ProvingError::WrongParameter( + "AIRs in LogUpMode::Gkr must be proven with multi_prove_gkr".to_string(), + )); + } + Ok(multi) + } + + /// [`Self::multi_prove`] for [`LogUpMode::Gkr`] tables: same round + /// structure with the batch GKR phase active, returning the + /// [`GkrMultiProof`] wrapper (per-table proofs + batch GKR proof + column + /// claims). Every interacting AIR must be in GKR mode. + fn multi_prove_gkr( + air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + PI: Send + Sync + Clone, + Field: Copy + 'static, + FieldExtension: Copy + 'static, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, + { + let (multi, gkr_artifacts) = Self::multi_prove_impl( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + let artifacts = gkr_artifacts.ok_or_else(|| { + ProvingError::WrongParameter( + "multi_prove_gkr requires at least one AIR in LogUpMode::Gkr".to_string(), + ) + })?; + Ok(GkrMultiProof { + multi, + batch_gkr_proof: artifacts.batch_gkr_proof, + column_claims_by_table: artifacts.column_claims_by_table, + }) + } + + /// Shared driver behind [`Self::multi_prove`] / [`Self::multi_prove_gkr`]: + /// the artifacts component is `Some` exactly when the tables run in + /// [`LogUpMode::Gkr`]. + #[allow(clippy::type_complexity)] + fn multi_prove_impl( + mut air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result< + ( + MultiProof, + Option>, + ), + ProvingError, + > where FieldElement: AsBytes, FieldElement: AsBytes, @@ -2529,6 +2621,116 @@ pub trait IsStarkProver< Vec::new() }; + // ===================================================================== + // Round 1, Phase B′/B″ (GKR mode): batch GKR, column claims, γ + // ===================================================================== + // In LogUpMode::Gkr the per-table LogUp sums are proven here by ONE + // batch GKR on the SHARED transcript (before forking): summation trees + // over every interacting table's leaf fractions, then every table's + // column claims are absorbed and γ is sampled, and each table's + // challenge vector is extended with its bridge parameters. In standard + // mode `table_challenges` stays the shared `[z, α]` pair. + let gkr_mode = air_trace_pairs + .iter() + .any(|(air, _, _)| air.logup_mode() == LogUpMode::Gkr); + if gkr_mode + && air_trace_pairs.iter().any(|(air, _, _)| { + air.has_trace_interaction() && air.logup_mode() != LogUpMode::Gkr + }) + { + // Prover and verifier derive expectations from the AIR config; a + // mix of modes across interacting tables has no coherent transcript. + return Err(ProvingError::WrongParameter( + "all interacting tables must share LogUpMode::Gkr".to_string(), + )); + } + let mut table_challenges: Vec>> = + vec![lookup_challenges.clone(); num_airs]; + let mut gkr_artifacts: Option> = None; + if gkr_mode { + let gkr_indices: Vec = air_trace_pairs + .iter() + .enumerate() + .filter(|(_, (air, _, _))| air.has_trace_interaction()) + .map(|(idx, _)| idx) + .collect(); + + // Leaf fractions + summation trees per instance (parallel). + // Transient: consumed by the batch prove below, freed before the + // aux/LDE phases — this transience is where GKR's memory win over + // committed term columns lives. + let layers_per_instance: Vec>> = + crate::par::par_map_collect(0..gkr_indices.len(), |k| { + let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; + compute_logup_layers( + air.bus_interactions(), + &trace.columns_main(), + trace.num_rows(), + &lookup_challenges, + ) + }); + + let (batch_gkr_proof, shared_point, final_claims) = + gkr_prove_batch(layers_per_instance, transcript); + + // Per-instance finalize (parallel): column-claim MLE evaluations at + // the instance random point (the per-column ⟨l, col⟩ inner products + // are the heavy part). + let results: Vec> = + crate::par::par_map_collect(0..gkr_indices.len(), |k| { + let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; + let n_vars = trace.num_rows().trailing_zeros() as usize; + let random_point = instance_eval_point(&shared_point, n_vars); + let (n_claim, d_claim) = final_claims[k]; + let (root_n, root_d) = batch_gkr_proof.root_claims[k]; + let table_contribution = root_n + * root_d + .inv() + .expect("honest GKR root denominator is nonzero"); + finalize_logup_gkr_result( + air.bus_interactions(), + &trace.columns_main(), + random_point, + n_claim, + d_claim, + table_contribution, + ) + }); + + // Phase B″: bind every table's column claims into the shared + // transcript, THEN sample γ — the claims must precede the challenge + // that batches them. + for result in &results { + for (_, claim) in &result.column_claims { + transcript.append_field_element(claim); + } + } + let gamma: FieldElement = transcript.sample_field_element(); + + let mut column_claims_by_table: Vec< + Option>, + > = vec![None; num_airs]; + for (k, result) in results.into_iter().enumerate() { + let idx = gkr_indices[k]; + let trace_len = air_trace_pairs[idx].1.num_rows(); + let mut challenges = lookup_challenges.clone(); + extend_rap_challenges_with_bridge( + &mut challenges, + &result.column_claims, + &gamma, + trace_len, + &result.random_point, + ); + table_challenges[idx] = challenges; + column_claims_by_table[idx] = Some(result.column_claims); + } + + gkr_artifacts = Some(GkrProverArtifacts { + batch_gkr_proof, + column_claims_by_table, + }); + } + // ===================================================================== // Phase C + Rounds 2-4: Forked per table // ===================================================================== @@ -2571,13 +2773,16 @@ pub trait IsStarkProver< } #[cfg(feature = "parallel")] - let aux_iter = air_trace_pairs.par_iter_mut(); + let aux_iter = air_trace_pairs.par_iter_mut().enumerate(); #[cfg(not(feature = "parallel"))] - let aux_iter = air_trace_pairs.iter_mut(); + let aux_iter = air_trace_pairs.iter_mut().enumerate(); + // Per-table challenge vectors: the shared [z, α] pair in standard mode, + // the bridge-extended vector in GKR mode (the GKR aux build reads its + // γ powers, Δ and random point from it). let bus_inputs_vec: Vec>> = aux_iter - .map(|(air, trace, _)| { + .map(|(idx, (air, trace, _))| { if air.has_aux_trace() { - air.build_auxiliary_trace(*trace, &lookup_challenges) + air.build_auxiliary_trace(*trace, &table_challenges[idx]) } else { None } @@ -2852,8 +3057,10 @@ pub trait IsStarkProver< #[cfg(not(feature = "cuda"))] let main_iter = main_commits.into_iter().zip(main_ldes); - for ((main_pack, aux_full), bus_public_inputs) in - main_iter.zip(aux_results).zip(bus_inputs_vec) + for (((main_pack, aux_full), bus_public_inputs), rap_challenges) in main_iter + .zip(aux_results) + .zip(bus_inputs_vec) + .zip(table_challenges) { #[cfg(feature = "cuda")] let ((main_commit, main_lde), gpu_main) = main_pack; @@ -2866,7 +3073,7 @@ pub trait IsStarkProver< commitments.push(Round1Commitments { main: main_commit, aux: aux_commit, - rap_challenges: lookup_challenges.clone(), + rap_challenges, bus_public_inputs, }); #[cfg(feature = "cuda")] @@ -3012,7 +3219,7 @@ pub trait IsStarkProver< }); } - Ok(MultiProof { proofs }) + Ok((MultiProof { proofs }, gkr_artifacts)) } /// Generate a STARK proof for a single AIR/trace. diff --git a/crypto/stark/src/test_utils.rs b/crypto/stark/src/test_utils.rs index f5cd19f80..6c2be2f72 100644 --- a/crypto/stark/src/test_utils.rs +++ b/crypto/stark/src/test_utils.rs @@ -1,6 +1,6 @@ //! Shared test helpers for the stark crate. -use crate::proof::stark::MultiProof; +use crate::proof::stark::{GkrMultiProof, MultiProof}; use crate::prover::{IsStarkProver, Prover, ProvingError}; use crate::trace::TraceTable; use crate::traits::AIR; @@ -36,3 +36,25 @@ where crate::storage_mode::StorageMode::Ram, ) } + +/// [`multi_prove_ram`] for [`crate::lookup::LogUpMode::Gkr`] tables. +pub fn multi_prove_gkr_ram( + air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), +) -> Result, ProvingError> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + Copy + 'static, + FieldExtension: IsField + Send + Sync + Copy + 'static, + PI: Send + Sync + Clone, + FieldElement: AsBytes + ByteConversion, + FieldElement: AsBytes + ByteConversion, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, +{ + Prover::::multi_prove_gkr( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ) +} diff --git a/crypto/stark/src/tests/bus_tests/gkr_mode_tests.rs b/crypto/stark/src/tests/bus_tests/gkr_mode_tests.rs new file mode 100644 index 000000000..55b1352d0 --- /dev/null +++ b/crypto/stark/src/tests/bus_tests/gkr_mode_tests.rs @@ -0,0 +1,344 @@ +//! LogUp-GKR mode: end-to-end completeness and soundness tests. +//! +//! The fixture mirrors `completeness_tests::test_multi_table_proof` (a CPU +//! table dispatching to ADD and MUL over buses), with every AIR switched to +//! [`LogUpMode::Gkr`]: proofs carry a batch GKR proof + column claims instead +//! of bus public inputs, and the verifier replays the batch GKR before the +//! per-table rounds. + +use crypto::fiat_shamir::default_transcript::DefaultTranscript; +use math::field::element::FieldElement; +use math::field::{ + extensions_goldilocks::Degree3GoldilocksExtensionField, goldilocks::GoldilocksField, +}; + +use crate::constraints::builder::EmptyConstraints; +use crate::examples::multi_table_lookup::{ + new_add_air_with_lookup, new_cpu_air_with_lookup, new_mul_air_with_lookup, +}; +use crate::lookup::{AirWithBuses, LogUpMode, NullBoundaryConstraintBuilder}; +use crate::proof::options::ProofOptions; +use crate::proof::stark::GkrMultiProof; +use crate::test_utils::multi_prove_gkr_ram; +use crate::trace::TraceTable; +use crate::traits::AIR; +use crate::verifier::{IsStarkVerifier, Verifier}; + +type F = GoldilocksField; +type E = Degree3GoldilocksExtensionField; +type FE = FieldElement; +type GkrAir = AirWithBuses; + +fn gkr_airs() -> (GkrAir, GkrAir, GkrAir) { + let proof_options = ProofOptions::default_test_options(); + ( + new_cpu_air_with_lookup(&proof_options).with_logup_mode(LogUpMode::Gkr), + new_add_air_with_lookup(&proof_options).with_logup_mode(LogUpMode::Gkr), + new_mul_air_with_lookup(&proof_options).with_logup_mode(LogUpMode::Gkr), + ) +} + +fn fixture_traces() -> (TraceTable, TraceTable, TraceTable) { + // CPU (8 rows) dispatches 4 additions and 4 multiplications; ADD and MUL + // (4 rows each) receive them with multiplicity 1. Mixed trace lengths + // exercise the batch GKR's staggered instance activation. + let cpu_trace = TraceTable::from_columns_main( + vec![ + // add_flag + [1u64, 0, 1, 0, 1, 1, 0, 0] + .iter() + .map(|&v| FE::from(v)) + .collect(), + // mul_flag + [0u64, 1, 0, 1, 0, 0, 1, 1] + .iter() + .map(|&v| FE::from(v)) + .collect(), + (1..=8).map(|v| FE::from(v as u64)).collect(), + (1..=8).map(|v| FE::from(10 * v as u64)).collect(), + [11u64, 40, 33, 160, 55, 66, 490, 640] + .iter() + .map(|&v| FE::from(v)) + .collect(), + ], + 1, + ); + let add_trace = TraceTable::from_columns_main( + vec![ + [1u64, 3, 5, 6].iter().map(|&v| FE::from(v)).collect(), + [10u64, 30, 50, 60].iter().map(|&v| FE::from(v)).collect(), + [11u64, 33, 55, 66].iter().map(|&v| FE::from(v)).collect(), + vec![FE::one(); 4], + ], + 1, + ); + let mul_trace = TraceTable::from_columns_main( + vec![ + [2u64, 4, 7, 8].iter().map(|&v| FE::from(v)).collect(), + [20u64, 40, 70, 80].iter().map(|&v| FE::from(v)).collect(), + [40u64, 160, 490, 640] + .iter() + .map(|&v| FE::from(v)) + .collect(), + vec![FE::one(); 4], + ], + 1, + ); + (cpu_trace, add_trace, mul_trace) +} + +/// Prove the fixture in GKR mode. +fn prove_fixture() -> (GkrAir, GkrAir, GkrAir, GkrMultiProof) { + let (cpu_air, add_air, mul_air) = gkr_airs(); + let (mut cpu_trace, mut add_trace, mut mul_trace) = fixture_traces(); + let air_trace_pairs: Vec<( + &dyn AIR, + _, + _, + )> = vec![ + (&cpu_air, &mut cpu_trace, &()), + (&add_air, &mut add_trace, &()), + (&mul_air, &mut mul_trace, &()), + ]; + let proof = multi_prove_gkr_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])) + .expect("GKR prove succeeds"); + (cpu_air, add_air, mul_air, proof) +} + +fn verify( + cpu_air: &GkrAir, + add_air: &GkrAir, + mul_air: &GkrAir, + proof: &GkrMultiProof, +) -> bool { + let airs: Vec<&dyn AIR> = + vec![cpu_air, add_air, mul_air]; + Verifier::multi_verify_gkr( + &airs, + proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ) +} + +/// Valid GKR-mode multi-table proof is accepted. +#[test_log::test] +fn test_gkr_multi_table_prove_verify() { + let (cpu_air, add_air, mul_air, proof) = prove_fixture(); + assert_eq!(proof.multi.proofs.len(), 3); + // GKR proofs carry no bus public inputs; the wrapper carries claims for + // every (interacting) table. + assert!( + proof + .multi + .proofs + .iter() + .all(|p| p.bus_public_inputs.is_none()) + ); + assert!(proof.column_claims_by_table.iter().all(|c| c.is_some())); + assert!(verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// The GKR wrapper survives an rkyv roundtrip and still verifies. +#[test_log::test] +fn test_gkr_proof_rkyv_roundtrip() { + let (cpu_air, add_air, mul_air, proof) = prove_fixture(); + let bytes = rkyv::to_bytes::(&proof).unwrap(); + let deserialized: GkrMultiProof = + rkyv::from_bytes::<_, rkyv::rancor::Error>(&bytes).unwrap(); + assert!(verify(&cpu_air, &add_air, &mul_air, &deserialized)); +} + +/// A single-table GKR proof roundtrips (the batch degenerates to one instance). +#[test_log::test] +fn test_gkr_single_table_prove_verify() { + // A self-balancing table: one sender and one receiver interaction over + // the same bus and columns cancel row-by-row. + let proof_options = ProofOptions::default_test_options(); + let air = { + use crate::lookup::{AuxiliaryTraceBuildData, BusInteraction, Multiplicity, Packing}; + AirWithBuses::::new( + 2, + AuxiliaryTraceBuildData { + interactions: vec![ + BusInteraction::sender(1u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver( + 1u64, + Multiplicity::One, + Packing::Direct.columns(&[0]), + ), + ], + }, + &proof_options, + 1, + EmptyConstraints, + ) + .with_logup_mode(LogUpMode::Gkr) + }; + let mut trace = TraceTable::from_columns_main( + vec![ + (1..=8).map(|v| FE::from(v as u64)).collect(), + (11..=18).map(|v| FE::from(v as u64)).collect(), + ], + 1, + ); + let pairs: Vec<( + &dyn AIR, + _, + _, + )> = vec![(&air, &mut trace, &())]; + let proof = + multi_prove_gkr_ram(pairs, &mut DefaultTranscript::::new(&[])).expect("prove succeeds"); + let airs: Vec<&dyn AIR> = vec![&air]; + assert!(Verifier::multi_verify_gkr( + &airs, + &proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + )); +} + +/// Tampering a batch-GKR root claim (the table contribution) is rejected — +/// this is the GKR analogue of tampering `table_contribution`. +#[test_log::test] +fn test_gkr_tampered_root_claim_rejected() { + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + proof.batch_gkr_proof.root_claims[0].0 += FieldElement::::one(); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Tampering a column claim is rejected (the claims are transcript-bound +/// before γ is sampled, and the bridge pins them to the committed trace). +#[test_log::test] +fn test_gkr_tampered_column_claim_rejected() { + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + let claims = proof.column_claims_by_table[1] + .as_mut() + .expect("ADD table has claims"); + claims[0].1 += FieldElement::::one(); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Dropping a column claim (or the whole table entry) is rejected by the +/// exact index-set check. +#[test_log::test] +fn test_gkr_missing_column_claims_rejected() { + let (cpu_air, add_air, mul_air, proof) = prove_fixture(); + + let mut truncated = proof.clone(); + truncated.column_claims_by_table[1] + .as_mut() + .expect("ADD table has claims") + .pop(); + assert!(!verify(&cpu_air, &add_air, &mul_air, &truncated)); + + let mut missing = proof; + missing.column_claims_by_table[1] = None; + assert!(!verify(&cpu_air, &add_air, &mul_air, &missing)); +} + +/// Tampering a layer proof's child claims breaks the transcript-bound GKR +/// replay and is rejected. +#[test_log::test] +fn test_gkr_tampered_child_claims_rejected() { + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + let layer = proof + .batch_gkr_proof + .layer_proofs + .last_mut() + .expect("batch proof has layers"); + layer.child_claims_by_instance[0][0] += FieldElement::::one(); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Tampering the σ (bridge running sum) next-row OOD evaluation is rejected — +/// σ is the only next-row read, so this is the pruned block's single column. +#[test_log::test] +fn test_gkr_tampered_sigma_ood_rejected() { + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + let add_proof = &mut proof.multi.proofs[1]; + assert_eq!( + add_proof.trace_ood_next_evaluations.width, 1, + "pruned next-row block is exactly the σ column" + ); + let corrupted = *add_proof.trace_ood_next_evaluations.get(0, 0) + FieldElement::one(); + add_proof.trace_ood_next_evaluations.set(0, 0, corrupted); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Tampering the Lagrange-kernel current-row OOD evaluation is rejected (the +/// kernel is bound by its boundary constraint and the bridge's γ^K·l² term). +#[test_log::test] +fn test_gkr_tampered_kernel_ood_rejected() { + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + let add_proof = &mut proof.multi.proofs[1]; + // ADD has 4 main columns; aux col 0 (kernel) is full-width index 4. + let kernel_idx = 4usize; + let corrupted = *add_proof.trace_ood_evaluations.get(0, kernel_idx) + FieldElement::one(); + add_proof + .trace_ood_evaluations + .set(0, kernel_idx, corrupted); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Smuggling bus public inputs into a GKR-mode proof is rejected (they would +/// otherwise be absorbed into the fork transcript). +#[test_log::test] +fn test_gkr_smuggled_bus_public_inputs_rejected() { + use crate::lookup::BusPublicInputs; + let (cpu_air, add_air, mul_air, mut proof) = prove_fixture(); + proof.multi.proofs[0].bus_public_inputs = + Some(BusPublicInputs::from_contribution(FieldElement::zero())); + assert!(!verify(&cpu_air, &add_air, &mul_air, &proof)); +} + +/// Mode mismatches fail closed in both directions: GKR-mode AIRs refuse the +/// standard entry point, and standard-mode AIRs refuse the GKR entry point. +#[test_log::test] +fn test_gkr_mode_mismatch_rejected() { + let (cpu_air, add_air, mul_air, proof) = prove_fixture(); + + // GKR-mode AIRs through the standard verifier (with the inner MultiProof). + let airs: Vec<&dyn AIR> = + vec![&cpu_air, &add_air, &mul_air]; + assert!( + !Verifier::multi_verify( + &airs, + &proof.multi, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "GKR-mode AIRs must be rejected by the standard entry point" + ); + + // Standard-mode AIRs through the GKR verifier. + let proof_options = ProofOptions::default_test_options(); + let std_cpu = new_cpu_air_with_lookup(&proof_options); + let std_add = new_add_air_with_lookup(&proof_options); + let std_mul = new_mul_air_with_lookup(&proof_options); + let std_airs: Vec<&dyn AIR> = + vec![&std_cpu, &std_add, &std_mul]; + assert!( + !Verifier::multi_verify_gkr( + &std_airs, + &proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "standard-mode AIRs must be rejected by the GKR entry point" + ); +} + +/// A wrong expected bus balance is rejected (the fixture balances at zero). +#[test_log::test] +fn test_gkr_wrong_expected_balance_rejected() { + let (cpu_air, add_air, mul_air, proof) = prove_fixture(); + let airs: Vec<&dyn AIR> = + vec![&cpu_air, &add_air, &mul_air]; + assert!(!Verifier::multi_verify_gkr( + &airs, + &proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::one(), + )); +} diff --git a/crypto/stark/src/tests/bus_tests/mod.rs b/crypto/stark/src/tests/bus_tests/mod.rs index a57ca9aaf..9f696a580 100644 --- a/crypto/stark/src/tests/bus_tests/mod.rs +++ b/crypto/stark/src/tests/bus_tests/mod.rs @@ -1,6 +1,7 @@ //! Tests for LogUp bus interactions. pub mod bus_value_tests; pub mod completeness_tests; +pub mod gkr_mode_tests; pub mod multiplicity_tests; pub mod packing_tests; pub mod soundness_tests; diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index 64ae24363..a35407297 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -9,8 +9,13 @@ pub use crate::proof::view::PiDeserializer; use crate::{ config::Commitment, domain::new_verifier_domain, - lookup::{BusPublicInputs, LOGUP_CHALLENGE_ALPHA, LOGUP_NUM_CHALLENGES, compute_alpha_powers}, - proof::stark::{ArchivedMultiProof, MultiProof}, + gkr::{BatchGkrProof, gkr_verify_batch, instance_eval_point}, + logup_gkr::{extend_rap_challenges_with_bridge, reconstruct_and_verify_gkr_claims}, + lookup::{ + BusPublicInputs, LOGUP_CHALLENGE_ALPHA, LOGUP_NUM_CHALLENGES, LogUpMode, + compute_alpha_powers, + }, + proof::stark::{ArchivedMultiProof, GkrMultiProof, MultiProof}, proof::view::{ DeepPolynomialOpeningView, FriDecommitmentView, MultiProofView, PolynomialOpeningsView, ProofViewSource, StarkProofView, StarkTableView, @@ -1125,13 +1130,77 @@ pub trait IsStarkVerifier< /// The single verification implementation, shared by [`Self::multi_verify`] /// (owned) and [`Self::multi_verify_archived`] (archived), operating on - /// proof views rather than either's concrete type. + /// proof views rather than either's concrete type. Standard LogUp mode + /// only: AIRs in [`LogUpMode::Gkr`] must go through + /// [`Self::multi_verify_gkr`] (fail-closed here). fn multi_verify_views<'p>( airs: &[&dyn AIR], proofs: impl ProofViewSource<'p, Field, FieldExtension, PI>, transcript: &mut (impl IsStarkTranscript + Clone), expected_bus_balance: &FieldElement, ) -> bool + where + Field: 'p, + FieldExtension: 'p, + PI: 'p, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + { + if airs.iter().any(|air| air.logup_mode() == LogUpMode::Gkr) { + error!("AIRs in LogUpMode::Gkr must be verified with multi_verify_gkr"); + return false; + } + Self::multi_verify_views_impl(airs, proofs, transcript, expected_bus_balance, None) + } + + /// Verify a [`GkrMultiProof`] ([`LogUpMode::Gkr`] tables): the standard + /// per-table round structure plus the batch GKR replay, column-claim + /// binding, and the root-claim bus-balance check. Every interacting AIR + /// must be in GKR mode (fail-closed). + fn multi_verify_gkr( + airs: &[&dyn AIR], + gkr_proof: &GkrMultiProof, + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + ) -> bool + where + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + PI: Clone, + { + if airs + .iter() + .any(|air| air.has_trace_interaction() && air.logup_mode() != LogUpMode::Gkr) + { + error!("multi_verify_gkr requires every interacting AIR in LogUpMode::Gkr"); + return false; + } + Self::multi_verify_views_impl( + airs, + MultiProofView::Owned(&gkr_proof.multi), + transcript, + expected_bus_balance, + Some(( + &gkr_proof.batch_gkr_proof, + gkr_proof.column_claims_by_table.as_slice(), + )), + ) + } + + /// Shared driver behind [`Self::multi_verify_views`] (standard) and + /// [`Self::multi_verify_gkr`]; `gkr` carries the batch proof and per-table + /// column claims exactly when the tables run in [`LogUpMode::Gkr`]. + #[allow(clippy::type_complexity)] + fn multi_verify_views_impl<'p, 'g>( + airs: &[&dyn AIR], + proofs: impl ProofViewSource<'p, Field, FieldExtension, PI>, + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + gkr: Option<( + &'g BatchGkrProof, + &'g [Option)>>], + )>, + ) -> bool where Field: 'p, FieldExtension: 'p, @@ -1226,15 +1295,122 @@ pub trait IsStarkVerifier< Vec::new() }; + // ===================================================================== + // Round 1, Phase B′/B″ (GKR mode): batch GKR replay, claims, γ + // ===================================================================== + // Mirrors the prover exactly: replay the batch GKR on the SHARED + // transcript (deriving the random point and per-instance leaf claims + // from the transcript, never the proof), bind every table's column + // claims, sample γ, and build each table's extended challenge vector. + let mut table_challenges: Vec>> = + vec![lookup_challenges.clone(); airs.len()]; + if let Some((batch_proof, column_claims_by_table)) = gkr { + if column_claims_by_table.len() != airs.len() { + error!("column_claims_by_table length does not match table count"); + return false; + } + let trace_lengths: Vec = proofs + .view_iter() + .map(|proof| proof.trace_length()) + .collect(); + let gkr_indices: Vec = airs + .iter() + .enumerate() + .filter(|(_, air)| air.has_trace_interaction()) + .map(|(idx, _)| idx) + .collect(); + + // Instance sizes come from the proofs' trace lengths (nonzero was + // checked in Phase A); reject non-power-of-two outright. + let mut n_layers_by_instance = Vec::with_capacity(gkr_indices.len()); + for &idx in &gkr_indices { + let trace_length = trace_lengths[idx]; + if !trace_length.is_power_of_two() { + error!("Table {idx}: trace length {trace_length} is not a power of two"); + return false; + } + n_layers_by_instance.push(trace_length.trailing_zeros() as usize); + } + + let (shared_point, per_instance_claims) = + match gkr_verify_batch(batch_proof, &n_layers_by_instance, transcript) { + Ok(result) => result, + Err(e) => { + error!("batch GKR verification failed: {e:?}"); + return false; + } + }; + + // Per table: claims must exist, match the canonical column set, and + // be consistent with the transcript-derived leaf claims; absorb + // them in the prover's order (ascending table index). + for (k, &idx) in gkr_indices.iter().enumerate() { + let Some(claims) = column_claims_by_table[idx].as_ref() else { + error!("Table {idx}: interacting GKR table is missing column claims"); + return false; + }; + let (n_claim, d_claim) = &per_instance_claims[k]; + if !reconstruct_and_verify_gkr_claims( + n_claim, + d_claim, + claims, + airs[idx].bus_interactions(), + &lookup_challenges, + n_layers_by_instance[k], + ) { + error!("Table {idx}: GKR column-claim reconstruction failed"); + return false; + } + for (_, claim) in claims { + transcript.append_field_element(claim); + } + } + for (idx, claims) in column_claims_by_table.iter().enumerate() { + if !airs[idx].has_trace_interaction() && claims.is_some() { + error!("Table {idx}: non-interacting table carries column claims"); + return false; + } + } + + let gamma: FieldElement = transcript.sample_field_element(); + + for (k, &idx) in gkr_indices.iter().enumerate() { + let claims = column_claims_by_table[idx] + .as_ref() + .expect("presence checked above"); + let trace_length = trace_lengths[idx]; + let random_point = instance_eval_point(&shared_point, n_layers_by_instance[k]); + let mut challenges = lookup_challenges.clone(); + extend_rap_challenges_with_bridge( + &mut challenges, + claims, + &gamma, + trace_length, + &random_point, + ); + table_challenges[idx] = challenges; + } + } + // ===================================================================== // Validate bus_public_inputs presence against AIR layout // ===================================================================== - // A dishonest prover could omit bus_public_inputs entirely (None) to - // bypass the bus balance check. With circular constraints, there are no - // boundary constraints on LogUp columns, so the bus balance check is - // the only cross-table validation. + // Standard mode: a dishonest prover could omit bus_public_inputs + // entirely (None) to bypass the bus balance check; with circular + // constraints there are no boundary constraints on LogUp columns, so + // the bus balance check is the only cross-table validation. GKR mode: + // no table may carry bus_public_inputs at all — the balance check runs + // on the batch proof's root claims, and a smuggled L would be absorbed + // into the fork transcript below. for (idx, (air, proof)) in airs.iter().zip(proofs.view_iter()).enumerate() { + if gkr.is_some() { + if proof.has_bus_public_inputs() { + error!("Table {idx}: GKR-mode proof must not carry bus_public_inputs"); + return false; + } + continue; + } if air.has_trace_interaction() && !proof.has_bus_public_inputs() { error!( "Table {idx}: AIR has LogUp interactions but proof is missing bus_public_inputs" @@ -1282,13 +1458,15 @@ pub trait IsStarkVerifier< None => return false, }; - // Rounds 2-4: verify + // Rounds 2-4: verify (per-table challenge vector — the shared + // [z, α] pair in standard mode, the bridge-extended vector in GKR + // mode). if !Self::verify_rounds_2_to_4( *air, proof, &public_inputs, &mut table_transcript, - lookup_challenges.clone(), + table_challenges[idx].clone(), ) { error!( "Table {} failed verify_rounds_2_to_4 (num_constraints={}, trace_cols={})", @@ -1310,7 +1488,30 @@ pub trait IsStarkVerifier< // receiver contributions are computed externally (e.g. verifier-computed // COMMIT output bus), the target is the missing positive remainder. - if needs_lookup_challenges { + if let Some((batch_proof, _)) = gkr { + // GKR mode: the balance is the sum of the batch proof's root + // claims (each table's total contribution n/d). The root claims + // were transcript-bound and reduced to the committed traces by the + // batch replay above; a zero denominator is an outright reject. + let mut total = FieldElement::::zero(); + for (root_n, root_d) in &batch_proof.root_claims { + let Ok(inv) = root_d.inv() else { + error!("GKR root claim has a zero denominator"); + return false; + }; + total += root_n * inv; + } + if total != *expected_bus_balance { + #[cfg(not(feature = "test_fiat_shamir"))] + error!( + "LogUp bus does not balance (GKR root claims): total={:?}, target={:?}", + total, expected_bus_balance + ); + return false; + } + #[cfg(feature = "debug-checks")] + info!("Bus balance check PASSED (GKR)"); + } else if needs_lookup_challenges { let mut total = FieldElement::::zero(); for (air, proof) in airs.iter().zip(proofs.view_iter()) { if air.has_trace_interaction() From af41001452a5a415c1fb0edb0728e7bde79651ce Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 13:08:01 -0300 Subject: [PATCH 07/20] =?UTF-8?q?feat(prover):=20VM-level=20LogUp-GKR=20en?= =?UTF-8?q?ablement=20=E2=80=94=20GkrVmProof=20+=20explicit=20entry=20poin?= =?UTF-8?q?ts?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - VmAirs::new_with_logup_mode: every table constructor takes the mode (new delegates with Standard; no env magic, prover and verifier must construct with the same mode) - prove_gkr_with_options_and_inputs -> GkrVmProof (GkrMultiProof + the usual verifier metadata; the standard VmProof wire format is untouched) - verify_gkr_with_options: same statement absorption and verifier-computed COMMIT bus target as the standard path, checked against the GKR root claims via multi_verify_gkr - fix single-row tables (n_vars = 0): the extended challenge vector ends exactly at the empty random point; kernel boundary is l[0] = 1 Tests: full VM table set (CPU..PAGE/REGISTER/DECODE incl. preprocessed, 1-row tables) proves+verifies in GKR mode (sub_neg_result minimal); library-path e2e on a committing program exercises the NONZERO COMMIT bus target and rejects tampered public output. --- crypto/stark/src/lookup.rs | 4 +- prover/src/lib.rs | 366 ++++++++++++++++++++++++--- prover/src/test_utils.rs | 16 ++ prover/src/tests/prove_elfs_tests.rs | 102 ++++++++ 4 files changed, 448 insertions(+), 40 deletions(-) diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 949eed1aa..63232c51a 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1403,8 +1403,10 @@ where LogUpMode::Gkr => { let rp_start = crate::logup_gkr::logup_random_point_start(self.gkr_column_indices.len()); + // == is the single-row case: n_vars = 0, empty random + // point, l = [1] (kernel of the empty point), l[0] = 1. assert!( - rap_challenges.len() > rp_start, + rap_challenges.len() >= rp_start, "GKR-mode boundary constraints require the extended \ challenge vector (got {} challenges, random point starts at {rp_start})", rap_challenges.len(), diff --git a/prover/src/lib.rs b/prover/src/lib.rs index ff9601bb4..7d13b2ef1 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -34,6 +34,7 @@ use crypto::fiat_shamir::is_transcript::IsTranscript; use executor::elf::Elf; use executor::vm::execution::Executor; use math::field::element::FieldElement; +use stark::lookup::LogUpMode; use stark::prover::{IsStarkProver, Prover}; #[cfg(feature = "disk-spill")] use stark::storage_mode::StorageMode; @@ -178,6 +179,24 @@ pub struct VmProof { pub num_private_input_pages: usize, } +/// A complete VM proof bundle under [`LogUpMode::Gkr`] (experimental): the +/// [`stark::proof::stark::GkrMultiProof`] wrapper plus the same verifier +/// metadata as [`VmProof`]. A separate type on purpose — the standard +/// [`VmProof`] wire format is untouched by the GKR experiment. +#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct GkrVmProof { + /// The multi-table STARK proofs plus the batch GKR artifacts. + pub proof: stark::proof::stark::GkrMultiProof, + /// See [`VmProof::runtime_page_ranges`]. + pub runtime_page_ranges: Vec, + /// See [`VmProof::table_counts`]. + pub table_counts: TableCounts, + /// See [`VmProof::public_output`]. + pub public_output: Vec, + /// See [`VmProof::num_private_input_pages`]. + pub num_private_input_pages: usize, +} + /// The private-input bundle the recursion verifier guest consumes: an inner /// proof plus the DECODE/ELF-data-page commitments supplied instead of /// recomputed in-VM (see `bench_vs/lambda/recursion`), and the inner ELF bytes @@ -672,6 +691,7 @@ impl VmAirs { /// Fiat-Shamir); a consistent prover-supplied mismatch is NOT — such /// callers must bind identity externally (see `recursion::check_attestation`). #[allow(clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments)] pub fn new( elf: &Elf, proof_options: &ProofOptions, @@ -683,46 +703,104 @@ impl VmAirs { register_init: Option<&[u32]>, page_commitments: Option<&[(u64, Commitment)]>, register_preprocessed: Option<(Commitment, usize)>, + ) -> Self { + Self::new_with_logup_mode( + elf, + proof_options, + minimal_bitwise, + page_configs, + table_counts, + decode_commitment, + include_halt, + register_init, + page_commitments, + register_preprocessed, + LogUpMode::Standard, + ) + } + + /// [`Self::new`] with an explicit LogUp mode: [`LogUpMode::Gkr`] switches + /// every table to the batch-GKR LogUp path (2 aux columns, bridge + /// constraint). Prover and verifier must construct their `VmAirs` with the + /// same mode. + #[allow(clippy::too_many_arguments)] + pub fn new_with_logup_mode( + elf: &Elf, + proof_options: &ProofOptions, + minimal_bitwise: bool, + page_configs: &[crate::tables::page::PageConfig], + table_counts: &TableCounts, + decode_commitment: Option, + include_halt: bool, + register_init: Option<&[u32]>, + page_commitments: Option<&[(u64, Commitment)]>, + register_preprocessed: Option<(Commitment, usize)>, + logup_mode: LogUpMode, ) -> Self { let cpus: Vec<_> = (0..table_counts.cpu) .map(|i| { - Box::new(create_cpu_air(proof_options).with_name(&format!("CPU[{}]", i))) as VmAir + Box::new( + create_cpu_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("CPU[{}]", i)), + ) as VmAir }) .collect(); let bitwise: VmAir = if minimal_bitwise { - Box::new(create_bitwise_air(proof_options)) + Box::new(create_bitwise_air(proof_options).with_logup_mode(logup_mode)) } else { - Box::new(create_bitwise_air(proof_options).with_preprocessed( - bitwise::preprocessed_commitment(proof_options), - bitwise::NUM_PRECOMPUTED_COLS, - )) + Box::new( + create_bitwise_air(proof_options) + .with_logup_mode(logup_mode) + .with_preprocessed( + bitwise::preprocessed_commitment(proof_options), + bitwise::NUM_PRECOMPUTED_COLS, + ), + ) }; let lts: Vec<_> = (0..table_counts.lt) .map(|i| { - Box::new(create_lt_air(proof_options).with_name(&format!("LT[{}]", i))) as VmAir + Box::new( + create_lt_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("LT[{}]", i)), + ) as VmAir }) .collect(); let shifts: Vec<_> = (0..table_counts.shift) .map(|i| { - Box::new(create_shift_air(proof_options).with_name(&format!("SHIFT[{}]", i))) - as VmAir + Box::new( + create_shift_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("SHIFT[{}]", i)), + ) as VmAir }) .collect(); let memws: Vec<_> = (0..table_counts.memw) .map(|i| { - Box::new(create_memw_air(proof_options).with_name(&format!("MEMW[{}]", i))) as VmAir + Box::new( + create_memw_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("MEMW[{}]", i)), + ) as VmAir }) .collect(); let memw_aligneds: Vec<_> = (0..table_counts.memw_aligned) .map(|i| { Box::new( - create_memw_aligned_air(proof_options).with_name(&format!("MEMW_A[{}]", i)), + create_memw_aligned_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("MEMW_A[{}]", i)), ) as VmAir }) .collect(); let loads: Vec<_> = (0..table_counts.load) .map(|i| { - Box::new(create_load_air(proof_options).with_name(&format!("LOAD[{}]", i))) as VmAir + Box::new( + create_load_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("LOAD[{}]", i)), + ) as VmAir }) .collect(); let decode_root = decode_commitment.unwrap_or_else(|| { @@ -731,48 +809,70 @@ impl VmAirs { }); let decode: VmAir = Box::new( create_decode_air(proof_options) + .with_logup_mode(logup_mode) .with_preprocessed(decode_root, decode::NUM_PRECOMPUTED_COLS), ); let muls: Vec<_> = (0..table_counts.mul) .map(|i| { - Box::new(create_mul_air(proof_options).with_name(&format!("MUL[{}]", i))) as VmAir + Box::new( + create_mul_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("MUL[{}]", i)), + ) as VmAir }) .collect(); let dvrms: Vec<_> = (0..table_counts.dvrm) .map(|i| { - Box::new(create_dvrm_air(proof_options).with_name(&format!("DVRM[{}]", i))) as VmAir + Box::new( + create_dvrm_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("DVRM[{}]", i)), + ) as VmAir }) .collect(); let branches: Vec<_> = (0..table_counts.branch) .map(|i| { - Box::new(create_branch_air(proof_options).with_name(&format!("BRANCH[{}]", i))) - as VmAir + Box::new( + create_branch_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("BRANCH[{}]", i)), + ) as VmAir }) .collect(); - let halt: VmAir = Box::new(create_halt_air(proof_options)); - let commit: VmAir = Box::new(create_commit_air(proof_options)); - let keccak: VmAir = Box::new(create_keccak_air(proof_options)); - let keccak_rnd: VmAir = Box::new(create_keccak_rnd_air(proof_options)); - let keccak_rc: VmAir = Box::new(create_keccak_rc_air(proof_options).with_preprocessed( - tables::keccak_rc::preprocessed_commitment(proof_options), - tables::keccak_rc::NUM_PRECOMPUTED_COLS, - )); - let ecsm: VmAir = Box::new(create_ecsm_air(proof_options)); - let ecdas: VmAir = Box::new(create_ecdas_air(proof_options)); + let halt: VmAir = Box::new(create_halt_air(proof_options).with_logup_mode(logup_mode)); + let commit: VmAir = Box::new(create_commit_air(proof_options).with_logup_mode(logup_mode)); + let keccak: VmAir = Box::new(create_keccak_air(proof_options).with_logup_mode(logup_mode)); + let keccak_rnd: VmAir = + Box::new(create_keccak_rnd_air(proof_options).with_logup_mode(logup_mode)); + let keccak_rc: VmAir = Box::new( + create_keccak_rc_air(proof_options) + .with_logup_mode(logup_mode) + .with_preprocessed( + tables::keccak_rc::preprocessed_commitment(proof_options), + tables::keccak_rc::NUM_PRECOMPUTED_COLS, + ), + ); + let ecsm: VmAir = Box::new(create_ecsm_air(proof_options).with_logup_mode(logup_mode)); + let ecdas: VmAir = Box::new(create_ecdas_air(proof_options).with_logup_mode(logup_mode)); let register: VmAir = if let Some((commitment, num_preprocessed_cols)) = register_preprocessed { Box::new( create_register_air(proof_options) + .with_logup_mode(logup_mode) .with_preprocessed(commitment, num_preprocessed_cols), ) } else { let register_init = register_init .map(<[u32]>::to_vec) .unwrap_or_else(|| register::register_init_from_entry_point(elf.entry_point)); - Box::new(create_register_air(proof_options).with_preprocessed( - register::preprocessed_commitment(proof_options, ®ister_init), - register::NUM_PREPROCESSED_COLS, - )) + Box::new( + create_register_air(proof_options) + .with_logup_mode(logup_mode) + .with_preprocessed( + register::preprocessed_commitment(proof_options, ®ister_init), + register::NUM_PREPROCESSED_COLS, + ), + ) }; // Every zero-init page shares one preprocessed commitment: OFFSET is // page-relative and INIT is all-zero, so it depends only on @@ -785,7 +885,8 @@ impl VmAirs { let pages: Vec = page_configs .iter() .map(|config| -> VmAir { - let air = create_page_air(proof_options, config.page_base); + let air = + create_page_air(proof_options, config.page_base).with_logup_mode(logup_mode); if config.is_private_input { // Private-input pages: all columns are main trace (not preprocessed). // The verifier doesn't see the init values; correctness is enforced @@ -815,31 +916,46 @@ impl VmAirs { let memw_registers: Vec<_> = (0..table_counts.memw_register) .map(|i| { Box::new( - create_memw_register_air(proof_options).with_name(&format!("MEMW_R[{}]", i)), + create_memw_register_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("MEMW_R[{}]", i)), ) as VmAir }) .collect(); let eqs: Vec<_> = (0..table_counts.eq) .map(|i| { - Box::new(create_eq_air(proof_options).with_name(&format!("EQ[{}]", i))) as VmAir + Box::new( + create_eq_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("EQ[{}]", i)), + ) as VmAir }) .collect(); let bytewises: Vec<_> = (0..table_counts.bytewise) .map(|i| { - Box::new(create_bytewise_air(proof_options).with_name(&format!("BYTEWISE[{}]", i))) - as VmAir + Box::new( + create_bytewise_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("BYTEWISE[{}]", i)), + ) as VmAir }) .collect(); let stores: Vec<_> = (0..table_counts.store) .map(|i| { - Box::new(create_store_air(proof_options).with_name(&format!("STORE[{}]", i))) - as VmAir + Box::new( + create_store_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("STORE[{}]", i)), + ) as VmAir }) .collect(); let cpu32s: Vec<_> = (0..table_counts.cpu32) .map(|i| { - Box::new(create_cpu32_air(proof_options).with_name(&format!("CPU32[{}]", i))) - as VmAir + Box::new( + create_cpu32_air(proof_options) + .with_logup_mode(logup_mode) + .with_name(&format!("CPU32[{}]", i)), + ) as VmAir }) .collect(); @@ -1208,6 +1324,178 @@ pub fn prove_with_options_and_inputs( }) } +/// [`prove_with_options_and_inputs`] under [`LogUpMode::Gkr`] (experimental): +/// every table proves its LogUp sums with the batch GKR instead of committed +/// term/acc columns (2 aux columns per table). Returns the [`GkrVmProof`] +/// bundle; verify with [`verify_gkr_with_options`]. +pub fn prove_gkr_with_options_and_inputs( + elf_bytes: &[u8], + private_inputs: &[u8], + proof_options: &ProofOptions, + max_rows: &MaxRowsConfig, +) -> Result { + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + + let table_counts = traces.table_counts(); + let airs = VmAirs::new_with_logup_mode( + &program, + proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + LogUpMode::Gkr, + ); + + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut transcript, + StatementKind::Monolithic, + elf_bytes, + &traces.public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + + let proof = Prover::multi_prove_gkr( + airs.air_trace_pairs(&mut traces), + &mut transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + + Ok(GkrVmProof { + proof, + runtime_page_ranges, + table_counts, + public_output: traces.public_output_bytes.clone(), + num_private_input_pages, + }) +} + +/// Verify a [`GkrVmProof`] ([`LogUpMode::Gkr`] tables) with caller-specified +/// proof options. The GKR analogue of [`verify_with_options`]. +pub fn verify_gkr_with_options( + gkr_proof: &GkrVmProof, + elf_bytes: &[u8], + proof_options: &ProofOptions, +) -> Result { + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let elf_digest = statement::elf_digest(elf_bytes); + + gkr_proof.table_counts.validate()?; + { + let max_pages = crate::tables::page::max_private_input_pages(); + if gkr_proof.num_private_input_pages > max_pages { + return Err(Error::InvalidTableCounts(format!( + "num_private_input_pages ({}) exceeds max ({max_pages})", + gkr_proof.num_private_input_pages, + ))); + } + } + + let page_configs = Traces::page_configs_from_elf_and_runtime( + &program, + &gkr_proof.runtime_page_ranges, + gkr_proof.num_private_input_pages, + ); + + let expected_proof_count = + gkr_proof.table_counts.total() + FIXED_TABLE_COUNT + page_configs.len(); + if expected_proof_count != gkr_proof.proof.multi.proofs.len() { + return Err(Error::InvalidTableCounts(format!( + "table_counts total ({}) + {FIXED_TABLE_COUNT} fixed + {} pages = {}, but proof contains {} sub-proofs", + gkr_proof.table_counts.total(), + page_configs.len(), + expected_proof_count, + gkr_proof.proof.multi.proofs.len(), + ))); + } + + let airs = VmAirs::new_with_logup_mode( + &program, + proof_options, + false, + &page_configs, + &gkr_proof.table_counts, + None, + true, + None, + None, + None, + LogUpMode::Gkr, + ); + let air_refs = airs.air_refs(); + + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement_with_digest( + &mut transcript, + StatementKind::Monolithic, + &elf_digest, + &gkr_proof.public_output, + &gkr_proof.table_counts, + gkr_proof.num_private_input_pages, + &gkr_proof.runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + + // The COMMIT output bus stays an external receiver in GKR mode: the same + // verifier-computed target, checked against the GKR root claims. + let mut transcript_for_replay = transcript.clone(); + let expected_bus_balance = match compute_expected_commit_bus_balance_view( + &air_refs, + MultiProofView::Owned(&gkr_proof.proof.multi), + &gkr_proof.public_output, + 0, + &mut transcript_for_replay, + ) { + Some(balance) => balance, + None => return Ok(false), + }; + + Ok(Verifier::multi_verify_gkr( + &air_refs, + &gkr_proof.proof, + &mut transcript, + &expected_bus_balance, + )) +} + /// Verify a proof produced by [`prove`] using default proof options. /// /// Uses [`GoldilocksCubicProofOptions::with_blowup(2)`] for verification. diff --git a/prover/src/test_utils.rs b/prover/src/test_utils.rs index 6dd28ce71..9dd6d67dd 100644 --- a/prover/src/test_utils.rs +++ b/prover/src/test_utils.rs @@ -142,6 +142,22 @@ where ) } +/// [`multi_prove_ram`] for [`stark::lookup::LogUpMode::Gkr`] tables. +pub fn multi_prove_gkr_ram( + air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), +) -> Result, ProvingError> +where + PI: Send + Sync + Clone, +{ + Prover::::multi_prove_gkr( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + StorageMode::Ram, + ) +} + // ============================================================================= // Soundness regression helpers (negative AIR tests) // ============================================================================= diff --git a/prover/src/tests/prove_elfs_tests.rs b/prover/src/tests/prove_elfs_tests.rs index ffe9071b2..a8407b933 100644 --- a/prover/src/tests/prove_elfs_tests.rs +++ b/prover/src/tests/prove_elfs_tests.rs @@ -100,6 +100,64 @@ fn prove_and_verify_vm_minimal(elf: &Elf, traces: &mut Traces) -> bool { ) } +/// [`prove_and_verify_vm_minimal`] under `LogUpMode::Gkr`: the full VM table +/// set proven with the batch-GKR LogUp path (2 aux columns + bridge per +/// table) and verified with the GKR verifier, including the verifier-computed +/// COMMIT bus target against the GKR root claims. +fn prove_and_verify_vm_minimal_gkr(elf: &Elf, traces: &mut Traces) -> bool { + use stark::lookup::LogUpMode; + + let _ = env_logger::builder().is_test(true).try_init(); + let proof_options = ProofOptions::default_test_options(); + + let table_counts = traces.table_counts(); + let airs = VmAirs::new_with_logup_mode( + elf, + &proof_options, + true, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + LogUpMode::Gkr, + ); + + let air_trace_pairs = airs.air_trace_pairs(traces); + let gkr_proof = match crate::test_utils::multi_prove_gkr_ram( + air_trace_pairs, + &mut DefaultTranscript::::new(&[]), + ) { + Ok(proof) => proof, + Err(_) => return false, + }; + + let views: Vec> = gkr_proof + .multi + .proofs + .iter() + .map(StarkProofView::Owned) + .collect(); + let mut replay_transcript = DefaultTranscript::::new(&[]); + let expected_bus_balance = crate::compute_expected_commit_bus_balance_view( + &airs.air_refs(), + &views, + &traces.public_output_bytes, + 0, + &mut replay_transcript, + ) + .expect("fingerprint collision in test"); + + Verifier::multi_verify_gkr( + &airs.air_refs(), + &gkr_proof, + &mut DefaultTranscript::::new(&[]), + &expected_bus_balance, + ) +} + /// Like [`crate::prove_with_options_and_inputs`] but trims the bitwise table to the /// rows the program uses instead of proving the full 2^20-row table (TEST ONLY). /// @@ -301,6 +359,20 @@ fn test_prove_elfs_sub_neg_result_fast() { ); } +/// The full VM table set proven and verified under LogUp-GKR mode +/// (`sub_neg_result`, minimal tables — the GKR twin of the fast test above). +#[test] +fn test_prove_elfs_sub_neg_result_gkr() { + let (elf, logs, instructions) = run_asm_elf("sub_neg_result"); + let mut traces = + Traces::from_logs_minimal(&logs, instructions.clone(), &Default::default()).unwrap(); + + assert!( + prove_and_verify_vm_minimal_gkr(&elf, &mut traces), + "GKR-mode proof verification failed for sub_neg_result program" + ); +} + #[test] fn test_prove_elfs_sub_underflow_fast() { let (elf, logs, instructions) = run_asm_elf("sub_underflow"); @@ -1442,6 +1514,36 @@ fn test_verify_rejects_tampered_public_output() { ); } +/// The GKR library path end-to-end: `prove_gkr_with_options_and_inputs` → +/// `verify_gkr_with_options` on a committing program, so the GKR bus balance +/// is checked against the NONZERO verifier-computed COMMIT target; a tampered +/// public output must be rejected. +#[test] +fn test_prove_verify_gkr_library_path_commit() { + let elf_bytes = crate::test_utils::asm_elf_bytes("test_commit_4"); + let proof_options = ProofOptions::default_test_options(); + let gkr_proof = crate::prove_gkr_with_options_and_inputs( + &elf_bytes, + &[], + &proof_options, + &Default::default(), + ) + .expect("GKR prover should succeed for test_commit_4"); + assert!( + crate::verify_gkr_with_options(&gkr_proof, &elf_bytes, &proof_options) + .expect("valid GKR commit proof should not error"), + "GKR proof of a committing program should verify" + ); + + let mut tampered = gkr_proof.clone(); + tampered.public_output[0] ^= 0x01; + assert!( + !crate::verify_gkr_with_options(&tampered, &elf_bytes, &proof_options) + .expect("verifier should not error on tampered public output"), + "GKR verifier should reject a tampered public output" + ); +} + /// Slow version using full bitwise table (2^20 rows) - production-safe /// /// This is the most comprehensive test covering all RV64IM instructions: From ebbba689e2db06b64793b185d4e6f6239fffeaad Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 14:08:52 -0300 Subject: [PATCH 08/20] feat(cli): --logup-mode flag, defaulting to gkr on this branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A plain `cli prove` — exactly what the /bench workflow runs — now proves with the batch-GKR LogUp path, so the PR-vs-main bench comparison is the GKR-vs-standard A/B on the same workload. `--logup-mode standard` keeps the production path reachable; `verify` takes the matching flag (the two modes have different proof bundle types). Timing still covers proof generation only, identically in both modes. Smoke-tested both modes end-to-end (prove + verify) on test_commit_4; GKR proof bundle is 20.8 MB vs 33.7 MB standard on that program. --- bin/cli/src/main.rs | 165 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 128 insertions(+), 37 deletions(-) diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index 1de36220b..51a06da33 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -6,7 +6,7 @@ use std::path::{Path, PathBuf}; use std::process::ExitCode; use std::time::Instant; -use clap::{Parser, Subcommand, ValueHint}; +use clap::{Parser, Subcommand, ValueEnum, ValueHint}; #[global_allocator] static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; @@ -19,6 +19,17 @@ use stark::proof::options::GoldilocksCubicProofOptions; const DEFAULT_CONTINUATION_EPOCH_SIZE_LOG2: u32 = 20; const MIN_CONTINUATION_EPOCH_SIZE_LOG2: u32 = 18; +/// LogUp mode for prove/verify. This experimental branch defaults to `gkr` so +/// a plain `cli prove` (what `/bench` runs) A/Bs the GKR LogUp path against +/// main's standard path on the same workload. +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +enum CliLogUpMode { + /// Committed batched term + accumulated columns (production path). + Standard, + /// Batch-GKR LogUp: 2 aux columns per table, bus sums proven by GKR. + Gkr, +} + /// Read a file into a buffer aligned for `rkyv::from_bytes`. A plain /// `Vec` from `std::fs::read` is align-1 by the type system even though /// the allocator happens to return well-aligned memory in practice — read @@ -180,6 +191,12 @@ enum Commands { #[arg(long, conflicts_with = "continuations")] elements: bool, + /// LogUp mode. Defaults to `gkr` on this experimental branch (see + /// `CliLogUpMode`); pass `--logup-mode standard` for the production + /// path. Not supported together with --continuations. + #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr, conflicts_with = "continuations")] + logup_mode: CliLogUpMode, + /// Prove with continuations (split execution into epochs; flat peak memory) #[arg(long)] continuations: bool, @@ -213,6 +230,11 @@ enum Commands { #[arg(long)] time: bool, + /// LogUp mode the proof was generated with (must match `prove`'s). + /// Defaults to `gkr`, matching this branch's `prove` default. + #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr, conflicts_with = "continuations")] + logup_mode: CliLogUpMode, + /// Verify a continuation proof bundle (produced by `prove --continuations`) #[arg(long)] continuations: bool, @@ -262,6 +284,7 @@ fn main() -> ExitCode { time, cycles, elements, + logup_mode, continuations, epoch_size_log2, } => { @@ -276,7 +299,16 @@ fn main() -> ExitCode { cycles, ) } else { - cmd_prove(elf, output, private_input, blowup, time, cycles, elements) + cmd_prove( + elf, + output, + private_input, + blowup, + time, + cycles, + elements, + logup_mode, + ) } } Commands::Verify { @@ -284,12 +316,13 @@ fn main() -> ExitCode { elf, blowup, time, + logup_mode, continuations, } => { if continuations { cmd_verify_continuation(proof, elf, blowup, time) } else { - cmd_verify(proof, elf, blowup, time) + cmd_verify(proof, elf, blowup, time, logup_mode) } } Commands::CountElements { elf, private_input } => cmd_count_elements(elf, private_input), @@ -542,6 +575,7 @@ fn cmd_execute( ExitCode::SUCCESS } +#[allow(clippy::too_many_arguments)] fn cmd_prove( elf_path: PathBuf, output_path: PathBuf, @@ -550,6 +584,7 @@ fn cmd_prove( time: bool, cycles: bool, elements: bool, + logup_mode: CliLogUpMode, ) -> ExitCode { eprintln!("Reading ELF file..."); let elf_data = match std::fs::read(&elf_path) { @@ -614,21 +649,57 @@ fn cmd_prove( } }; eprintln!( - "Generating proof (blowup={blowup}, queries={})...", + "Generating proof (blowup={blowup}, queries={}, logup={logup_mode:?})...", opts.fri_number_of_queries ); - let proof = prover::prove_with_options_and_inputs( - &elf_data, - &private_inputs, - &opts, - &Default::default(), - ); - let prove_elapsed = start.elapsed(); - let proof = match proof { - Ok(proof) => proof, - Err(e) => { - eprintln!("Proof generation failed: {}", e); - return ExitCode::FAILURE; + // Prove, timing the proof generation only (serialization happens after + // `prove_elapsed` is taken, identically in both modes). + let (prove_elapsed, bytes) = match logup_mode { + CliLogUpMode::Standard => { + let proof = prover::prove_with_options_and_inputs( + &elf_data, + &private_inputs, + &opts, + &Default::default(), + ); + let prove_elapsed = start.elapsed(); + let proof = match proof { + Ok(proof) => proof, + Err(e) => { + eprintln!("Proof generation failed: {}", e); + return ExitCode::FAILURE; + } + }; + match rkyv::to_bytes::(&proof) { + Ok(b) => (prove_elapsed, b), + Err(e) => { + eprintln!("Failed to serialize proof: {}", e); + return ExitCode::FAILURE; + } + } + } + CliLogUpMode::Gkr => { + let proof = prover::prove_gkr_with_options_and_inputs( + &elf_data, + &private_inputs, + &opts, + &Default::default(), + ); + let prove_elapsed = start.elapsed(); + let proof = match proof { + Ok(proof) => proof, + Err(e) => { + eprintln!("Proof generation failed: {}", e); + return ExitCode::FAILURE; + } + }; + match rkyv::to_bytes::(&proof) { + Ok(b) => (prove_elapsed, b), + Err(e) => { + eprintln!("Failed to serialize proof: {}", e); + return ExitCode::FAILURE; + } + } } }; @@ -642,14 +713,6 @@ fn cmd_prove( }; let mut writer = BufWriter::new(file); - let bytes = match rkyv::to_bytes::(&proof) { - Ok(b) => b, - Err(e) => { - eprintln!("Failed to serialize proof: {}", e); - return ExitCode::FAILURE; - } - }; - if let Err(e) = writer.write_all(&bytes) { eprintln!("Failed to write proof: {}", e); return ExitCode::FAILURE; @@ -674,7 +737,13 @@ fn cmd_prove( ExitCode::SUCCESS } -fn cmd_verify(proof_path: PathBuf, elf_path: PathBuf, blowup: u8, time: bool) -> ExitCode { +fn cmd_verify( + proof_path: PathBuf, + elf_path: PathBuf, + blowup: u8, + time: bool, + logup_mode: CliLogUpMode, +) -> ExitCode { eprintln!("Reading ELF file..."); let elf_data = match std::fs::read(&elf_path) { Ok(data) => data, @@ -693,16 +762,7 @@ fn cmd_verify(proof_path: PathBuf, elf_path: PathBuf, blowup: u8, time: bool) -> } }; - let proof: VmProof = match rkyv::from_bytes::(&proof_bytes) { - Ok(p) => p, - Err(e) => { - eprintln!("Failed to deserialize proof: {}", e); - return ExitCode::FAILURE; - } - }; - - eprintln!("Verifying proof..."); - let start = Instant::now(); + eprintln!("Verifying proof (logup={logup_mode:?})..."); let opts = match GoldilocksCubicProofOptions::with_blowup(blowup) { Ok(opts) => opts, Err(e) => { @@ -710,8 +770,39 @@ fn cmd_verify(proof_path: PathBuf, elf_path: PathBuf, blowup: u8, time: bool) -> return ExitCode::FAILURE; } }; - let result = prover::verify_with_options(&proof, &elf_data, &opts, None, None); - let verify_elapsed = start.elapsed(); + let start = Instant::now(); + let (result, verify_elapsed) = match logup_mode { + CliLogUpMode::Standard => { + let proof: VmProof = + match rkyv::from_bytes::(&proof_bytes) { + Ok(p) => p, + Err(e) => { + eprintln!("Failed to deserialize proof: {}", e); + return ExitCode::FAILURE; + } + }; + let result = prover::verify_with_options(&proof, &elf_data, &opts, None, None); + (result, start.elapsed()) + } + CliLogUpMode::Gkr => { + let proof: prover::GkrVmProof = match rkyv::from_bytes::< + prover::GkrVmProof, + rkyv::rancor::Error, + >(&proof_bytes) + { + Ok(p) => p, + Err(e) => { + eprintln!( + "Failed to deserialize proof (was it proven with --logup-mode gkr?): {}", + e + ); + return ExitCode::FAILURE; + } + }; + let result = prover::verify_gkr_with_options(&proof, &elf_data, &opts); + (result, start.elapsed()) + } + }; let result = match result { Ok(valid) => valid, Err(e) => { From d4ceb2e5c397312bc4e33d849ded92fdf893dff1 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 14:59:22 -0300 Subject: [PATCH 09/20] perf(gkr): parallelize the batch sumcheck and cut redundant work The +28.5% prove-time regression on the ethrex-20tx bench was dominated by serial work in the GKR phase (which runs between multi_prove's parallel phases and inherits none of their table parallelism): - batch sumcheck round loop: per-instance round evals now run in parallel across instances, the per-round gate/eq sums use parallel reductions above 2^10 pairs, and the post-challenge claim/fold updates run in parallel across instances (field ops are exact, so any reduction order is value-identical; the transcript sequence is untouched). Previously only fold_table was parallel. - PerInstanceTables leaf-layer even/odd splits parallelized. - compute_lagrange_kernel: doubling algorithm (1 mul + 1 sub per output pair, O(N)) replacing the per-bit full-table passes (O(n*N) ext mults, ~20x the necessary work at production sizes); values are identical (regression test vs the per-bit product definition). - the GKR path no longer materializes column-major clones of the main segment (previously 3x per table): leaf fractions, the column-claim inner products (now ONE row-major pass computing all K claims with per-chunk partial sums), and the aux batched column all read borrowed row-major rows in place via a new BusValue::accumulate_fingerprint_row. - hoisted the per-instance-per-round inversion of 2; instruments spans for the GKR phases (r1_gkr_trees / r1_gkr_batch_prove / r1_gkr_finalize). Local ethrex-5tx phase timings (10-core M-series): trees 3.13s -> 2.12s, batch prove 4.63s -> 2.26s, finalize 1.39s -> 0.25s, GKR aux build 1.90s -> 0.38s. The wins should be larger on the many-core bench box, where the serial batch phase was the Amdahl bottleneck. --- crypto/stark/src/gkr.rs | 526 +++++++++++++++++----------- crypto/stark/src/lagrange_kernel.rs | 80 +++-- crypto/stark/src/logup_gkr.rs | 113 +++--- crypto/stark/src/lookup.rs | 59 ++++ crypto/stark/src/prover.rs | 16 +- prover/src/lib.rs | 18 + 6 files changed, 521 insertions(+), 291 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index 0838630ed..5ae11d179 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -1431,85 +1431,103 @@ pub fn gkr_prove_batch( // For each active instance, build bookkeeping tables. // We run the sumcheck manually, combining round polynomials across instances. - let mut per_instance_tables: Vec> = active_instances - .iter() - .map(|&i| { - let tree_layer_idx = n_remaining - 1; - let child = &layers_per_instance[i][tree_layer_idx]; + // Even/odd split of a layer vector: parallel for the large leaf + // layers (element-wise clones dominate the batch setup otherwise). + fn split_even_odd( + src: &[FieldElement], + parent_size: usize, + ) -> (Vec>, Vec>) { + #[cfg(feature = "parallel")] + if parent_size >= 1 << 10 { + let even: Vec<_> = (0..parent_size) + .into_par_iter() + .map(|j| src[2 * j].clone()) + .collect(); + let odd: Vec<_> = (0..parent_size) + .into_par_iter() + .map(|j| src[2 * j + 1].clone()) + .collect(); + return (even, odd); + } + let even: Vec<_> = (0..parent_size).map(|j| src[2 * j].clone()).collect(); + let odd: Vec<_> = (0..parent_size).map(|j| src[2 * j + 1].clone()).collect(); + (even, odd) + } - let parent_size = match child { - Layer::LogUpGeneric { denominators, .. } - | Layer::LogUpSingles { denominators } => denominators.len() / 2, - }; + let build_instance_tables = |&i: &usize| { + let tree_layer_idx = n_remaining - 1; + let child = &layers_per_instance[i][tree_layer_idx]; - let (nl_table, nr_table, is_singles) = match child { - Layer::LogUpGeneric { numerators, .. } => { - let nl: Vec<_> = (0..parent_size) - .map(|j| numerators[2 * j].clone()) - .collect(); - let nr: Vec<_> = (0..parent_size) - .map(|j| numerators[2 * j + 1].clone()) - .collect(); - (nl, nr, false) - } - Layer::LogUpSingles { .. } => { - // Singles: numerators are all 1 - (vec![], vec![], true) - } - }; + let parent_size = match child { + Layer::LogUpGeneric { denominators, .. } + | Layer::LogUpSingles { denominators } => denominators.len() / 2, + }; - let denominators = match child { - Layer::LogUpGeneric { denominators, .. } - | Layer::LogUpSingles { denominators } => denominators, - }; + let (nl_table, nr_table, is_singles) = match child { + Layer::LogUpGeneric { numerators, .. } => { + let (nl, nr) = split_even_odd(numerators, parent_size); + (nl, nr, false) + } + Layer::LogUpSingles { .. } => { + // Singles: numerators are all 1 + (vec![], vec![], true) + } + }; - let dl_table: Vec<_> = (0..parent_size) - .map(|j| denominators[2 * j].clone()) - .collect(); - let dr_table: Vec<_> = (0..parent_size) - .map(|j| denominators[2 * j + 1].clone()) - .collect(); + let denominators = match child { + Layer::LogUpGeneric { denominators, .. } + | Layer::LogUpSingles { denominators } => denominators, + }; - let my_parent_num_vars = parent_num_vars_by_instance - [active_instances.iter().position(|&x| x == i).unwrap()]; - // current_point must have at least parent_num_vars coordinates - // (it grows by max_parent_vars+1 each layer, matching the tree doubling). - debug_assert!( - current_point.len() >= my_parent_num_vars, - "current_point.len()={} < parent_num_vars={}", - current_point.len(), - my_parent_num_vars - ); - let inst_point = instance_eval_point(¤t_point, my_parent_num_vars); - let use_svo = my_parent_num_vars >= SVO_THRESHOLD; - let svo_suffix_len = if use_svo { my_parent_num_vars / 2 } else { 0 }; - - let (eq_table, eq_prefix, eq_suffix) = if use_svo { - let suffix = compute_eq_evals(&inst_point[..svo_suffix_len]); - let prefix = - compute_eq_evals(&inst_point[svo_suffix_len..my_parent_num_vars]); - (Vec::new(), prefix, suffix) - } else { - (compute_eq_evals(&inst_point), Vec::new(), Vec::new()) - }; + let (dl_table, dr_table) = split_even_odd(denominators, parent_size); - PerInstanceTables { - nl_table, - nr_table, - dl_table, - dr_table, - eq_table, - eq_correction: FieldElement::one(), - is_singles, - parent_num_vars: my_parent_num_vars, - instance_point: inst_point, - use_svo, - svo_suffix_len, - eq_prefix, - eq_suffix, - } - }) + let my_parent_num_vars = parent_num_vars_by_instance + [active_instances.iter().position(|&x| x == i).unwrap()]; + // current_point must have at least parent_num_vars coordinates + // (it grows by max_parent_vars+1 each layer, matching the tree doubling). + debug_assert!( + current_point.len() >= my_parent_num_vars, + "current_point.len()={} < parent_num_vars={}", + current_point.len(), + my_parent_num_vars + ); + let inst_point = instance_eval_point(¤t_point, my_parent_num_vars); + let use_svo = my_parent_num_vars >= SVO_THRESHOLD; + let svo_suffix_len = if use_svo { my_parent_num_vars / 2 } else { 0 }; + + let (eq_table, eq_prefix, eq_suffix) = if use_svo { + let suffix = compute_eq_evals(&inst_point[..svo_suffix_len]); + let prefix = compute_eq_evals(&inst_point[svo_suffix_len..my_parent_num_vars]); + (Vec::new(), prefix, suffix) + } else { + (compute_eq_evals(&inst_point), Vec::new(), Vec::new()) + }; + + PerInstanceTables { + nl_table, + nr_table, + dl_table, + dr_table, + eq_table, + eq_correction: FieldElement::one(), + is_singles, + parent_num_vars: my_parent_num_vars, + instance_point: inst_point, + use_svo, + svo_suffix_len, + eq_prefix, + eq_suffix, + } + }; + + #[cfg(feature = "parallel")] + let mut per_instance_tables: Vec> = active_instances + .par_iter() + .map(build_instance_tables) .collect(); + #[cfg(not(feature = "parallel"))] + let mut per_instance_tables: Vec> = + active_instances.iter().map(build_instance_tables).collect(); let mut round_polys = Vec::with_capacity(max_parent_vars); let mut challenges = Vec::with_capacity(max_parent_vars); @@ -1524,67 +1542,51 @@ pub fn gkr_prove_batch( sum }; - // Per-instance round polynomial evals, used to update combined_claims - // via O(1) interpolation instead of O(N) table scan. - let mut per_instance_evals: Vec<[FieldElement; 4]> = vec![ - [ - FieldElement::zero(), - FieldElement::zero(), - FieldElement::zero(), - FieldElement::zero() - ]; - per_instance_tables.len() - ]; + // Constants hoisted out of the round loop (the old code inverted 2 + // per instance per round). + let inv_two = FieldElement::::from(2u64) + .inv() + .expect("2 is invertible"); + /// Above this pair count the per-round gate/eq sums and folds use + /// parallel reductions (field ops are exact, so any reduction + /// order is value-identical). + #[cfg(feature = "parallel")] + const PAR_SUM_THRESHOLD: usize = 1 << 10; for round_idx in 0..max_parent_vars { - // For each active instance, compute the round poly contribution - let mut batch_s0 = FieldElement::::zero(); - let mut batch_s1 = FieldElement::::zero(); - let mut batch_s2 = FieldElement::::zero(); - let mut batch_s3 = FieldElement::::zero(); - let mut alpha_pow = FieldElement::::one(); - - for (idx, tables) in per_instance_tables.iter_mut().enumerate() { + // Per-instance round evals S_i(t) at t = 0,1,2,3. Instances are + // independent given the (read-only) combined claims, so this — + // the dominant work of the whole batch prove — runs in + // parallel across instances, with parallel inner sums for the + // large early rounds. + let compute_instance = |idx: usize, + tables: &mut PerInstanceTables| + -> [FieldElement; 4] { let n_unused = max_parent_vars - tables.parent_num_vars; if round_idx < n_unused { - // This instance hasn't started yet — constant polynomial = claim/2 - let half_claim = - &combined_claims[idx] * &FieldElement::::from(2u64).inv().unwrap(); - // S(0) = S(1) = half_claim, S(2) = S(3) = half_claim (constant) - per_instance_evals[idx] = [ - half_claim.clone(), + // This instance hasn't started yet — constant polynomial = claim/2. + let half_claim = &combined_claims[idx] * &inv_two; + return [ half_claim.clone(), half_claim.clone(), half_claim.clone(), + half_claim, ]; - batch_s0 = &batch_s0 + &(&alpha_pow * &half_claim); - batch_s1 = &batch_s1 + &(&alpha_pow * &half_claim); - batch_s2 = &batch_s2 + &(&alpha_pow * &half_claim); - batch_s3 = &batch_s3 + &(&alpha_pow * &half_claim); - alpha_pow = &alpha_pow * &sumcheck_alpha; - continue; } let instance_round = round_idx - n_unused; let half = tables.nl_table.len().max(tables.dl_table.len()) / 2; if half == 0 { - // Already reduced to constant - let half_claim = - &combined_claims[idx] * &FieldElement::::from(2u64).inv().unwrap(); - per_instance_evals[idx] = [ - half_claim.clone(), + // Already reduced to constant. + let half_claim = &combined_claims[idx] * &inv_two; + return [ half_claim.clone(), half_claim.clone(), half_claim.clone(), + half_claim, ]; - batch_s0 = &batch_s0 + &(&alpha_pow * &half_claim); - batch_s1 = &batch_s1 + &(&alpha_pow * &half_claim); - batch_s2 = &batch_s2 + &(&alpha_pow * &half_claim); - batch_s3 = &batch_s3 + &(&alpha_pow * &half_claim); - alpha_pow = &alpha_pow * &sumcheck_alpha; - continue; } // Eq polynomial factoring (same as single-instance prover). @@ -1628,73 +1630,127 @@ pub fn gkr_prove_batch( }; // Compute h(0) and h(2) using SVO or standard path. - let (raw_h0, raw_h2) = - if tables.use_svo && instance_round < tables.svo_suffix_len { - // SVO suffix round: nested eq_suffix × (eq_prefix × gate) loop - let suffix_half = tables.eq_suffix.len() / 2; - let prefix_size = tables.eq_prefix.len(); - - // Pre-halve eq_suffix - for j in 0..suffix_half { - tables.eq_suffix[j] = - &tables.eq_suffix[2 * j] + &tables.eq_suffix[2 * j + 1]; + let (raw_h0, raw_h2) = if tables.use_svo + && instance_round < tables.svo_suffix_len + { + // SVO suffix round: nested eq_suffix × (eq_prefix × gate) loop + let suffix_half = tables.eq_suffix.len() / 2; + let prefix_size = tables.eq_prefix.len(); + + // Pre-halve eq_suffix + for j in 0..suffix_half { + tables.eq_suffix[j] = + &tables.eq_suffix[2 * j] + &tables.eq_suffix[2 * j + 1]; + } + tables.eq_suffix.truncate(suffix_half); + + // Eq mutations are done; sum over an immutable view so the + // gate closures can run in parallel. + let view: &PerInstanceTables = tables; + let suffix_sum = |suffix_idx: usize| -> (FieldElement, FieldElement) { + let eq_s = &view.eq_suffix[suffix_idx]; + let mut ch0 = FieldElement::::zero(); + let mut ch2 = FieldElement::::zero(); + #[allow(clippy::needless_range_loop)] + for prefix_idx in 0..prefix_size { + let j = prefix_idx * suffix_half + suffix_idx; + let eq_p = &view.eq_prefix[prefix_idx]; + let [g0, g2] = if view.is_singles { + gate_singles(view, j) + } else { + gate_generic(view, j) + }; + ch0 = &ch0 + &(eq_p * &g0); + ch2 = &ch2 + &(eq_p * &g2); } - tables.eq_suffix.truncate(suffix_half); + (eq_s * &ch0, eq_s * &ch2) + }; + #[cfg(feature = "parallel")] + if suffix_half * prefix_size >= PAR_SUM_THRESHOLD { + (0..suffix_half).into_par_iter().map(suffix_sum).reduce( + || (FieldElement::::zero(), FieldElement::::zero()), + |(a0, a2), (b0, b2)| (&a0 + &b0, &a2 + &b2), + ) + } else { let mut h0 = FieldElement::::zero(); let mut h2 = FieldElement::::zero(); for suffix_idx in 0..suffix_half { - let eq_s = &tables.eq_suffix[suffix_idx]; - let mut ch0 = FieldElement::::zero(); - let mut ch2 = FieldElement::::zero(); - #[allow(clippy::needless_range_loop)] - for prefix_idx in 0..prefix_size { - let j = prefix_idx * suffix_half + suffix_idx; - let eq_p = &tables.eq_prefix[prefix_idx]; - let [g0, g2] = if tables.is_singles { - gate_singles(tables, j) - } else { - gate_generic(tables, j) - }; - ch0 = &ch0 + &(eq_p * &g0); - ch2 = &ch2 + &(eq_p * &g2); - } - h0 = &h0 + &(eq_s * &ch0); - h2 = &h2 + &(eq_s * &ch2); + let (c0, c2) = suffix_sum(suffix_idx); + h0 = &h0 + &c0; + h2 = &h2 + &c2; } (h0, h2) - } else { - // Standard path (non-SVO, or SVO prefix rounds after suffix is exhausted) - if tables.use_svo && instance_round == tables.svo_suffix_len { - // Transition: absorb remaining eq_suffix into eq_correction - // and switch to eq_prefix as eq_table. - debug_assert_eq!(tables.eq_suffix.len(), 1); - tables.eq_correction = &tables.eq_correction * &tables.eq_suffix[0]; - tables.eq_table = std::mem::take(&mut tables.eq_prefix); - tables.use_svo = false; + } + #[cfg(not(feature = "parallel"))] + { + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for suffix_idx in 0..suffix_half { + let (c0, c2) = suffix_sum(suffix_idx); + h0 = &h0 + &c0; + h2 = &h2 + &c2; } + (h0, h2) + } + } else { + // Standard path (non-SVO, or SVO prefix rounds after suffix is exhausted) + if tables.use_svo && instance_round == tables.svo_suffix_len { + // Transition: absorb remaining eq_suffix into eq_correction + // and switch to eq_prefix as eq_table. + debug_assert_eq!(tables.eq_suffix.len(), 1); + tables.eq_correction = &tables.eq_correction * &tables.eq_suffix[0]; + tables.eq_table = std::mem::take(&mut tables.eq_prefix); + tables.use_svo = false; + } + + // Pre-halve eq_table + for j in 0..half { + tables.eq_table[j] = + &tables.eq_table[2 * j] + &tables.eq_table[2 * j + 1]; + } + tables.eq_table.truncate(half); + + // Eq mutations are done; sum over an immutable view. + let view: &PerInstanceTables = tables; + let pair_sum = |j: usize| -> (FieldElement, FieldElement) { + let eq_rem = &view.eq_table[j]; + let [g0, g2] = if view.is_singles { + gate_singles(view, j) + } else { + gate_generic(view, j) + }; + (eq_rem * &g0, eq_rem * &g2) + }; - // Pre-halve eq_table + #[cfg(feature = "parallel")] + if half >= PAR_SUM_THRESHOLD { + (0..half).into_par_iter().map(pair_sum).reduce( + || (FieldElement::::zero(), FieldElement::::zero()), + |(a0, a2), (b0, b2)| (&a0 + &b0, &a2 + &b2), + ) + } else { + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); for j in 0..half { - tables.eq_table[j] = - &tables.eq_table[2 * j] + &tables.eq_table[2 * j + 1]; + let (c0, c2) = pair_sum(j); + h0 = &h0 + &c0; + h2 = &h2 + &c2; } - tables.eq_table.truncate(half); - + (h0, h2) + } + #[cfg(not(feature = "parallel"))] + { let mut h0 = FieldElement::::zero(); let mut h2 = FieldElement::::zero(); for j in 0..half { - let eq_rem = &tables.eq_table[j]; - let [g0, g2] = if tables.is_singles { - gate_singles(tables, j) - } else { - gate_generic(tables, j) - }; - h0 = &h0 + &(eq_rem * &g0); - h2 = &h2 + &(eq_rem * &g2); + let (c0, c2) = pair_sum(j); + h0 = &h0 + &c0; + h2 = &h2 + &c2; } (h0, h2) - }; + } + }; // Apply eq_correction let total_h0 = &tables.eq_correction * &raw_h0; @@ -1718,11 +1774,33 @@ pub fn gkr_prove_batch( - &FieldElement::::from(2u64); let s3 = &eq_at_3 * &h3; - per_instance_evals[idx] = [s0.clone(), s1.clone(), s2.clone(), s3.clone()]; - batch_s0 = &batch_s0 + &(&alpha_pow * &s0); - batch_s1 = &batch_s1 + &(&alpha_pow * &s1); - batch_s2 = &batch_s2 + &(&alpha_pow * &s2); - batch_s3 = &batch_s3 + &(&alpha_pow * &s3); + [s0, s1, s2, s3] + }; + + #[cfg(feature = "parallel")] + let per_instance_evals: Vec<[FieldElement; 4]> = per_instance_tables + .par_iter_mut() + .enumerate() + .map(|(idx, tables)| compute_instance(idx, tables)) + .collect(); + #[cfg(not(feature = "parallel"))] + let per_instance_evals: Vec<[FieldElement; 4]> = per_instance_tables + .iter_mut() + .enumerate() + .map(|(idx, tables)| compute_instance(idx, tables)) + .collect(); + + // α-combine the per-instance evals (sequential; O(instances)). + let mut batch_s0 = FieldElement::::zero(); + let mut batch_s1 = FieldElement::::zero(); + let mut batch_s2 = FieldElement::::zero(); + let mut batch_s3 = FieldElement::::zero(); + let mut alpha_pow = FieldElement::::one(); + for [s0, s1, s2, s3] in &per_instance_evals { + batch_s0 = &batch_s0 + &(&alpha_pow * s0); + batch_s1 = &batch_s1 + &(&alpha_pow * s1); + batch_s2 = &batch_s2 + &(&alpha_pow * s2); + batch_s3 = &batch_s3 + &(&alpha_pow * s3); alpha_pow = &alpha_pow * &sumcheck_alpha; } @@ -1738,56 +1816,76 @@ pub fn gkr_prove_batch( _round_combined_claim = round_poly.evaluate(&challenge); // Update per-instance: fold tables, update eq_correction, and - // update combined_claims via O(1) polynomial evaluation (not O(N) table scan). - for (idx, tables) in per_instance_tables.iter_mut().enumerate() { - // Update combined_claims from saved per-instance round poly evals. - // S_i(challenge) via degree-3 Lagrange interpolation at {0,1,2,3}. - let [ref si0, ref si1, ref si2, ref si3] = per_instance_evals[idx]; - let instance_poly = - RoundPoly::new(vec![si0.clone(), si1.clone(), si2.clone(), si3.clone()]); - combined_claims[idx] = instance_poly.evaluate(&challenge); - - let n_unused = max_parent_vars - tables.parent_num_vars; - if round_idx < n_unused || tables.dl_table.len() / 2 == 0 { - continue; - } - - let instance_round = round_idx - n_unused; - let half = tables.dl_table.len() / 2; - - // Update eq_correction using instance-specific eval point - let r_round = &tables.instance_point[instance_round]; - let one = FieldElement::::one(); - let eq_update = - &(r_round * &challenge) + &(&(&one - r_round) * &(&one - &challenge)); - tables.eq_correction = &tables.eq_correction * &eq_update; + // update combined_claims via O(1) polynomial evaluation (not + // O(N) table scan). Instances are independent — parallel. + let update_instance = + |tables: &mut PerInstanceTables, + claim: &mut FieldElement, + evals: &[FieldElement; 4]| { + // S_i(challenge) via degree-3 Lagrange interpolation at {0,1,2,3}. + let [si0, si1, si2, si3] = evals; + let instance_poly = RoundPoly::new(vec![ + si0.clone(), + si1.clone(), + si2.clone(), + si3.clone(), + ]); + *claim = instance_poly.evaluate(&challenge); - // Fold gate tables - let fold_table = |table: &mut Vec>| { - #[cfg(feature = "parallel")] - if half >= 256 { - let folded: Vec> = table - .par_chunks(2) - .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) - .collect(); - *table = folded; + let n_unused = max_parent_vars - tables.parent_num_vars; + if round_idx < n_unused || tables.dl_table.len() / 2 == 0 { return; } - for j in 0..half { - let left = &table[2 * j]; - let right = &table[2 * j + 1]; - table[j] = left + &(&challenge * &(right - left)); + + let instance_round = round_idx - n_unused; + let half = tables.dl_table.len() / 2; + + // Update eq_correction using instance-specific eval point + let r_round = &tables.instance_point[instance_round]; + let one = FieldElement::::one(); + let eq_update = + &(r_round * &challenge) + &(&(&one - r_round) * &(&one - &challenge)); + tables.eq_correction = &tables.eq_correction * &eq_update; + + // Fold gate tables + let fold_table = |table: &mut Vec>| { + #[cfg(feature = "parallel")] + if half >= 256 { + let folded: Vec> = table + .par_chunks(2) + .map(|pair| &pair[0] + &(&challenge * &(&pair[1] - &pair[0]))) + .collect(); + *table = folded; + return; + } + for j in 0..half { + let left = &table[2 * j]; + let right = &table[2 * j + 1]; + table[j] = left + &(&challenge * &(right - left)); + } + table.truncate(half); + }; + + if !tables.is_singles { + fold_table(&mut tables.nl_table); + fold_table(&mut tables.nr_table); } - table.truncate(half); + fold_table(&mut tables.dl_table); + fold_table(&mut tables.dr_table); }; - if !tables.is_singles { - fold_table(&mut tables.nl_table); - fold_table(&mut tables.nr_table); - } - fold_table(&mut tables.dl_table); - fold_table(&mut tables.dr_table); - } + #[cfg(feature = "parallel")] + per_instance_tables + .par_iter_mut() + .zip(combined_claims.par_iter_mut()) + .zip(per_instance_evals.par_iter()) + .for_each(|((tables, claim), evals)| update_instance(tables, claim, evals)); + #[cfg(not(feature = "parallel"))] + per_instance_tables + .iter_mut() + .zip(combined_claims.iter_mut()) + .zip(per_instance_evals.iter()) + .for_each(|((tables, claim), evals)| update_instance(tables, claim, evals)); round_polys.push(round_poly); challenges.push(challenge); diff --git a/crypto/stark/src/lagrange_kernel.rs b/crypto/stark/src/lagrange_kernel.rs index bd5a832fa..099b10d9a 100644 --- a/crypto/stark/src/lagrange_kernel.rs +++ b/crypto/stark/src/lagrange_kernel.rs @@ -17,50 +17,41 @@ use rayon::prelude::*; /// claims back to committed univariate traces. It satisfies the partition-of-unity /// property: sum_{i=0}^{N-1} s[i] = 1 for any r. /// -/// Uses the butterfly algorithm in O(N) time and O(N) space. +/// Uses the doubling algorithm: each coordinate doubles the table with one +/// multiplication and one subtraction per output pair, O(N) total work (the +/// per-bit full-table passes would be O(n·N)). pub fn compute_lagrange_kernel(r: &[FieldElement]) -> Vec> { let n = r.len(); - let big_n = 1usize << n; - let one = FieldElement::::one(); - - let mut s = vec![one.clone(); big_n]; - - #[allow(clippy::needless_range_loop)] - for j in 0..n { + let mut s = vec![FieldElement::::one()]; + + // Process coordinates from LAST to FIRST, interleaving each step: + // next[2t+1] = s[t]·r_j + // next[2t] = s[t] − next[2t+1] (= s[t]·(1 − r_j)) + // After processing r_{n-1}..r_0, bit l of the index selects the r_l + // factor — identical values (and order) to the per-bit product + // definition s[i] = ∏_l [r_l·b_l(i) + (1−r_l)·(1−b_l(i))]. + for j in (0..n).rev() { let rj = &r[j]; - let one_minus_rj = &one - rj; + let m = s.len(); + let mut next = vec![FieldElement::::zero(); 2 * m]; #[cfg(feature = "parallel")] - { - if big_n >= 1024 { - s.par_iter_mut().enumerate().for_each(|(i, s_i)| { - if (i >> j) & 1 == 1 { - *s_i = &*s_i * rj; - } else { - *s_i = &*s_i * &one_minus_rj; - } + if m >= 1024 { + next.par_chunks_mut(2) + .zip(s.par_iter()) + .for_each(|(pair, s_t)| { + pair[1] = s_t * rj; + pair[0] = s_t - &pair[1]; }); - } else { - for i in 0..big_n { - if (i >> j) & 1 == 1 { - s[i] = &s[i] * rj; - } else { - s[i] = &s[i] * &one_minus_rj; - } - } - } + s = next; + continue; } - #[cfg(not(feature = "parallel"))] - { - for i in 0..big_n { - if (i >> j) & 1 == 1 { - s[i] = &s[i] * rj; - } else { - s[i] = &s[i] * &one_minus_rj; - } - } + for (t, s_t) in s.iter().enumerate() { + next[2 * t + 1] = s_t * rj; + next[2 * t] = s_t - &next[2 * t + 1]; } + s = next; } s @@ -213,6 +204,25 @@ mod tests { assert_eq!(sum, FE::one()); } + /// The doubling algorithm must match the per-bit product definition + /// s[i] = ∏_l [r_l·b_l(i) + (1−r_l)·(1−b_l(i))] exactly. + #[test] + fn test_kernel_doubling_matches_naive_definition() { + for &(n, seed) in &[(1usize, 3u64), (2, 7), (3, 11), (5, 13), (8, 17)] { + let r: Vec = (0..n).map(|l| FE::from(seed + 31 * l as u64)).collect(); + let fast = compute_lagrange_kernel(&r); + let one = FE::one(); + for (i, v) in fast.iter().enumerate() { + let mut expected = FE::one(); + for (l, r_l) in r.iter().enumerate() { + let bit = (i >> l) & 1; + expected *= if bit == 1 { *r_l } else { &one - r_l }; + } + assert_eq!(*v, expected, "kernel mismatch at n={n}, i={i}"); + } + } + } + #[test] fn test_lagrange_kernel_partition_of_unity() { // For any random r, the sum of kernel values must equal 1. diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index f88fda228..b06b9ceac 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -25,8 +25,9 @@ use math::field::{ use rayon::prelude::*; use crate::gkr::{Layer, gen_layers}; -use crate::lagrange_kernel::{compute_lagrange_kernel, eval_mle_base_with_kernel}; +use crate::lagrange_kernel::compute_lagrange_kernel; use crate::lookup::{BusInteraction, LOGUP_CHALLENGE_ALPHA, PackingShifts, compute_alpha_powers}; +use crate::table::Table; use crate::trace::TraceTable; // ============================================================================= @@ -101,11 +102,11 @@ pub fn extract_column_indices(interactions: &[BusInteraction]) -> Vec { // ============================================================================= /// The fingerprint `z − (bus_id + Σ αⁱ·elementᵢ)` for one interaction at one -/// row. `α⁰ = 1`: the bus-id term is embedded directly, no multiply. +/// row (a borrowed row-major slice). `α⁰ = 1`: the bus-id term is embedded +/// directly, no multiply. fn compute_fingerprint_at_row( interaction: &BusInteraction, - main_segment_cols: &[Vec>], - row: usize, + row: &[FieldElement], z: &FieldElement, alpha_powers: &[FieldElement], shifts: &PackingShifts, @@ -117,9 +118,8 @@ where let mut lc = FieldElement::::from(interaction.bus_id); let mut alpha_offset = 1; for bv in &interaction.values { - alpha_offset += bv.accumulate_fingerprint( - main_segment_cols, - row, + alpha_offset += bv.accumulate_fingerprint_row( + |col| &row[col], alpha_powers, alpha_offset, &mut lc, @@ -144,7 +144,7 @@ where /// Returns `(numerators, denominators)`, each of length `trace_len`. pub fn compute_logup_leaf_fractions( interactions: &[BusInteraction], - main_segment_cols: &[Vec>], + main: &Table, trace_len: usize, challenges: &[FieldElement], ) -> (Vec>, Vec>) @@ -167,19 +167,15 @@ where let alpha_powers = compute_alpha_powers(alpha, max_bus_elements); let shifts = PackingShifts::::new(); - let leaf_at_row = |row: usize| { + // Rows are read in place from the row-major main table (borrowed slices, + // no column-major transpose/clone of the segment). + let leaf_at_row = |row_idx: usize| { + let row = main.get_row(row_idx); let mut running_n = FieldElement::::zero(); let mut running_d = FieldElement::::one(); for inter in interactions { - let fp = compute_fingerprint_at_row( - inter, - main_segment_cols, - row, - z, - &alpha_powers, - &shifts, - ); - let m = inter.multiplicity.evaluate_at_row(main_segment_cols, row); + let fp = compute_fingerprint_at_row(inter, row, z, &alpha_powers, &shifts); + let m = inter.multiplicity.evaluate_from(|col| row[col].clone()); // m·d is F×E (base operand LEFT); the sign resolves as add vs neg. let cross = &m * &running_d; let cross = if inter.is_sender { cross } else { -cross }; @@ -202,7 +198,7 @@ where /// fractions, then pairwise fraction-summation layers up to the root. pub fn compute_logup_layers( interactions: &[BusInteraction], - main_segment_cols: &[Vec>], + main: &Table, trace_len: usize, challenges: &[FieldElement], ) -> Vec> @@ -211,7 +207,7 @@ where E: IsField + Send + Sync, { let (numerators, denominators) = - compute_logup_leaf_fractions(interactions, main_segment_cols, trace_len, challenges); + compute_logup_leaf_fractions(interactions, main, trace_len, challenges); gen_layers(Layer::LogUpGeneric { numerators, denominators, @@ -245,7 +241,7 @@ pub struct LogUpGkrResult { /// use — the aux-trace build recomputes it to keep this result small). pub fn finalize_logup_gkr_result( interactions: &[BusInteraction], - main_segment_cols: &[Vec>], + main: &Table, random_point: Vec>, n_claim: FieldElement, d_claim: FieldElement, @@ -257,18 +253,51 @@ where { let col_indices = extract_column_indices(interactions); let kernel = compute_lagrange_kernel(&random_point); + let k = col_indices.len(); + let num_rows = main.height; + + // All K inner products ⟨kernel, col_j⟩ in ONE row-major pass (per-chunk + // partial accumulators, reduced by field addition — value-identical to + // per-column sums, no column-major transpose of the segment). + let row_partial = |acc: &mut [FieldElement], row_idx: usize| { + let row = main.get_row(row_idx); + let k_r = &kernel[row_idx]; + for (j, &col_idx) in col_indices.iter().enumerate() { + // F×E (base operand LEFT). + acc[j] = &acc[j] + &(&row[col_idx] * k_r); + } + }; #[cfg(feature = "parallel")] - let col_iter = col_indices.into_par_iter(); + let claims: Vec> = (0..num_rows) + .into_par_iter() + .fold( + || vec![FieldElement::::zero(); k], + |mut acc, row_idx| { + row_partial(&mut acc, row_idx); + acc + }, + ) + .reduce( + || vec![FieldElement::::zero(); k], + |mut a, b| { + for (x, y) in a.iter_mut().zip(b) { + *x = &*x + &y; + } + a + }, + ); #[cfg(not(feature = "parallel"))] - let col_iter = col_indices.into_iter(); + let claims: Vec> = { + let mut acc = vec![FieldElement::::zero(); k]; + for row_idx in 0..num_rows { + row_partial(&mut acc, row_idx); + } + acc + }; - let column_claims: Vec<(usize, FieldElement)> = col_iter - .map(|col_idx| { - let claim = eval_mle_base_with_kernel(&main_segment_cols[col_idx], &kernel); - (col_idx, claim) - }) - .collect(); + let column_claims: Vec<(usize, FieldElement)> = + col_indices.into_iter().zip(claims).collect(); LogUpGkrResult { table_contribution, @@ -388,15 +417,16 @@ pub(crate) fn build_gkr_aux_columns( let bridge_offset = &challenges[LOGUP_BRIDGE_OFFSET_IDX]; let kernel = compute_lagrange_kernel(random_point); - let main_segment_cols = trace.columns_main(); - // batched[i] = Σ_j γ^j·col_j[i] + γ^K·l[i] — row-parallel, matches the - // bridge constraint's batched sum exactly. - let batched_at_row = |row: usize| { - let mut acc = &gamma_powers[k] * &kernel[row]; + // batched[i] = Σ_j γ^j·col_j[i] + γ^K·l[i] — row-parallel over borrowed + // row-major rows (no column-major clone of the main segment), matching + // the bridge constraint's batched sum exactly. + let batched_at_row = |row_idx: usize| { + let row = trace.main_table.get_row(row_idx); + let mut acc = &gamma_powers[k] * &kernel[row_idx]; for (j, &col_idx) in column_indices.iter().enumerate() { // col·γ is F×E (base operand LEFT). - acc += &main_segment_cols[col_idx][row] * &gamma_powers[j]; + acc += &row[col_idx] * &gamma_powers[j]; } acc }; @@ -541,7 +571,7 @@ mod tests { fn leaf_fractions_single_sender() { let trace_len = 4; let col0: Vec = [10u64, 20, 30, 40].iter().map(|&v| FE::from(v)).collect(); - let main_segment_cols = vec![col0.clone()]; + let main_segment_cols = Table::from_columns(vec![col0.clone()]); let interactions = vec![BusInteraction::sender( 1u64, Multiplicity::One, @@ -573,7 +603,7 @@ mod tests { let trace_len = 4; let col0: Vec = [5u64, 6, 7, 8].iter().map(|&v| FE::from(v)).collect(); let col1: Vec = [2u64, 0, 1, 3].iter().map(|&v| FE::from(v)).collect(); - let main_segment_cols = vec![col0.clone(), col1.clone()]; + let main_segment_cols = Table::from_columns(vec![col0.clone(), col1.clone()]); let interactions = vec![BusInteraction::receiver( 0u64, Multiplicity::Column(1), @@ -606,7 +636,7 @@ mod tests { let trace_len = 2; let col0: Vec = [10u64, 20].iter().map(|&v| FE::from(v)).collect(); let col1: Vec = [30u64, 40].iter().map(|&v| FE::from(v)).collect(); - let main_segment_cols = vec![col0.clone(), col1.clone()]; + let main_segment_cols = Table::from_columns(vec![col0.clone(), col1.clone()]); let interactions = vec![ BusInteraction::sender(0u64, Multiplicity::One, Packing::Direct.columns(&[0])), BusInteraction::receiver(1u64, Multiplicity::One, Packing::Direct.columns(&[1])), @@ -721,7 +751,10 @@ mod tests { .iter() .map(|&v| FE::from(v)) .collect(); - let main_segment_cols = vec![col0, col1, col2]; + // The standard term-column path takes column-major vectors; the GKR + // path reads the row-major table — same data, both representations. + let columns = vec![col0, col1, col2]; + let main_segment_cols = Table::from_columns(columns.clone()); let interactions = vec![ BusInteraction::sender(3u64, Multiplicity::One, Packing::Direct.columns(&[0])), @@ -735,7 +768,7 @@ mod tests { for inter in &interactions { let terms = crate::lookup::compute_logup_term_column::( &[inter], - &main_segment_cols, + &columns, trace_len, &challenges, "oracle", diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 63232c51a..74cbc08f3 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -699,6 +699,65 @@ impl BusValue { } } + /// Row-accessor variant of [`Self::accumulate_fingerprint`]: reads column + /// values through `get_col` (e.g. a borrowed row-major trace row) instead + /// of column-major vectors. Value-identical, including the zero-skip on + /// `Linear` elements. + pub fn accumulate_fingerprint_row<'a, F, E>( + &self, + get_col: impl Fn(usize) -> &'a FieldElement, + alpha_powers: &[FieldElement], + alpha_offset: usize, + acc: &mut FieldElement, + shifts: &PackingShifts, + ) -> usize + where + F: IsField + IsSubFieldOf + 'a, + E: IsField, + { + match self { + BusValue::Packed { + start_column, + packing, + } => packing.accumulate_fingerprint_with( + *start_column, + get_col, + alpha_powers, + alpha_offset, + acc, + shifts, + ), + BusValue::Linear(terms) => { + let mut result = FieldElement::::zero(); + for term in terms { + match term { + LinearTerm::Column { + coefficient, + column, + } => { + let coeff = FieldElement::::from(*coefficient); + result += get_col(*column) * coeff; + } + LinearTerm::ColumnUnsigned { + coefficient, + column, + } => { + let coeff = FieldElement::::from(*coefficient); + result += get_col(*column) * coeff; + } + LinearTerm::Constant(value) => { + result += FieldElement::::from(*value); + } + } + } + if result != FieldElement::::zero() { + *acc += &result * &alpha_powers[alpha_offset]; + } + 1 + } + } + } + /// Computes the bus element value from column values. /// /// # Arguments diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 333956895..bc9ee36e1 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2659,23 +2659,33 @@ pub trait IsStarkProver< // Transient: consumed by the batch prove below, freed before the // aux/LDE phases — this transience is where GKR's memory win over // committed term columns lives. + #[cfg(feature = "instruments")] + let __sp_gkr = crate::instruments::span("r1_gkr_trees"); let layers_per_instance: Vec>> = crate::par::par_map_collect(0..gkr_indices.len(), |k| { let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; compute_logup_layers( air.bus_interactions(), - &trace.columns_main(), + &trace.main_table, trace.num_rows(), &lookup_challenges, ) }); + #[cfg(feature = "instruments")] + drop(__sp_gkr); + #[cfg(feature = "instruments")] + let __sp_gkr = crate::instruments::span("r1_gkr_batch_prove"); let (batch_gkr_proof, shared_point, final_claims) = gkr_prove_batch(layers_per_instance, transcript); + #[cfg(feature = "instruments")] + drop(__sp_gkr); // Per-instance finalize (parallel): column-claim MLE evaluations at // the instance random point (the per-column ⟨l, col⟩ inner products // are the heavy part). + #[cfg(feature = "instruments")] + let __sp_gkr = crate::instruments::span("r1_gkr_finalize"); let results: Vec> = crate::par::par_map_collect(0..gkr_indices.len(), |k| { let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; @@ -2689,13 +2699,15 @@ pub trait IsStarkProver< .expect("honest GKR root denominator is nonzero"); finalize_logup_gkr_result( air.bus_interactions(), - &trace.columns_main(), + &trace.main_table, random_point, n_claim, d_claim, table_contribution, ) }); + #[cfg(feature = "instruments")] + drop(__sp_gkr); // Phase B″: bind every table's column claims into the shared // transcript, THEN sample γ — the claims must precede the challenge diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 7d13b2ef1..de78950f3 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -1334,6 +1334,13 @@ pub fn prove_gkr_with_options_and_inputs( proof_options: &ProofOptions, max_rows: &MaxRowsConfig, ) -> Result { + // Wall-clock span tree (same output as the standard path's timeline, so + // GKR-vs-standard phase breakdowns compare directly under `instruments`). + #[cfg(feature = "instruments")] + stark::instruments::reset_timeline(); + #[cfg(feature = "instruments")] + let __root = stark::instruments::span("prove_total"); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; let executor = Executor::new(&program, private_inputs.to_vec()) .map_err(|e| Error::Execution(format!("{e}")))?; @@ -1399,6 +1406,17 @@ pub fn prove_gkr_with_options_and_inputs( ) .map_err(|e| Error::Prover(format!("{e:?}")))?; + #[cfg(feature = "instruments")] + { + drop(__root); + let spans = stark::instruments::take_timeline(); + print!("{}", stark::instruments::format_timeline(&spans)); + if let Ok(path) = std::env::var("LAMBDA_VM_TIMELINE_JSON") { + let _ = std::fs::write(&path, stark::instruments::timeline_json(&spans)); + println!("[timeline] wrote {path}"); + } + } + Ok(GkrVmProof { proof, runtime_page_ranges, From d4e64fa987423fffa444132a0b9a8ded3978ef7d Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 15:37:00 -0300 Subject: [PATCH 10/20] feat(recursion): GKR proofs through the recursion guest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recursion pipeline could not consume GkrVmProof at all (host-side, owned-deserialize only). Now: - Verifier::multi_verify_gkr_views: the GKR verify entry over proof VIEWS — per-table STARK proofs owned or rkyv-archived (read in place, #769-style), batch GKR proof + column claims borrowed. - verify_gkr_proof_parts: single GKR VM-proof verification impl (the analogue of verify_proof_parts), taking supplied DECODE/page roots; verify_gkr_with_options funnels into it. - GkrGuestInput + recursion::encode_gkr_guest_input: the guest blob on the same magic-prefixed wire format; verify_gkr_recursion_blob reads the archive in place and materializes only small metadata plus the KB-scale GKR artifacts (batch proof, column claims) — the MB-scale per-table data stays zero-copy. - recursion::verify_and_attest_gkr_blob: attests the SAME program_id(elf, roots) || public_output as the standard path (LogUp mode changes bus-sum proving, not program identity), so check_attestation needs no GKR variant. - guest: 'gkr' feature (mutually exclusive with 'continuation') + recursion-gkr-{min,blowup2}-bench bins; Makefile builds recursion-gkr-{min,blowup2}.elf. - dump mode: RECURSION_DUMP_GKR=1 on test_dump_recursion_input proves the inner in GKR mode and encodes the GkrGuestInput blob. - host roundtrip test mirrors the guest path end-to-end and pins the cross-format rejects (standard blob vs GKR verifier and vice versa). --- Makefile | 14 +- bench_vs/lambda/recursion/Cargo.toml | 14 ++ bench_vs/lambda/recursion/src/main.rs | 9 +- crypto/stark/src/verifier.rs | 37 ++++- prover/src/lib.rs | 182 +++++++++++++++++++---- prover/src/recursion.rs | 52 +++++++ prover/src/tests/recursion_smoke_test.rs | 75 ++++++++++ 7 files changed, 351 insertions(+), 32 deletions(-) diff --git a/Makefile b/Makefile index 4abed4a90..4003d4aef 100644 --- a/Makefile +++ b/Makefile @@ -70,8 +70,12 @@ RECURSION_VERIFIER_PRESETS := min blowup2 blowup4 blowup8 # `continuation` feature: verify a multi-epoch ContinuationProof bundle instead # of a monolithic VmProof. Only the presets the benchmarks actually measure. RECURSION_CONT_PRESETS := min blowup2 blowup4 +# `gkr` feature: verify a LogUpMode::Gkr inner proof (GkrGuestInput blob). +# Experimental — only the presets the GKR guest-cycle measurements use. +RECURSION_GKR_PRESETS := min blowup2 RECURSION_VERIFIER_ARTIFACTS := $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-, $(addsuffix .elf, $(RECURSION_VERIFIER_PRESETS))) \ - $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-cont-, $(addsuffix .elf, $(RECURSION_CONT_PRESETS))) + $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-cont-, $(addsuffix .elf, $(RECURSION_CONT_PRESETS))) \ + $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-gkr-, $(addsuffix .elf, $(RECURSION_GKR_PRESETS))) # Override with: make ... SYSROOT_DIR=$HOME/.lambda-vm-sysroot # to install the sysroot in a user-writable location and avoid sudo. @@ -245,6 +249,14 @@ $(RECURSION_ARTIFACTS_DIR)/recursion-cont-$(1).elf: FORCE | prepare-sysroot $(RE endef $(foreach preset,$(RECURSION_CONT_PRESETS),$(eval $(call recursion_cont_verifier_rule,$(preset)))) +# GKR variants: same crate, `gkr` feature on top of the preset feature -> +# recursion-gkr--bench -> recursion-gkr-.elf. +define recursion_gkr_verifier_rule +$(RECURSION_ARTIFACTS_DIR)/recursion-gkr-$(1).elf: FORCE | prepare-sysroot $(RECURSION_ARTIFACTS_DIR) + $$(call build_guest_elf,$$(RECURSION_GUESTS_DIR)/recursion,recursion-gkr-$(1)-bench,--features "gkr $(1)") +endef +$(foreach preset,$(RECURSION_GKR_PRESETS),$(eval $(call recursion_gkr_verifier_rule,$(preset)))) + clean-asm: -rm -rf $(ASM_ARTIFACTS_DIR) diff --git a/bench_vs/lambda/recursion/Cargo.toml b/bench_vs/lambda/recursion/Cargo.toml index f612949a8..a44e1163f 100644 --- a/bench_vs/lambda/recursion/Cargo.toml +++ b/bench_vs/lambda/recursion/Cargo.toml @@ -21,6 +21,10 @@ blowup8 = [] # memory-bounded inner prove) instead of a monolithic VmProof. Selects the # `recursion-cont--bench` bins below. continuation = [] +# Orthogonal to the presets (and mutually exclusive with `continuation`): +# verify a LogUpMode::Gkr inner proof (GkrGuestInput blob). Selects the +# `recursion-gkr--bench` bins below. +gkr = [] # One distinctly named binary per preset (selected by its feature) so a parallel # `make -j` builds them to different filenames — structurally race-free, no cp @@ -60,6 +64,16 @@ name = "recursion-cont-blowup4-bench" path = "src/main.rs" required-features = ["continuation", "blowup4"] +[[bin]] +name = "recursion-gkr-min-bench" +path = "src/main.rs" +required-features = ["gkr", "min"] + +[[bin]] +name = "recursion-gkr-blowup2-bench" +path = "src/main.rs" +required-features = ["gkr", "blowup2"] + [dependencies] lambda-vm-prover = { path = "../../../prover", default-features = false, features = [ "profile-markers", diff --git a/bench_vs/lambda/recursion/src/main.rs b/bench_vs/lambda/recursion/src/main.rs index 1a846109b..7c4052de4 100644 --- a/bench_vs/lambda/recursion/src/main.rs +++ b/bench_vs/lambda/recursion/src/main.rs @@ -53,6 +53,8 @@ compile_error!("select exactly one of the `min`/`blowup2`/`blowup4`/`blowup8` fe all(feature = "blowup4", feature = "blowup8"), ))] compile_error!("select exactly one of the `min`/`blowup2`/`blowup4`/`blowup8` features"); +#[cfg(all(feature = "continuation", feature = "gkr"))] +compile_error!("`continuation` and `gkr` are mutually exclusive input layouts"); /// The build preset fixing the inner `ProofOptions` (see the module docs). #[cfg(feature = "min")] @@ -88,7 +90,7 @@ pub fn main() -> ! { // not self-enforcing here. let options = PRESET.options(); - #[cfg(not(feature = "continuation"))] + #[cfg(not(any(feature = "continuation", feature = "gkr")))] let attestation = lambda_vm_prover::recursion::verify_and_attest_blob(blob, &options) .expect("verify errored") .expect("inner proof failed verification"); @@ -98,6 +100,11 @@ pub fn main() -> ! { .expect("verify errored") .expect("inner continuation proof failed verification"); + #[cfg(feature = "gkr")] + let attestation = lambda_vm_prover::recursion::verify_and_attest_gkr_blob(blob, &options) + .expect("verify errored") + .expect("inner GKR proof failed verification"); + lambda_vm_syscalls::syscalls::commit(&attestation); lambda_vm_syscalls::syscalls::sys_halt(); } diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index a35407297..1b438ceb6 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -1167,6 +1167,36 @@ pub trait IsStarkVerifier< FieldElement: AsBytes + Sync + Send, FieldElement: AsBytes + Sync + Send, PI: Clone, + { + Self::multi_verify_gkr_views( + airs, + MultiProofView::Owned(&gkr_proof.multi), + &gkr_proof.batch_gkr_proof, + gkr_proof.column_claims_by_table.as_slice(), + transcript, + expected_bus_balance, + ) + } + + /// [`Self::multi_verify_gkr`] over proof VIEWS: the per-table STARK proofs + /// come in as a [`ProofViewSource`] (owned or rkyv-archived, read in + /// place), while the (small) batch GKR proof and column claims come in as + /// borrowed values — the recursion guest deserializes only those and keeps + /// the heavy per-table data zero-copy. + fn multi_verify_gkr_views<'p>( + airs: &[&dyn AIR], + proofs: impl ProofViewSource<'p, Field, FieldExtension, PI>, + batch_gkr_proof: &BatchGkrProof, + column_claims_by_table: &[Option>], + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + ) -> bool + where + Field: 'p, + FieldExtension: 'p, + PI: 'p, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, { if airs .iter() @@ -1177,13 +1207,10 @@ pub trait IsStarkVerifier< } Self::multi_verify_views_impl( airs, - MultiProofView::Owned(&gkr_proof.multi), + proofs, transcript, expected_bus_balance, - Some(( - &gkr_proof.batch_gkr_proof, - gkr_proof.column_claims_by_table.as_slice(), - )), + Some((batch_gkr_proof, column_claims_by_table)), ) } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index de78950f3..b0da75cdc 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -213,6 +213,18 @@ pub struct GuestInput { pub page_commitments: Vec<(u64, Commitment)>, } +/// [`GuestInput`] for a [`LogUpMode::Gkr`] inner proof: the monolithic +/// [`VmProof`] replaced by the [`GkrVmProof`] bundle. Same wire format +/// (magic-prefixed rkyv archive); the guest is feature-pinned to one layout +/// and a blob of the other kind fails the bytecheck validation. +#[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct GkrGuestInput { + pub gkr_proof: GkrVmProof, + pub inner_elf: Vec, + pub decode_commitment: Commitment, + pub page_commitments: Vec<(u64, Commitment)>, +} + // ============================================================================ // Recursion-input wire format: aligning magic prefix + rkyv archive // ============================================================================ @@ -1435,46 +1447,81 @@ pub fn verify_gkr_with_options( ) -> Result { let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; let elf_digest = statement::elf_digest(elf_bytes); + verify_gkr_proof_parts( + MultiProofView::Owned(&gkr_proof.proof.multi), + &gkr_proof.proof.batch_gkr_proof, + &gkr_proof.proof.column_claims_by_table, + &gkr_proof.table_counts, + &gkr_proof.runtime_page_ranges, + gkr_proof.num_private_input_pages, + &gkr_proof.public_output, + &program, + &elf_digest, + proof_options, + None, + None, + ) +} - gkr_proof.table_counts.validate()?; +/// The single GKR VM-proof verification implementation, given the proof's +/// metadata fields plus an already-parsed ELF and its digest — the GKR +/// analogue of [`verify_proof_parts`]. Both [`verify_gkr_with_options`] +/// (owned) and [`verify_gkr_recursion_blob`] (guest blob) funnel here: the +/// per-table STARK proofs come in as a [`MultiProofView`] (owned or archived, +/// read in place), while the small GKR artifacts (batch proof + column +/// claims) come in as borrowed values the guest deserializes cheaply. +#[allow(clippy::too_many_arguments)] +fn verify_gkr_proof_parts( + proofs: MultiProofView<'_, F, E, ()>, + batch_gkr_proof: &stark::gkr::BatchGkrProof, + column_claims_by_table: &[Option>], + table_counts: &TableCounts, + runtime_page_ranges: &[RuntimePageRange], + num_private_input_pages: usize, + public_output: &[u8], + program: &Elf, + elf_digest: &[u8; 32], + proof_options: &ProofOptions, + decode_commitment: Option, + page_commitments: Option<&[(u64, Commitment)]>, +) -> Result { + table_counts.validate()?; { let max_pages = crate::tables::page::max_private_input_pages(); - if gkr_proof.num_private_input_pages > max_pages { + if num_private_input_pages > max_pages { return Err(Error::InvalidTableCounts(format!( - "num_private_input_pages ({}) exceeds max ({max_pages})", - gkr_proof.num_private_input_pages, + "num_private_input_pages ({num_private_input_pages}) exceeds max ({max_pages})", ))); } } let page_configs = Traces::page_configs_from_elf_and_runtime( - &program, - &gkr_proof.runtime_page_ranges, - gkr_proof.num_private_input_pages, + program, + runtime_page_ranges, + num_private_input_pages, ); - let expected_proof_count = - gkr_proof.table_counts.total() + FIXED_TABLE_COUNT + page_configs.len(); - if expected_proof_count != gkr_proof.proof.multi.proofs.len() { + let expected_proof_count = table_counts.total() + FIXED_TABLE_COUNT + page_configs.len(); + if expected_proof_count != proofs.len() { return Err(Error::InvalidTableCounts(format!( "table_counts total ({}) + {FIXED_TABLE_COUNT} fixed + {} pages = {}, but proof contains {} sub-proofs", - gkr_proof.table_counts.total(), + table_counts.total(), page_configs.len(), expected_proof_count, - gkr_proof.proof.multi.proofs.len(), + proofs.len(), ))); } let airs = VmAirs::new_with_logup_mode( - &program, + program, proof_options, false, &page_configs, - &gkr_proof.table_counts, - None, + table_counts, + decode_commitment, true, None, - None, + page_commitments, None, LogUpMode::Gkr, ); @@ -1484,11 +1531,11 @@ pub fn verify_gkr_with_options( absorb_statement_with_digest( &mut transcript, StatementKind::Monolithic, - &elf_digest, - &gkr_proof.public_output, - &gkr_proof.table_counts, - gkr_proof.num_private_input_pages, - &gkr_proof.runtime_page_ranges, + elf_digest, + public_output, + table_counts, + num_private_input_pages, + runtime_page_ranges, proof_options.fri_final_poly_log_degree, ); @@ -1497,8 +1544,8 @@ pub fn verify_gkr_with_options( let mut transcript_for_replay = transcript.clone(); let expected_bus_balance = match compute_expected_commit_bus_balance_view( &air_refs, - MultiProofView::Owned(&gkr_proof.proof.multi), - &gkr_proof.public_output, + proofs, + public_output, 0, &mut transcript_for_replay, ) { @@ -1506,14 +1553,99 @@ pub fn verify_gkr_with_options( None => return Ok(false), }; - Ok(Verifier::multi_verify_gkr( + Ok(Verifier::multi_verify_gkr_views( &air_refs, - &gkr_proof.proof, + proofs, + batch_gkr_proof, + column_claims_by_table, &mut transcript, &expected_bus_balance, )) } +/// Verify a GKR recursion-input blob produced by +/// [`recursion::encode_gkr_guest_input`] — the GKR analogue of +/// [`verify_recursion_blob`]. The per-table STARK proofs are verified +/// **in place** from the archive (zero-copy, like the standard path); only +/// the small metadata and the GKR artifacts (batch proof + column claims, +/// KBs against the MB-scale proof data) are materialized. +pub fn verify_gkr_recursion_blob<'a>( + blob: &'a [u8], + proof_options: &ProofOptions, +) -> Result, Error> { + use rkyv::rancor::Error as RkyvError; + + let mut aligned_fallback = rkyv::util::AlignedVec::::new(); + let (archived, archive_bytes, archive_base): (&ArchivedGkrGuestInput, &[u8], *const u8) = + access_recursion_archive(blob, &mut aligned_fallback)?; + + // Materialize only the small metadata; the per-table proofs stay in the + // buffer. + let table_counts: TableCounts = + rkyv::deserialize::(&archived.gkr_proof.table_counts) + .map_err(|e| Error::Execution(format!("rkyv deserialize table_counts failed: {e}")))?; + let runtime_page_ranges: Vec = rkyv::deserialize::< + Vec, + RkyvError, + >(&archived.gkr_proof.runtime_page_ranges) + .map_err(|e| Error::Execution(format!("rkyv deserialize page ranges failed: {e}")))?; + let page_commitments: Vec<(u64, Commitment)> = rkyv::deserialize::< + Vec<(u64, Commitment)>, + RkyvError, + >(&archived.page_commitments) + .map_err(|e| Error::Execution(format!("rkyv deserialize page commitments failed: {e}")))?; + // The GKR artifacts are small (root claims + layer round polys + column + // claims — KBs); an owned deserialize here keeps the batch-GKR verifier + // and claim reconstruction unchanged while the heavy data stays archived. + let batch_gkr_proof: stark::gkr::BatchGkrProof = + rkyv::deserialize::<_, RkyvError>(&archived.gkr_proof.proof.batch_gkr_proof) + .map_err(|e| Error::Execution(format!("rkyv deserialize batch proof failed: {e}")))?; + let column_claims_by_table: Vec>> = + rkyv::deserialize::<_, RkyvError>(&archived.gkr_proof.proof.column_claims_by_table) + .map_err(|e| Error::Execution(format!("rkyv deserialize column claims failed: {e}")))?; + let num_private_input_pages = archived.gkr_proof.num_private_input_pages.to_native() as usize; + let inner_elf: &[u8] = archived.inner_elf.as_slice(); + let public_output: &[u8] = archived.gkr_proof.public_output.as_slice(); + let decode_commitment: Commitment = archived.decode_commitment; + + // Rebase the returned slices onto the caller's buffer (`archived` may + // point into the aligned fallback copy) — same as the standard path. + let rebase = |s: &[u8]| -> &'a [u8] { + let offset = s.as_ptr() as usize - archive_base as usize; + &archive_bytes[offset..offset + s.len()] + }; + let inner_elf_rebased = rebase(inner_elf); + let public_output_rebased = rebase(public_output); + + let program = Elf::load(inner_elf).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let elf_digest = statement::elf_digest(inner_elf); + + let ok = verify_gkr_proof_parts( + MultiProofView::Archived(&archived.gkr_proof.proof.multi), + &batch_gkr_proof, + &column_claims_by_table, + &table_counts, + &runtime_page_ranges, + num_private_input_pages, + public_output, + &program, + &elf_digest, + proof_options, + Some(decode_commitment), + Some(&page_commitments), + )?; + + Ok(RecursionVerification { + ok, + public_output: public_output_rebased, + inner_elf: inner_elf_rebased, + decode_commitment, + page_commitments, + elf_digest, + entry_point: program.entry_point, + }) +} + /// Verify a proof produced by [`prove`] using default proof options. /// /// Uses [`GoldilocksCubicProofOptions::with_blowup(2)`] for verification. diff --git a/prover/src/recursion.rs b/prover/src/recursion.rs index efca722c9..7cf5b7421 100644 --- a/prover/src/recursion.rs +++ b/prover/src/recursion.rs @@ -150,6 +150,33 @@ pub fn encode_guest_input( }) } +/// Build the GKR guest's private-input blob for `inner_proof` of `inner_elf`: +/// precomputes the roots and rkyv-encodes a [`crate::GkrGuestInput`] behind +/// the standard aligning prefix — [`encode_guest_input`] for a +/// [`crate::GkrVmProof`] inner. +pub fn encode_gkr_guest_input( + inner_proof: &crate::GkrVmProof, + inner_elf: &[u8], + opts: &ProofOptions, +) -> Result, Error> { + let (decode_commitment, page_commitments) = precomputed_commitments(inner_elf, opts)?; + let input = crate::GkrGuestInput { + gkr_proof: inner_proof.clone(), + inner_elf: inner_elf.to_vec(), + decode_commitment, + page_commitments, + }; + let archive = rkyv::to_bytes::(&input) + .map_err(|e| Error::Execution(format!("rkyv encode failed: {e}")))?; + let mut blob = Vec::with_capacity(crate::RECURSION_INPUT_PREFIX_LEN + archive.len()); + blob.extend_from_slice(&crate::RECURSION_INPUT_MAGIC); + blob.extend_from_slice(&crate::RECURSION_INPUT_VERSION.to_le_bytes()); + blob.extend_from_slice(&[0u8; 4]); // reserved + debug_assert_eq!(blob.len(), crate::RECURSION_INPUT_PREFIX_LEN); + blob.extend_from_slice(&archive); + Ok(blob) +} + /// The continuation guest's private-input layout (the `continuation` guest /// feature). Mirrors [`crate::GuestInput`] with the monolithic proof replaced /// by the bundle and the PAGE roots replaced by the global-memory genesis @@ -289,6 +316,31 @@ pub fn verify_and_attest_blob( Ok(Some(attestation)) } +/// [`verify_and_attest_blob`] for a [`LogUpMode::Gkr`](stark::lookup::LogUpMode) +/// inner proof: verifies a [`crate::GkrGuestInput`] blob +/// ([`encode_gkr_guest_input`]) — per-table proofs in place from the archive, +/// only the small GKR artifacts materialized — and attests the SAME +/// `program_id(elf, roots) || inner_public_output` as the standard path (the +/// LogUp mode changes how bus sums are proven, not the program identity). +pub fn verify_and_attest_gkr_blob( + blob: &[u8], + proof_options: &ProofOptions, +) -> Result>, Error> { + let verification = crate::verify_gkr_recursion_blob(blob, proof_options)?; + if !verification.ok { + return Ok(None); + } + let id = program_id_from_digest( + &verification.elf_digest, + verification.entry_point, + &verification.decode_commitment, + &verification.page_commitments, + ); + let mut attestation = id.to_vec(); + attestation.extend_from_slice(verification.public_output); + Ok(Some(attestation)) +} + /// [`verify_and_attest_blob`]'s logic for a continuation bundle: takes the /// wire-format blob ([`encode_continuation_guest_input`]) and does the /// intended `continuation` guest's whole job in one call — verify every diff --git a/prover/src/tests/recursion_smoke_test.rs b/prover/src/tests/recursion_smoke_test.rs index 1f4800011..a162a36b2 100644 --- a/prover/src/tests/recursion_smoke_test.rs +++ b/prover/src/tests/recursion_smoke_test.rs @@ -594,6 +594,59 @@ fn run_recursion_pipeline( ); } +/// The GKR analogue of `test_recursion_blob_decodes_and_verifies_on_host`: +/// prove the inner in `LogUpMode::Gkr`, encode a `GkrGuestInput` blob, mirror +/// the `gkr` guest's verify+attest on the host (per-table proofs read in +/// place from the archive), and run the consumer check. The attestation uses +/// the SAME `program_id` as the standard path, so `check_attestation` needs +/// no GKR variant. +#[test] +fn test_gkr_recursion_blob_decodes_and_verifies_on_host() { + let root = workspace_root(); + let empty_elf_bytes = read_guest_elf(&root, "empty"); + + let inner_proof = crate::prove_gkr_with_options_and_inputs( + &empty_elf_bytes, + &[], + &MIN_PROOF_OPTIONS, + &crate::MaxRowsConfig::default(), + ) + .expect("inner GKR prove should succeed"); + let blob = + recursion::encode_gkr_guest_input(&inner_proof, &empty_elf_bytes, &MIN_PROOF_OPTIONS) + .expect("recursion::encode_gkr_guest_input failed"); + assert!( + blob.len() <= executor::vm::memory::MAX_PRIVATE_INPUT_SIZE as usize, + "GKR recursion input exceeds MAX_PRIVATE_INPUT_SIZE" + ); + + let attestation = recursion::verify_and_attest_gkr_blob(&blob, &MIN_PROOF_OPTIONS) + .expect("verify_and_attest_gkr_blob errored") + .expect("GKR proof did not survive the rkyv round-trip"); + + let output = recursion::check_attestation(&attestation, &empty_elf_bytes, &MIN_PROOF_OPTIONS) + .expect("check_attestation errored") + .expect("attestation must match the trusted inner ELF"); + assert_eq!( + output, inner_proof.public_output, + "attested public output must equal the inner proof's public output" + ); + + // A standard-mode blob fed to the GKR verifier (and vice versa) must fail + // the bytecheck validation, not misparse: the two GuestInput layouts are + // distinct archived types. + let (_, std_blob) = + prove_inner_and_encode_blob("gkr-cross", &empty_elf_bytes, &[], &MIN_PROOF_OPTIONS); + assert!( + crate::verify_gkr_recursion_blob(&std_blob, &MIN_PROOF_OPTIONS).is_err(), + "standard blob must be rejected by the GKR blob verifier" + ); + assert!( + crate::verify_recursion_blob(&blob, &MIN_PROOF_OPTIONS).is_err(), + "GKR blob must be rejected by the standard blob verifier" + ); +} + /// Decode the blob on the host and mirror the guest's verify+attest, then run /// the consumer check — a cheap guard on the encode/decode/attest contract /// without running the VM. @@ -889,6 +942,10 @@ fn test_recursion_prove_1query() { /// * `RECURSION_DUMP_EPOCH_LOG2` (int, default unset = monolithic) — prove via /// continuations with `2^n`-cycle epochs and encode a /// [`recursion::ContinuationGuestInput`] blob for `recursion-cont-.elf`. +/// * `RECURSION_DUMP_GKR=1` (default unset) — prove the inner in +/// `LogUpMode::Gkr` and encode a [`crate::GkrGuestInput`] blob for +/// `recursion-gkr-.elf`. Mutually exclusive with +/// `RECURSION_DUMP_EPOCH_LOG2`. #[test] #[ignore = "diagnostic: writes recursion private input to /tmp/recursion_input.bin"] fn test_dump_recursion_input() { @@ -968,6 +1025,24 @@ fn test_dump_recursion_input() { .expect("recursion::encode_continuation_guest_input failed"); (blob, Some((expected_id, expected_output))) } + Err(_) if std::env::var("RECURSION_DUMP_GKR").is_ok_and(|v| v == "1") => { + let opts = preset.options(); + eprintln!( + "[dump-input] proving inner in LogUpMode::Gkr (blowup={}, fri_queries={}) ...", + opts.blowup_factor, opts.fri_number_of_queries + ); + let inner_proof = crate::prove_gkr_with_options_and_inputs( + &inner_elf_bytes, + &inner_input, + &opts, + &crate::MaxRowsConfig::default(), + ) + .expect("inner GKR prove should succeed"); + let blob = recursion::encode_gkr_guest_input(&inner_proof, &inner_elf_bytes, &opts) + .expect("recursion::encode_gkr_guest_input failed"); + eprintln!("[dump-input] GKR rkyv blob: {} bytes", blob.len()); + (blob, None) + } Err(_) => { let (_inner_proof, blob) = prove_inner_and_encode_blob( "dump-input", From f8a9236c4c36908629426789be9db3a3177cea8b Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 16:39:20 -0300 Subject: [PATCH 11/20] feat(continuation): LogUp-GKR mode for the continuation pipeline Every multi_prove/multi_verify site in the continuation path is now mode-parametrized: epoch proofs (VM tables + the epoch-local L2G bookend), the cross-epoch global-memory proof (per-epoch L2G airs + one GLOBAL_MEMORY air per touched page), and their verifiers all thread a LogUpMode through the AIR builders (with_logup_mode) and dispatch to the standard or GKR prove/verify entry points. - GkrContinuationProof = the standard ContinuationProof (unchanged wire format) + per-epoch and global GkrProofExtras (batch proof + column claims, KB-scale). Containing the standard bundle as `base` keeps every view, L2G binding check, and root precompute working unchanged. - prove_continuation_gkr / verify_continuation_gkr (host, owned) and verify_gkr_continuation_archived (guest: per-table proofs read in place off the archived base bundle, only the GKR extras deserialized). - recursion: GkrContinuationGuestInput + encode_gkr_continuation_guest_input + verify_gkr_continuation_and_attest -- same program_id attestation as the standard continuation path (consumer flow unchanged). - guest: the gkr and continuation features now compose (fourth input layout); recursion-gkr-cont-{min,blowup2}.elf Makefile targets. - CLI: --logup-mode now applies to --continuations on prove and verify. - dump mode: RECURSION_DUMP_GKR=1 + RECURSION_DUMP_EPOCH_LOG2 dumps a GKR continuation blob. - host roundtrip test: multi-epoch GKR continuation proves, verifies owned + archived, attests, and passes the consumer program_id check. --- Makefile | 16 +- bench_vs/lambda/recursion/Cargo.toml | 16 +- bench_vs/lambda/recursion/src/main.rs | 14 +- bin/cli/src/main.rs | 130 +++++--- prover/src/continuation.rs | 400 +++++++++++++++++++---- prover/src/recursion.rs | 95 ++++++ prover/src/tests/recursion_smoke_test.rs | 165 ++++++++-- 7 files changed, 700 insertions(+), 136 deletions(-) diff --git a/Makefile b/Makefile index 4003d4aef..ab8e1de88 100644 --- a/Makefile +++ b/Makefile @@ -70,12 +70,14 @@ RECURSION_VERIFIER_PRESETS := min blowup2 blowup4 blowup8 # `continuation` feature: verify a multi-epoch ContinuationProof bundle instead # of a monolithic VmProof. Only the presets the benchmarks actually measure. RECURSION_CONT_PRESETS := min blowup2 blowup4 -# `gkr` feature: verify a LogUpMode::Gkr inner proof (GkrGuestInput blob). -# Experimental — only the presets the GKR guest-cycle measurements use. +# `gkr` feature: verify a LogUpMode::Gkr inner proof (GkrGuestInput blob; +# with `continuation`, a GkrContinuationProof bundle). Experimental — only the +# presets the GKR guest-cycle measurements use. RECURSION_GKR_PRESETS := min blowup2 RECURSION_VERIFIER_ARTIFACTS := $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-, $(addsuffix .elf, $(RECURSION_VERIFIER_PRESETS))) \ $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-cont-, $(addsuffix .elf, $(RECURSION_CONT_PRESETS))) \ - $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-gkr-, $(addsuffix .elf, $(RECURSION_GKR_PRESETS))) + $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-gkr-, $(addsuffix .elf, $(RECURSION_GKR_PRESETS))) \ + $(addprefix $(RECURSION_ARTIFACTS_DIR)/recursion-gkr-cont-, $(addsuffix .elf, $(RECURSION_GKR_PRESETS))) # Override with: make ... SYSROOT_DIR=$HOME/.lambda-vm-sysroot # to install the sysroot in a user-writable location and avoid sudo. @@ -257,6 +259,14 @@ $(RECURSION_ARTIFACTS_DIR)/recursion-gkr-$(1).elf: FORCE | prepare-sysroot $(REC endef $(foreach preset,$(RECURSION_GKR_PRESETS),$(eval $(call recursion_gkr_verifier_rule,$(preset)))) +# GKR continuation variants: `gkr` + `continuation` on top of the preset -> +# recursion-gkr-cont--bench -> recursion-gkr-cont-.elf. +define recursion_gkr_cont_verifier_rule +$(RECURSION_ARTIFACTS_DIR)/recursion-gkr-cont-$(1).elf: FORCE | prepare-sysroot $(RECURSION_ARTIFACTS_DIR) + $$(call build_guest_elf,$$(RECURSION_GUESTS_DIR)/recursion,recursion-gkr-cont-$(1)-bench,--features "gkr continuation $(1)") +endef +$(foreach preset,$(RECURSION_GKR_PRESETS),$(eval $(call recursion_gkr_cont_verifier_rule,$(preset)))) + clean-asm: -rm -rf $(ASM_ARTIFACTS_DIR) diff --git a/bench_vs/lambda/recursion/Cargo.toml b/bench_vs/lambda/recursion/Cargo.toml index a44e1163f..fbb249a4b 100644 --- a/bench_vs/lambda/recursion/Cargo.toml +++ b/bench_vs/lambda/recursion/Cargo.toml @@ -21,9 +21,9 @@ blowup8 = [] # memory-bounded inner prove) instead of a monolithic VmProof. Selects the # `recursion-cont--bench` bins below. continuation = [] -# Orthogonal to the presets (and mutually exclusive with `continuation`): -# verify a LogUpMode::Gkr inner proof (GkrGuestInput blob). Selects the -# `recursion-gkr--bench` bins below. +# Orthogonal to the presets: verify a LogUpMode::Gkr inner proof +# (GkrGuestInput blob). Composes with `continuation` to select the +# GkrContinuationGuestInput layout (`recursion-gkr-cont--bench`). gkr = [] # One distinctly named binary per preset (selected by its feature) so a parallel @@ -74,6 +74,16 @@ name = "recursion-gkr-blowup2-bench" path = "src/main.rs" required-features = ["gkr", "blowup2"] +[[bin]] +name = "recursion-gkr-cont-min-bench" +path = "src/main.rs" +required-features = ["gkr", "continuation", "min"] + +[[bin]] +name = "recursion-gkr-cont-blowup2-bench" +path = "src/main.rs" +required-features = ["gkr", "continuation", "blowup2"] + [dependencies] lambda-vm-prover = { path = "../../../prover", default-features = false, features = [ "profile-markers", diff --git a/bench_vs/lambda/recursion/src/main.rs b/bench_vs/lambda/recursion/src/main.rs index 7c4052de4..4bcb01898 100644 --- a/bench_vs/lambda/recursion/src/main.rs +++ b/bench_vs/lambda/recursion/src/main.rs @@ -53,8 +53,8 @@ compile_error!("select exactly one of the `min`/`blowup2`/`blowup4`/`blowup8` fe all(feature = "blowup4", feature = "blowup8"), ))] compile_error!("select exactly one of the `min`/`blowup2`/`blowup4`/`blowup8` features"); -#[cfg(all(feature = "continuation", feature = "gkr"))] -compile_error!("`continuation` and `gkr` are mutually exclusive input layouts"); +// `continuation` and `gkr` compose: together they select the +// `GkrContinuationGuestInput` layout (a GKR-mode continuation bundle). /// The build preset fixing the inner `ProofOptions` (see the module docs). #[cfg(feature = "min")] @@ -95,16 +95,22 @@ pub fn main() -> ! { .expect("verify errored") .expect("inner proof failed verification"); - #[cfg(feature = "continuation")] + #[cfg(all(feature = "continuation", not(feature = "gkr")))] let attestation = lambda_vm_prover::recursion::verify_continuation_and_attest(blob, &options) .expect("verify errored") .expect("inner continuation proof failed verification"); - #[cfg(feature = "gkr")] + #[cfg(all(feature = "gkr", not(feature = "continuation")))] let attestation = lambda_vm_prover::recursion::verify_and_attest_gkr_blob(blob, &options) .expect("verify errored") .expect("inner GKR proof failed verification"); + #[cfg(all(feature = "gkr", feature = "continuation"))] + let attestation = + lambda_vm_prover::recursion::verify_gkr_continuation_and_attest(blob, &options) + .expect("verify errored") + .expect("inner GKR continuation proof failed verification"); + lambda_vm_syscalls::syscalls::commit(&attestation); lambda_vm_syscalls::syscalls::sys_halt(); } diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index 51a06da33..c2fd9392a 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -193,8 +193,8 @@ enum Commands { /// LogUp mode. Defaults to `gkr` on this experimental branch (see /// `CliLogUpMode`); pass `--logup-mode standard` for the production - /// path. Not supported together with --continuations. - #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr, conflicts_with = "continuations")] + /// path. Applies to --continuations too. + #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr)] logup_mode: CliLogUpMode, /// Prove with continuations (split execution into epochs; flat peak memory) @@ -232,7 +232,7 @@ enum Commands { /// LogUp mode the proof was generated with (must match `prove`'s). /// Defaults to `gkr`, matching this branch's `prove` default. - #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr, conflicts_with = "continuations")] + #[arg(long, value_enum, default_value_t = CliLogUpMode::Gkr)] logup_mode: CliLogUpMode, /// Verify a continuation proof bundle (produced by `prove --continuations`) @@ -297,6 +297,7 @@ fn main() -> ExitCode { blowup, time, cycles, + logup_mode, ) } else { cmd_prove( @@ -320,7 +321,7 @@ fn main() -> ExitCode { continuations, } => { if continuations { - cmd_verify_continuation(proof, elf, blowup, time) + cmd_verify_continuation(proof, elf, blowup, time, logup_mode) } else { cmd_verify(proof, elf, blowup, time, logup_mode) } @@ -823,6 +824,7 @@ fn cmd_verify( } } +#[allow(clippy::too_many_arguments)] fn cmd_prove_continuation( elf_path: PathBuf, output_path: PathBuf, @@ -831,6 +833,7 @@ fn cmd_prove_continuation( blowup: u8, time: bool, cycles: bool, + logup_mode: CliLogUpMode, ) -> ExitCode { eprintln!("Reading ELF file..."); let elf_data = match std::fs::read(&elf_path) { @@ -879,22 +882,56 @@ fn cmd_prove_continuation( }; eprintln!( - "Generating continuation proof (blowup={blowup}, epoch_size_log2={epoch_size_log2}, epoch_size={epoch_size})...", + "Generating continuation proof (blowup={blowup}, epoch_size_log2={epoch_size_log2}, epoch_size={epoch_size}, logup={logup_mode:?})...", ); + // Prove + serialize per mode; timing covers proof generation only. let start = Instant::now(); - let bundle = match prover::continuation::prove_continuation( - &elf_data, - &private_inputs, - epoch_size_log2, - &opts, - ) { - Ok(b) => b, - Err(e) => { - eprintln!("Continuation proof generation failed: {}", e); - return ExitCode::FAILURE; + let (prove_elapsed, num_epochs, bytes) = match logup_mode { + CliLogUpMode::Standard => { + let bundle = match prover::continuation::prove_continuation( + &elf_data, + &private_inputs, + epoch_size_log2, + &opts, + ) { + Ok(b) => b, + Err(e) => { + eprintln!("Continuation proof generation failed: {}", e); + return ExitCode::FAILURE; + } + }; + let prove_elapsed = start.elapsed(); + match rkyv::to_bytes::(&bundle) { + Ok(b) => (prove_elapsed, bundle.num_epochs(), b), + Err(e) => { + eprintln!("Failed to serialize proof: {}", e); + return ExitCode::FAILURE; + } + } + } + CliLogUpMode::Gkr => { + let bundle = match prover::continuation::prove_continuation_gkr( + &elf_data, + &private_inputs, + epoch_size_log2, + &opts, + ) { + Ok(b) => b, + Err(e) => { + eprintln!("Continuation proof generation failed: {}", e); + return ExitCode::FAILURE; + } + }; + let prove_elapsed = start.elapsed(); + match rkyv::to_bytes::(&bundle) { + Ok(b) => (prove_elapsed, bundle.num_epochs(), b), + Err(e) => { + eprintln!("Failed to serialize proof: {}", e); + return ExitCode::FAILURE; + } + } } }; - let prove_elapsed = start.elapsed(); eprintln!("Writing proof..."); let file = match File::create(&output_path) { @@ -905,13 +942,6 @@ fn cmd_prove_continuation( } }; let mut writer = BufWriter::new(file); - let bytes = match rkyv::to_bytes::(&bundle) { - Ok(b) => b, - Err(e) => { - eprintln!("Failed to serialize proof: {}", e); - return ExitCode::FAILURE; - } - }; if let Err(e) = writer.write_all(&bytes) { eprintln!("Failed to write proof: {}", e); return ExitCode::FAILURE; @@ -921,7 +951,7 @@ fn cmd_prove_continuation( if let Some(c) = cycle_count { println!("Cycles: {}", c); } - println!("Epochs: {}", bundle.num_epochs()); + println!("Epochs: {}", num_epochs); if time { println!("Proving time: {:.3}s", prove_elapsed.as_secs_f64()); } @@ -933,6 +963,7 @@ fn cmd_verify_continuation( elf_path: PathBuf, blowup: u8, time: bool, + logup_mode: CliLogUpMode, ) -> ExitCode { eprintln!("Reading ELF file..."); let elf_data = match std::fs::read(&elf_path) { @@ -951,17 +982,6 @@ fn cmd_verify_continuation( return ExitCode::FAILURE; } }; - let bundle: prover::continuation::ContinuationProof = - match rkyv::from_bytes::( - &proof_bytes, - ) { - Ok(p) => p, - Err(e) => { - eprintln!("Failed to deserialize proof: {}", e); - return ExitCode::FAILURE; - } - }; - let opts = match GoldilocksCubicProofOptions::with_blowup(blowup) { Ok(opts) => opts, Err(e) => { @@ -970,10 +990,44 @@ fn cmd_verify_continuation( } }; - eprintln!("Verifying continuation proof..."); + eprintln!("Verifying continuation proof (logup={logup_mode:?})..."); let start = Instant::now(); - let result = prover::continuation::verify_continuation(&elf_data, &bundle, &opts); - let verify_elapsed = start.elapsed(); + let (result, verify_elapsed) = match logup_mode { + CliLogUpMode::Standard => { + let bundle: prover::continuation::ContinuationProof = match rkyv::from_bytes::< + prover::continuation::ContinuationProof, + rkyv::rancor::Error, + >(&proof_bytes) + { + Ok(p) => p, + Err(e) => { + eprintln!("Failed to deserialize proof: {}", e); + return ExitCode::FAILURE; + } + }; + let result = prover::continuation::verify_continuation(&elf_data, &bundle, &opts); + (result, start.elapsed()) + } + CliLogUpMode::Gkr => { + let bundle: prover::continuation::GkrContinuationProof = match rkyv::from_bytes::< + prover::continuation::GkrContinuationProof, + rkyv::rancor::Error, + >( + &proof_bytes + ) { + Ok(p) => p, + Err(e) => { + eprintln!( + "Failed to deserialize proof (was it proven with --logup-mode gkr?): {}", + e + ); + return ExitCode::FAILURE; + } + }; + let result = prover::continuation::verify_continuation_gkr(&elf_data, &bundle, &opts); + (result, start.elapsed()) + } + }; match result { Ok(Some(output)) => { diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 169cd7278..dfbe64e42 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -55,9 +55,12 @@ use executor::vm::execution::Executor; use math::field::element::FieldElement; use stark::config::Commitment; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet, EmptyConstraints}; -use stark::lookup::{AirWithBuses, AuxiliaryTraceBuildData, NullBoundaryConstraintBuilder}; +use stark::gkr::BatchGkrProof; +use stark::lookup::{ + AirWithBuses, AuxiliaryTraceBuildData, LogUpMode, NullBoundaryConstraintBuilder, +}; use stark::proof::options::ProofOptions; -use stark::proof::stark::MultiProof; +use stark::proof::stark::{GkrColumnClaims, MultiProof}; use stark::proof::view::MultiProofView; use stark::prover::{IsStarkProver, Prover}; use stark::trace::TraceTable; @@ -79,6 +82,9 @@ use crate::{ type F = GoldilocksField; type E = GoldilocksExtension; type AirRef<'a> = &'a dyn AIR; +/// [`prove_global`]'s result: the proof plus its GKR artifacts (Some exactly +/// under [`LogUpMode::Gkr`]). +type GlobalProveResult = Result<(MultiProof, Option), Error>; /// Fresh transcript seeded with the epoch's statement (ELF, public output, table /// layout) and `epoch_label` (its position). The epoch's prove, verify, and @@ -162,6 +168,7 @@ impl ConstraintSet for L2gMemoryConstraints { fn l2g_global_air( opts: &ProofOptions, epoch_label: u64, + mode: LogUpMode, ) -> AirWithBuses { AirWithBuses::new( local_to_global::cols::NUM_COLUMNS, @@ -172,6 +179,7 @@ fn l2g_global_air( 1, EmptyConstraints, ) + .with_logup_mode(mode) } /// Local-to-global AIR on the epoch-local Memory bus (used inside an epoch proof). @@ -183,6 +191,7 @@ fn l2g_global_air( fn l2g_memory_air( opts: &ProofOptions, epoch_label: u64, + mode: LogUpMode, ) -> AirWithBuses { let interactions = [ local_to_global::memory_bus_interactions(), @@ -196,6 +205,7 @@ fn l2g_memory_air( 1, L2gMemoryConstraints, ) + .with_logup_mode(mode) } /// GLOBAL_MEMORY AIR for one touched page (the cross-epoch analog of PAGE). @@ -221,6 +231,7 @@ fn global_memory_air( opts: &ProofOptions, config: &PageConfig, preprocessed: Option, + mode: LogUpMode, ) -> AirWithBuses { let air = AirWithBuses::new( global_memory::cols::NUM_COLUMNS, @@ -230,7 +241,8 @@ fn global_memory_air( opts, 1, EmptyConstraints, - ); + ) + .with_logup_mode(mode); if config.is_private_input { return air; } @@ -447,6 +459,46 @@ impl ContinuationProof { } } +/// The GKR artifacts of one multi-table proof under +/// [`LogUpMode::Gkr`]: the batch GKR proof plus the per-table column claims — +/// the two fields a [`stark::proof::stark::GkrMultiProof`] carries beyond its +/// inner [`MultiProof`]. KB-scale, so the archived verify path deserializes +/// these while the per-table proofs stay zero-copy. +#[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct GkrProofExtras { + pub batch_gkr_proof: BatchGkrProof, + pub column_claims_by_table: Vec>>, +} + +/// A continuation proof under [`LogUpMode::Gkr`]: the standard +/// [`ContinuationProof`] layout (per-epoch proofs are plain [`MultiProof`]s +/// with no bus public inputs) plus each proof's GKR artifacts. Containing the +/// standard bundle as `base` keeps every view, binding check, and root +/// precompute working unchanged; the standard bundle's wire format is +/// untouched (this is a separate top-level type, like +/// [`stark::proof::stark::GkrMultiProof`] vs [`MultiProof`]). +#[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct GkrContinuationProof { + base: ContinuationProof, + /// Per-epoch GKR artifacts, aligned with `base.epochs`. + epoch_gkr: Vec, + /// The global-memory proof's GKR artifacts. + global_gkr: GkrProofExtras, +} + +impl GkrContinuationProof { + /// Number of epochs the execution was split into. + pub fn num_epochs(&self) -> usize { + self.base.num_epochs() + } + + /// The standard-layout bundle inside (per-epoch + global proofs and the + /// public statement fields) — what root precomputes and consumers read. + pub fn base(&self) -> &ContinuationProof { + &self.base + } +} + /// Borrowed view over an [`EpochProof`] (owned or archived-in-place). Lets /// `verify_epoch` take a single argument again instead of the field-by-field /// parameter list the owned/archived split used to force on every caller: @@ -599,6 +651,7 @@ fn build_epoch_airs( reg_fini: &[u32], is_final: bool, decode_commitment: Option, + mode: LogUpMode, ) -> VmAirs { // Continuation epochs preprocess FINI = R_{i+1} too (not just INIT = R_i), so the // final register file is a verifier-known public value bound by the REG-C2 @@ -608,7 +661,7 @@ fn build_epoch_airs( register::compute_precomputed_commitment_with_fini(opts, register_init, reg_fini), register::NUM_PREPROCESSED_COLS_WITH_FINI, )); - VmAirs::new( + VmAirs::new_with_logup_mode( elf, opts, false, @@ -619,6 +672,7 @@ fn build_epoch_airs( None, None, register_preprocessed, + mode, ) } @@ -635,7 +689,8 @@ fn prove_epoch( is_final: bool, boundary: &[CellBoundary], opts: &ProofOptions, -) -> Result { + mode: LogUpMode, +) -> Result<(EpochProof, Option), Error> { // Count this L2G table's range-check lookups into the BITWISE table so its // AreBytes/IsHalfword multiplicities balance the range-check senders. crate::tables::bitwise::update_multiplicities( @@ -668,6 +723,7 @@ fn prove_epoch( ®_fini, is_final, None, + mode, ); let label = start.label; @@ -682,7 +738,7 @@ fn prove_epoch( ) }; - let l2g_air = l2g_memory_air(opts, label); + let l2g_air = l2g_memory_air(opts, label, mode); // Build this epoch's L2G table from the cross-epoch boundary so it is identical // to the one the global proof commits (the commitment binding compares their // roots). It is appended to the proof below, not through `air_trace_pairs`. @@ -690,13 +746,34 @@ fn prove_epoch( let mut pairs = airs.air_trace_pairs(&mut traces); pairs.push((&l2g_air, &mut l2g_trace, &())); - let proof = Prover::multi_prove( - pairs, - &mut seed(), - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - ) - .map_err(|e| Error::Prover(format!("{e:?}")))?; + let (proof, gkr_extras) = match mode { + LogUpMode::Standard => ( + Prover::multi_prove( + pairs, + &mut seed(), + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?, + None, + ), + LogUpMode::Gkr => { + let gkr = Prover::multi_prove_gkr( + pairs, + &mut seed(), + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + ( + gkr.multi, + Some(GkrProofExtras { + batch_gkr_proof: gkr.batch_gkr_proof, + column_claims_by_table: gkr.column_claims_by_table, + }), + ) + } + }; let l2g_root = proof .proofs @@ -706,14 +783,17 @@ fn prove_epoch( })? .lde_trace_main_merkle_root; - Ok(EpochProof { - proof, - public_output, - table_counts, - runtime_page_ranges, - reg_fini, - l2g_root, - }) + Ok(( + EpochProof { + proof, + public_output, + table_counts, + runtime_page_ranges, + reg_fini, + l2g_root, + }, + gkr_extras, + )) } /// Verify one epoch using ONLY the epoch's public statement fields (via @@ -738,6 +818,7 @@ fn verify_epoch( label: u64, opts: &ProofOptions, decode_commitment: Option, + gkr: Option<&GkrProofExtras>, ) -> Result { let table_counts = epoch.table_counts()?; // Reject degenerate table counts (mirrors the monolithic verifier). @@ -763,6 +844,11 @@ fn verify_epoch( let runtime_page_ranges = epoch.runtime_page_ranges()?; let public_output = epoch.public_output(); + let mode = if gkr.is_some() { + LogUpMode::Gkr + } else { + LogUpMode::Standard + }; let airs = build_epoch_airs( elf, opts, @@ -772,8 +858,9 @@ fn verify_epoch( ®_fini, is_final, decode_commitment, + mode, ); - let l2g_air = l2g_memory_air(opts, label); + let l2g_air = l2g_memory_air(opts, label, mode); let mut refs = airs.air_refs(); refs.push(&l2g_air); @@ -806,7 +893,18 @@ fn verify_epoch( None => return Ok(false), }; - if !Verifier::multi_verify_views(&refs, proof, &mut seed(), &expected) { + let verified = match gkr { + None => Verifier::multi_verify_views(&refs, proof, &mut seed(), &expected), + Some(extras) => Verifier::multi_verify_gkr_views( + &refs, + proof, + &extras.batch_gkr_proof, + &extras.column_claims_by_table, + &mut seed(), + &expected, + ), + }; + if !verified { return Ok(false); } @@ -829,7 +927,8 @@ fn prove_global( page_bases: &[u64], num_private_input_pages: usize, opts: &ProofOptions, -) -> Result, Error> { + mode: LogUpMode, +) -> GlobalProveResult { // Each cell's final state (boundaries are in epoch order, so the last fini wins). let mut final_state: global_memory::FiniStateMap = HashMap::new(); for epoch in boundaries { @@ -862,11 +961,11 @@ fn prove_global( // One L2G air per epoch, each carrying its own 1-based `fini_epoch` constant. let l2g_airs: Vec<_> = (0..boundaries.len()) - .map(|i| l2g_global_air(opts, local_to_global::epoch_label(i as u64))) + .map(|i| l2g_global_air(opts, local_to_global::epoch_label(i as u64), mode)) .collect(); let gm_airs: Vec<_> = gm_configs .iter() - .map(|config| global_memory_air(opts, config, None)) + .map(|config| global_memory_air(opts, config, None, mode)) .collect(); let mut pairs: Vec<(AirRef, &mut TraceTable, &())> = l2g_airs @@ -878,19 +977,41 @@ fn prove_global( pairs.push((air as AirRef, trace, &())); } - Prover::multi_prove( - pairs, - &mut global_transcript( - elf_bytes, - boundaries.len(), - num_private_input_pages, - opts.fri_final_poly_log_degree, - page_bases, - ), - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - ) - .map_err(|e| Error::Prover(format!("{e:?}"))) + let mut transcript = global_transcript( + elf_bytes, + boundaries.len(), + num_private_input_pages, + opts.fri_final_poly_log_degree, + page_bases, + ); + match mode { + LogUpMode::Standard => Ok(( + Prover::multi_prove( + pairs, + &mut transcript, + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?, + None, + )), + LogUpMode::Gkr => { + let gkr = Prover::multi_prove_gkr( + pairs, + &mut transcript, + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + Ok(( + gkr.multi, + Some(GkrProofExtras { + batch_gkr_proof: gkr.batch_gkr_proof, + column_claims_by_table: gkr.column_claims_by_table, + }), + )) + } + } } #[allow(clippy::too_many_arguments)] @@ -903,11 +1024,17 @@ fn verify_global( num_private_input_pages: usize, opts: &ProofOptions, page_genesis_commitments: Option<&[(u64, Commitment)]>, + gkr: Option<&GkrProofExtras>, ) -> bool { + let mode = if gkr.is_some() { + LogUpMode::Gkr + } else { + LogUpMode::Standard + }; // One L2G air per epoch, each with its own 1-based `fini_epoch` constant — // must match the order/labels the global proof committed in `prove_global`. let l2g_airs: Vec<_> = (0..num_epochs) - .map(|i| l2g_global_air(opts, local_to_global::epoch_label(i as u64))) + .map(|i| l2g_global_air(opts, local_to_global::epoch_label(i as u64), mode)) .collect(); // Rebuild the genesis configs FROM THE ELF (no private bytes) and recompute their // commitments: this is the binding for ELF/runtime pages — a prover that claimed @@ -956,7 +1083,7 @@ fn verify_global( } else { Some(zero_init_root) }; - global_memory_air(opts, config, preprocessed) + global_memory_air(opts, config, preprocessed, mode) }) .collect(); @@ -965,18 +1092,24 @@ fn verify_global( refs.push(air as AirRef); } - Verifier::multi_verify_views( - &refs, - proof, - &mut global_transcript( - elf_bytes, - num_epochs, - num_private_input_pages, - opts.fri_final_poly_log_degree, - page_bases, + let mut transcript = global_transcript( + elf_bytes, + num_epochs, + num_private_input_pages, + opts.fri_final_poly_log_degree, + page_bases, + ); + match gkr { + None => Verifier::multi_verify_views(&refs, proof, &mut transcript, &FieldElement::zero()), + Some(extras) => Verifier::multi_verify_gkr_views( + &refs, + proof, + &extras.batch_gkr_proof, + &extras.column_claims_by_table, + &mut transcript, + &FieldElement::zero(), ), - &FieldElement::zero(), - ) + } } /// Prove a full continuation and return a self-contained [`ContinuationProof`] @@ -998,6 +1131,58 @@ pub fn prove_continuation( epoch_size_log2: u32, opts: &ProofOptions, ) -> Result { + let (bundle, _extras) = prove_continuation_impl( + elf_bytes, + private_inputs, + epoch_size_log2, + opts, + LogUpMode::Standard, + )?; + Ok(bundle) +} + +/// [`prove_continuation`] under [`LogUpMode::Gkr`]: every epoch and the +/// global-memory proof prove their LogUp sums with the batch GKR; the bundle +/// carries each proof's GKR artifacts alongside the standard layout. +pub fn prove_continuation_gkr( + elf_bytes: &[u8], + private_inputs: &[u8], + epoch_size_log2: u32, + opts: &ProofOptions, +) -> Result { + let (base, extras) = prove_continuation_impl( + elf_bytes, + private_inputs, + epoch_size_log2, + opts, + LogUpMode::Gkr, + )?; + let (epoch_gkr, global_gkr) = extras.ok_or_else(|| { + Error::ContinuationInvariant("GKR prove returned no GKR artifacts".to_string()) + })?; + Ok(GkrContinuationProof { + base, + epoch_gkr, + global_gkr, + }) +} + +/// Shared driver behind [`prove_continuation`] / [`prove_continuation_gkr`]; +/// the extras component is `Some` exactly under [`LogUpMode::Gkr`]. +#[allow(clippy::type_complexity)] +fn prove_continuation_impl( + elf_bytes: &[u8], + private_inputs: &[u8], + epoch_size_log2: u32, + opts: &ProofOptions, + mode: LogUpMode, +) -> Result< + ( + ContinuationProof, + Option<(Vec, GkrProofExtras)>, + ), + Error, +> { if epoch_size_log2 < 2 { return Err(Error::InvalidContinuationEpochSize( "epoch_size_log2 must be at least 2 (4 cycles)".to_string(), @@ -1021,6 +1206,7 @@ pub fn prove_continuation( local_to_global::genesis_provenance(image.iter().map(|(a, v)| (a, v as u64))); let mut epochs: Vec = Vec::new(); + let mut epoch_extras: Vec = Vec::new(); // Full per-epoch boundaries, kept prover-local for `prove_global` (L2G traces + // final-state). Deliberately NOT stored in `EpochProof`/the bundle — `CellBoundary` // holds cell values (private-input bytes for private reads); only the value-free @@ -1098,7 +1284,12 @@ pub fn prove_continuation( register_init: ®ister_init, label, }; - let epoch = prove_epoch(&elf, elf_bytes, &start, traces, is_final, &boundary, opts)?; + let (epoch, extras) = prove_epoch( + &elf, elf_bytes, &start, traces, is_final, &boundary, opts, mode, + )?; + if let Some(extras) = extras { + epoch_extras.push(extras); + } prev_fini = Some(epoch.reg_fini.clone()); // Carry the image forward: this epoch's fini is the next epoch's init. @@ -1120,21 +1311,36 @@ pub fn prove_continuation( // SINGLE source of truth: the same page-base list drives the committed GLOBAL_MEMORY // tables and is shipped in the bundle, so the two can never diverge in set or order. let touched_page_bases = touched_page_bases(&all_boundaries); - let global = prove_global( + let (global, global_extras) = prove_global( &all_boundaries, elf_bytes, &init_page_data, &touched_page_bases, num_private_input_pages, opts, + mode, )?; - Ok(ContinuationProof { - epochs, - global, - num_private_input_pages, - touched_page_bases, - }) + let extras = match mode { + LogUpMode::Standard => None, + LogUpMode::Gkr => Some(( + epoch_extras, + global_extras.ok_or_else(|| { + Error::ContinuationInvariant( + "GKR global prove returned no GKR artifacts".to_string(), + ) + })?, + )), + }; + Ok(( + ContinuationProof { + epochs, + global, + num_private_input_pages, + touched_page_bases, + }, + extras, + )) } /// Verify a [`ContinuationProof`] using ONLY the bundle and the ELF — nothing from @@ -1185,6 +1391,7 @@ pub fn verify_continuation_with_roots( opts, decode_commitment, page_genesis_commitments, + None, )?; Ok(result.map(|(public_output, _entry_point)| public_output)) } @@ -1210,6 +1417,64 @@ pub(crate) fn verify_continuation_archived( opts, Some(decode_commitment), Some(page_genesis_commitments), + None, + ) +} + +/// [`verify_continuation`] for a [`GkrContinuationProof`]: every epoch and the +/// global proof verified with the batch-GKR replay (the GKR analogue, full +/// ELF-root recompute — the trustless host path). +pub fn verify_continuation_gkr( + elf_bytes: &[u8], + bundle: &GkrContinuationProof, + opts: &ProofOptions, +) -> Result>, Error> { + let result = verify_continuation_view( + ContinuationProofView::Owned(&bundle.base), + elf_bytes, + opts, + None, + None, + Some((&bundle.epoch_gkr, &bundle.global_gkr)), + )?; + Ok(result.map(|(public_output, _entry_point)| public_output)) +} + +/// [`verify_continuation_archived`]'s GKR counterpart, for the recursion +/// `gkr`+`continuation` guest: per-epoch/global proofs read in place via the +/// archived `base` bundle; only the KB-scale per-proof GKR artifacts are +/// deserialized. +pub(crate) fn verify_gkr_continuation_archived( + archived: &ArchivedGkrContinuationProof, + elf_bytes: &[u8], + opts: &ProofOptions, + decode_commitment: Commitment, + page_genesis_commitments: &[(u64, Commitment)], +) -> Result, u64)>, Error> { + use rkyv::rancor::Error as RkyvError; + + let epoch_gkr: Vec = archived + .epoch_gkr + .as_slice() + .iter() + .map(|e| { + rkyv::deserialize::(e).map_err(|err| { + Error::Execution(format!("rkyv deserialize epoch GKR extras failed: {err}")) + }) + }) + .collect::>()?; + let global_gkr: GkrProofExtras = rkyv::deserialize::<_, RkyvError>(&archived.global_gkr) + .map_err(|err| { + Error::Execution(format!("rkyv deserialize global GKR extras failed: {err}")) + })?; + + verify_continuation_view( + ContinuationProofView::Archived(&archived.base), + elf_bytes, + opts, + Some(decode_commitment), + Some(page_genesis_commitments), + Some((&epoch_gkr, &global_gkr)), ) } @@ -1224,6 +1489,7 @@ fn verify_continuation_view( opts: &ProofOptions, decode_commitment: Option, page_genesis_commitments: Option<&[(u64, Commitment)]>, + gkr: Option<(&[GkrProofExtras], &GkrProofExtras)>, ) -> Result, u64)>, Error> { // Bound the claimed private-input page count before using it to size/allocate AIRs // (mirrors `verify_with_options`). The count is also bound into the global proof's @@ -1245,6 +1511,18 @@ fn verify_continuation_view( return Ok(None); } + // GKR bundles must carry exactly one extras entry per epoch (structural — + // a malformed bundle, not a failed proof). + if let Some((epoch_gkr, _)) = gkr + && epoch_gkr.len() != n + { + return Err(Error::MalformedContinuationBundle(format!( + "GKR bundle has {} epoch extras for {} epochs", + epoch_gkr.len(), + n, + ))); + } + // Reject a malformed bundle up front. `reg_fini` is prover-supplied (deserialized, // untrusted) and is indexed by `NUM_REGISTER_ADDRESSES` when building each epoch's // preprocessed REGISTER commitment, so a wrong length would otherwise panic the @@ -1277,6 +1555,7 @@ fn verify_continuation_view( label, opts, decode_commitment, + gkr.map(|(epoch_gkr, _)| &epoch_gkr[index]), )? { return Ok(None); } @@ -1338,6 +1617,7 @@ fn verify_continuation_view( num_private_input_pages, opts, page_genesis_commitments, + gkr.map(|(_, global_gkr)| global_gkr), ) { return Ok(None); } diff --git a/prover/src/recursion.rs b/prover/src/recursion.rs index 7cf5b7421..7c38bd368 100644 --- a/prover/src/recursion.rs +++ b/prover/src/recursion.rs @@ -220,6 +220,47 @@ pub fn encode_continuation_guest_input( Ok(blob) } +/// The GKR continuation guest's private-input layout (the `gkr` + +/// `continuation` guest features): [`ContinuationGuestInput`] with the bundle +/// replaced by its [`crate::continuation::GkrContinuationProof`] counterpart. +/// Same magic-prefixed wire format; a blob of another layout fails the +/// bytecheck validation. +#[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct GkrContinuationGuestInput { + pub bundle: crate::continuation::GkrContinuationProof, + pub inner_elf: Vec, + pub decode_commitment: Commitment, + pub page_commitments: Vec<(u64, Commitment)>, +} + +/// Build the GKR continuation guest's private-input blob for `bundle` of +/// `inner_elf` — [`encode_continuation_guest_input`] for a +/// [`crate::continuation::GkrContinuationProof`]. Roots are precomputed over +/// the bundle's standard `base` (the touched-page set is mode-independent). +pub fn encode_gkr_continuation_guest_input( + bundle: crate::continuation::GkrContinuationProof, + inner_elf: &[u8], + opts: &ProofOptions, +) -> Result, Error> { + let (decode_commitment, page_commitments) = + crate::continuation::continuation_precomputed_commitments(inner_elf, bundle.base(), opts)?; + let input = GkrContinuationGuestInput { + bundle, + inner_elf: inner_elf.to_vec(), + decode_commitment, + page_commitments, + }; + let archive = rkyv::to_bytes::(&input) + .map_err(|e| Error::Execution(format!("rkyv encode failed: {e}")))?; + let mut blob = Vec::with_capacity(crate::RECURSION_INPUT_PREFIX_LEN + archive.len()); + blob.extend_from_slice(&crate::RECURSION_INPUT_MAGIC); + blob.extend_from_slice(&crate::RECURSION_INPUT_VERSION.to_le_bytes()); + blob.extend_from_slice(&[0u8; 4]); // reserved + debug_assert_eq!(blob.len(), crate::RECURSION_INPUT_PREFIX_LEN); + blob.extend_from_slice(&archive); + Ok(blob) +} + /// Domain tag for [`program_id`]. const PROGRAM_ID_TAG: &[u8] = b"LAMBDAVM_PROGRAM_ID_V1"; @@ -408,6 +449,60 @@ pub fn verify_continuation_and_attest( Ok(Some(attestation)) } +/// [`verify_continuation_and_attest`] for a GKR continuation bundle +/// ([`encode_gkr_continuation_guest_input`]): epochs + global proof verified +/// with the batch-GKR replay, per-table proofs read in place from the archive +/// (only the KB-scale GKR artifacts are deserialized), then the SAME +/// `program_id(elf, roots) || public_output` attestation as the standard +/// continuation path. +pub fn verify_gkr_continuation_and_attest( + blob: &[u8], + proof_options: &ProofOptions, +) -> Result>, Error> { + use rkyv::rancor::Error as RkyvError; + + let archive_bytes = crate::recursion_archive_bytes(blob).ok_or_else(|| { + Error::Execution(String::from( + "GKR continuation recursion blob: bad magic or version", + )) + })?; + let mut aligned_fallback = rkyv::util::AlignedVec::<{ crate::RECURSION_INPUT_ALIGN }>::new(); + let archive: &[u8] = + if (archive_bytes.as_ptr() as usize).is_multiple_of(crate::RECURSION_INPUT_ALIGN) { + archive_bytes + } else { + aligned_fallback.extend_from_slice(archive_bytes); + &aligned_fallback + }; + let archived = rkyv::access::(archive) + .map_err(|e| Error::Execution(format!("GKR continuation blob validation failed: {e}")))?; + + let page_commitments: Vec<(u64, Commitment)> = rkyv::deserialize::< + Vec<(u64, Commitment)>, + RkyvError, + >(&archived.page_commitments) + .map_err(|e| Error::Execution(format!("rkyv deserialize page commitments failed: {e}")))?; + let decode_commitment: Commitment = archived.decode_commitment; + let inner_elf: &[u8] = archived.inner_elf.as_slice(); + + let Some((public_output, entry_point)) = crate::continuation::verify_gkr_continuation_archived( + &archived.bundle, + inner_elf, + proof_options, + decode_commitment, + &page_commitments, + )? + else { + return Ok(None); + }; + + let digest = elf_digest(inner_elf); + let id = program_id_from_digest(&digest, entry_point, &decode_commitment, &page_commitments); + let mut attestation = id.to_vec(); + attestation.extend_from_slice(&public_output); + Ok(Some(attestation)) +} + /// Split committed attestation bytes into `(program_id, inner_public_output)`. /// `None` if too short to contain an id. pub fn split_attestation(committed: &[u8]) -> Option<([u8; 32], &[u8])> { diff --git a/prover/src/tests/recursion_smoke_test.rs b/prover/src/tests/recursion_smoke_test.rs index a162a36b2..4f07a5b5a 100644 --- a/prover/src/tests/recursion_smoke_test.rs +++ b/prover/src/tests/recursion_smoke_test.rs @@ -594,6 +594,75 @@ fn run_recursion_pipeline( ); } +/// The GKR continuation analogue of the host roundtrip tests: prove a small +/// continuation in `LogUpMode::Gkr` (multiple epochs), verify the bundle on +/// the host (owned path), then mirror the `gkr`+`continuation` guest exactly — +/// encode the `GkrContinuationGuestInput` blob, verify+attest over the +/// archive, and run the consumer check. +#[test] +fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { + let root = workspace_root(); + let empty_elf_bytes = read_guest_elf(&root, "empty"); + + // Epoch 2^6 = 64 cycles: small enough that `empty` splits into several + // epochs, exercising the cross-epoch GKR batches. + let bundle = + crate::continuation::prove_continuation_gkr(&empty_elf_bytes, &[], 6, &MIN_PROOF_OPTIONS) + .expect("GKR continuation prove should succeed"); + assert!( + bundle.num_epochs() > 1, + "fixture must split into multiple epochs to exercise the cross-epoch path \ + (got {})", + bundle.num_epochs() + ); + + // Host verify (owned path, full ELF-root recompute). + let expected_output = + crate::continuation::verify_continuation_gkr(&empty_elf_bytes, &bundle, &MIN_PROOF_OPTIONS) + .expect("verify_continuation_gkr errored") + .expect("GKR continuation bundle must verify"); + + // Guest path: blob → archived verify+attest → consumer check. + let blob = recursion::encode_gkr_continuation_guest_input( + bundle, + &empty_elf_bytes, + &MIN_PROOF_OPTIONS, + ) + .expect("encode_gkr_continuation_guest_input failed"); + assert!( + blob.len() <= executor::vm::memory::MAX_PRIVATE_INPUT_SIZE as usize, + "GKR continuation input exceeds MAX_PRIVATE_INPUT_SIZE" + ); + let attestation = recursion::verify_gkr_continuation_and_attest(&blob, &MIN_PROOF_OPTIONS) + .expect("verify_gkr_continuation_and_attest errored") + .expect("GKR continuation proof did not survive the rkyv round-trip"); + + let (id, output) = recursion::split_attestation(&attestation).expect("attestation too short"); + assert_eq!(output, expected_output, "attested output mismatch"); + // The continuation program_id folds the bundle-dependent touched-page + // genesis roots, so recompute over the bundle the consumer holds (the + // standard consumer flow for continuations). + let archive_bytes = crate::recursion_archive_bytes(&blob).expect("wire prefix"); + let mut aligned = rkyv::util::AlignedVec::<{ crate::RECURSION_INPUT_ALIGN }>::new(); + aligned.extend_from_slice(archive_bytes); + let archived = + rkyv::access::(&aligned) + .expect("blob validates"); + let owned_bundle: crate::continuation::GkrContinuationProof = + rkyv::deserialize::<_, rkyv::rancor::Error>(&archived.bundle).expect("bundle deserializes"); + let (expected_decode, expected_pages) = + crate::continuation::continuation_precomputed_commitments( + &empty_elf_bytes, + owned_bundle.base(), + &MIN_PROOF_OPTIONS, + ) + .expect("continuation_precomputed_commitments errored"); + let expected_id = + recursion::program_id_from_elf(&empty_elf_bytes, &expected_decode, &expected_pages) + .expect("program_id_from_elf errored"); + assert_eq!(id, expected_id, "attested program_id mismatch"); +} + /// The GKR analogue of `test_recursion_blob_decodes_and_verifies_on_host`: /// prove the inner in `LogUpMode::Gkr`, encode a `GkrGuestInput` blob, mirror /// the `gkr` guest's verify+attest on the host (per-table proofs read in @@ -943,9 +1012,10 @@ fn test_recursion_prove_1query() { /// continuations with `2^n`-cycle epochs and encode a /// [`recursion::ContinuationGuestInput`] blob for `recursion-cont-.elf`. /// * `RECURSION_DUMP_GKR=1` (default unset) — prove the inner in -/// `LogUpMode::Gkr` and encode a [`crate::GkrGuestInput`] blob for -/// `recursion-gkr-.elf`. Mutually exclusive with -/// `RECURSION_DUMP_EPOCH_LOG2`. +/// `LogUpMode::Gkr`: a [`crate::GkrGuestInput`] blob for +/// `recursion-gkr-.elf`, or (with `RECURSION_DUMP_EPOCH_LOG2`) a +/// [`recursion::GkrContinuationGuestInput`] blob for +/// `recursion-gkr-cont-.elf`. #[test] #[ignore = "diagnostic: writes recursion private input to /tmp/recursion_input.bin"] fn test_dump_recursion_input() { @@ -993,37 +1063,76 @@ fn test_dump_recursion_input() { .parse() .unwrap_or_else(|e| panic!("bad RECURSION_DUMP_EPOCH_LOG2 '{s}': {e}")); let opts = preset.options(); + let gkr = std::env::var("RECURSION_DUMP_GKR").is_ok_and(|v| v == "1"); eprintln!( - "[dump-input] proving inner continuation (blowup={}, fri_queries={}, epoch=2^{epoch_log2}) ...", + "[dump-input] proving inner continuation (blowup={}, fri_queries={}, epoch=2^{epoch_log2}, gkr={gkr}) ...", opts.blowup_factor, opts.fri_number_of_queries ); - let bundle = crate::continuation::prove_continuation( - &inner_elf_bytes, - &inner_input, - epoch_log2, - &opts, - ) - .expect("inner continuation prove should succeed"); - eprintln!("[dump-input] continuation epochs: {}", bundle.num_epochs()); - - let expected_output = - crate::continuation::verify_continuation(&inner_elf_bytes, &bundle, &opts) - .expect("verify_continuation errored") - .expect("continuation bundle must verify on host before dumping"); - let (expected_decode, expected_pages) = - crate::continuation::continuation_precomputed_commitments( + if gkr { + let bundle = crate::continuation::prove_continuation_gkr( + &inner_elf_bytes, + &inner_input, + epoch_log2, + &opts, + ) + .expect("inner GKR continuation prove should succeed"); + eprintln!("[dump-input] continuation epochs: {}", bundle.num_epochs()); + + let expected_output = + crate::continuation::verify_continuation_gkr(&inner_elf_bytes, &bundle, &opts) + .expect("verify_continuation_gkr errored") + .expect("GKR continuation bundle must verify on host before dumping"); + let (expected_decode, expected_pages) = + crate::continuation::continuation_precomputed_commitments( + &inner_elf_bytes, + bundle.base(), + &opts, + ) + .expect("continuation_precomputed_commitments errored"); + let expected_id = recursion::program_id_from_elf( + &inner_elf_bytes, + &expected_decode, + &expected_pages, + ) + .expect("program_id_from_elf errored"); + + let blob = + recursion::encode_gkr_continuation_guest_input(bundle, &inner_elf_bytes, &opts) + .expect("recursion::encode_gkr_continuation_guest_input failed"); + (blob, Some((expected_id, expected_output))) + } else { + let bundle = crate::continuation::prove_continuation( &inner_elf_bytes, - &bundle, + &inner_input, + epoch_log2, &opts, ) - .expect("continuation_precomputed_commitments errored"); - let expected_id = - recursion::program_id_from_elf(&inner_elf_bytes, &expected_decode, &expected_pages) - .expect("program_id_from_elf errored"); - - let blob = recursion::encode_continuation_guest_input(bundle, &inner_elf_bytes, &opts) - .expect("recursion::encode_continuation_guest_input failed"); - (blob, Some((expected_id, expected_output))) + .expect("inner continuation prove should succeed"); + eprintln!("[dump-input] continuation epochs: {}", bundle.num_epochs()); + + let expected_output = + crate::continuation::verify_continuation(&inner_elf_bytes, &bundle, &opts) + .expect("verify_continuation errored") + .expect("continuation bundle must verify on host before dumping"); + let (expected_decode, expected_pages) = + crate::continuation::continuation_precomputed_commitments( + &inner_elf_bytes, + &bundle, + &opts, + ) + .expect("continuation_precomputed_commitments errored"); + let expected_id = recursion::program_id_from_elf( + &inner_elf_bytes, + &expected_decode, + &expected_pages, + ) + .expect("program_id_from_elf errored"); + + let blob = + recursion::encode_continuation_guest_input(bundle, &inner_elf_bytes, &opts) + .expect("recursion::encode_continuation_guest_input failed"); + (blob, Some((expected_id, expected_output))) + } } Err(_) if std::env::var("RECURSION_DUMP_GKR").is_ok_and(|v| v == "1") => { let opts = preset.options(); From 96111909a0b7f54a89d81624ed2bb6805887d9af Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 16:41:51 -0300 Subject: [PATCH 12/20] test(recursion): shrink the GKR continuation fixture epoch to 2^3 cycles The empty guest finishes inside a 2^6-cycle epoch, so the multi-epoch assert fired; 2^3 matches the standard continuation tests' fixture size. --- prover/src/tests/recursion_smoke_test.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/prover/src/tests/recursion_smoke_test.rs b/prover/src/tests/recursion_smoke_test.rs index 4f07a5b5a..d464b19bb 100644 --- a/prover/src/tests/recursion_smoke_test.rs +++ b/prover/src/tests/recursion_smoke_test.rs @@ -604,10 +604,11 @@ fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { let root = workspace_root(); let empty_elf_bytes = read_guest_elf(&root, "empty"); - // Epoch 2^6 = 64 cycles: small enough that `empty` splits into several - // epochs, exercising the cross-epoch GKR batches. + // Epoch 2^3 = 8 cycles: small enough that `empty` splits into several + // epochs, exercising the cross-epoch GKR batches (the standard + // continuation tests use 2^3 with the same fixture). let bundle = - crate::continuation::prove_continuation_gkr(&empty_elf_bytes, &[], 6, &MIN_PROOF_OPTIONS) + crate::continuation::prove_continuation_gkr(&empty_elf_bytes, &[], 3, &MIN_PROOF_OPTIONS) .expect("GKR continuation prove should succeed"); assert!( bundle.num_epochs() > 1, From b9686712fc5748ed9df89d8cfed9a7e33d4c7e39 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 16:43:41 -0300 Subject: [PATCH 13/20] test(recursion): use the fibonacci fixture for the GKR continuation roundtrip The empty guest finishes inside even a 2^3-cycle epoch; fibonacci(10) at 2^4-cycle epochs splits into multiple epochs (the same fixture the standard continuation roundtrip test uses). --- prover/src/tests/recursion_smoke_test.rs | 27 ++++++++++++++---------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/prover/src/tests/recursion_smoke_test.rs b/prover/src/tests/recursion_smoke_test.rs index d464b19bb..defec8ef4 100644 --- a/prover/src/tests/recursion_smoke_test.rs +++ b/prover/src/tests/recursion_smoke_test.rs @@ -602,14 +602,19 @@ fn run_recursion_pipeline( #[test] fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { let root = workspace_root(); - let empty_elf_bytes = read_guest_elf(&root, "empty"); + let fib_elf_bytes = read_guest_elf(&root, "fibonacci"); + let inner_input = 10u64.to_le_bytes(); - // Epoch 2^3 = 8 cycles: small enough that `empty` splits into several - // epochs, exercising the cross-epoch GKR batches (the standard - // continuation tests use 2^3 with the same fixture). - let bundle = - crate::continuation::prove_continuation_gkr(&empty_elf_bytes, &[], 3, &MIN_PROOF_OPTIONS) - .expect("GKR continuation prove should succeed"); + // fibonacci(10) at 2^4-cycle epochs splits into several epochs (same + // fixture as the standard continuation roundtrip test), exercising the + // cross-epoch GKR batches. + let bundle = crate::continuation::prove_continuation_gkr( + &fib_elf_bytes, + &inner_input, + 4, + &MIN_PROOF_OPTIONS, + ) + .expect("GKR continuation prove should succeed"); assert!( bundle.num_epochs() > 1, "fixture must split into multiple epochs to exercise the cross-epoch path \ @@ -619,14 +624,14 @@ fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { // Host verify (owned path, full ELF-root recompute). let expected_output = - crate::continuation::verify_continuation_gkr(&empty_elf_bytes, &bundle, &MIN_PROOF_OPTIONS) + crate::continuation::verify_continuation_gkr(&fib_elf_bytes, &bundle, &MIN_PROOF_OPTIONS) .expect("verify_continuation_gkr errored") .expect("GKR continuation bundle must verify"); // Guest path: blob → archived verify+attest → consumer check. let blob = recursion::encode_gkr_continuation_guest_input( bundle, - &empty_elf_bytes, + &fib_elf_bytes, &MIN_PROOF_OPTIONS, ) .expect("encode_gkr_continuation_guest_input failed"); @@ -653,13 +658,13 @@ fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { rkyv::deserialize::<_, rkyv::rancor::Error>(&archived.bundle).expect("bundle deserializes"); let (expected_decode, expected_pages) = crate::continuation::continuation_precomputed_commitments( - &empty_elf_bytes, + &fib_elf_bytes, owned_bundle.base(), &MIN_PROOF_OPTIONS, ) .expect("continuation_precomputed_commitments errored"); let expected_id = - recursion::program_id_from_elf(&empty_elf_bytes, &expected_decode, &expected_pages) + recursion::program_id_from_elf(&fib_elf_bytes, &expected_decode, &expected_pages) .expect("program_id_from_elf errored"); assert_eq!(id, expected_id, "attested program_id mismatch"); } From 1bf27083c22e00e6d6443c66ef1eb1d10fa7faf1 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:09:44 -0300 Subject: [PATCH 14/20] docs(gkr): design for the linear input layer (leaf-binding fix, shape 2b) --- thoughts/logup-gkr/input-layer-design.md | 147 +++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 thoughts/logup-gkr/input-layer-design.md diff --git a/thoughts/logup-gkr/input-layer-design.md b/thoughts/logup-gkr/input-layer-design.md new file mode 100644 index 000000000..c609cda39 --- /dev/null +++ b/thoughts/logup-gkr/input-layer-design.md @@ -0,0 +1,147 @@ +# Linear input layer for LogUp-GKR (the leaf-binding fix, shape "2b") + +> Design for closing the multi-interaction leaf-binding gap (port-plan.md §6) +> by extending each table's summation tree down to per-interaction leaves — +> the Papini–Hàbock LogUp-GKR shape. Acceptance test: +> `reconstruct_multi_interaction_rejects_fabricated_leaf_claims` un-ignored +> and passing. + +## The change in one paragraph + +Today each table is ONE GKR instance over N leaves, where leaf i is the +cross-multiplied fraction `Σ_k ±m_k(i)/fp_k(i)` — a degree-K function of the +trace columns, which is why the verifier cannot check the leaf claims +(fail-open). The fix: extend the tree by `log2(K̂)` layers (K̂ = K padded to a +power of two), so the instance has `K̂·N` leaves indexed `(i, k)` with leaf +value `(±m_k(i), fp_k(i))` for k < K and the fraction identity `(0, 1)` for +padding. Leaves are now LINEAR in the trace columns, so the verifier +reconstructs the final claims EXACTLY from the column claims — no new proof +fields, no rational cross-check special cases. + +## Fact 1: the N-sized layer is bit-identical to today's tree + +Fraction-pair addition `(a,b)+(c,d) = (ad+cb, bd)` is associative at the PAIR +level (not just as rationals): any association order over +`[f_1..f_K, (0,1)…]` yields the identical `(n, d)` pair, and `(0,1)` is the +identity. Therefore the extended tree's layer at size N — the balanced +pairwise sum over each row's K̂ interaction leaves — equals EXACTLY the output +of today's `compute_logup_leaf_fractions` (a sequential fold). Consequence: + +- Every layer from size N up to the root is UNCHANGED — same values, same + materialized `gen_layers` tree, same batch sumcheck code, same memory. +- The root (bus balance) is unchanged → the cross-mode oracle + (`gkr_root_matches_standard_table_contribution`) still holds as-is. +- The ONLY new prover work is the `log2(K̂)` deep layers between K̂·N and N. + +## Fact 2: variable ordering — interaction bits LOW + +Leaf index = `i·K̂ + k` (k in the low bits). Then: + +- The deep layers pair adjacent k's, so the tree "absorbs" the interaction + sum first and reaches today's per-row fractions at size N (Fact 1 needs + this ordering). +- The final input-layer evaluation point splits as `(κ, ρ)`: κ = the low + `log2(K̂)` coordinates (interaction bits), ρ = the high `log2(N)` + coordinates (row bits). +- **The bridge is untouched.** Column claims remain row-MLEs `⟨l, col⟩` at + the ROW point ρ; the Lagrange-kernel/σ columns, their constraints, and the + aux layout stay exactly as shipped. κ never touches committed data — it + only enters the verifier's reconstruction weights (Fact 3). +- Bookkeeping: instance n_vars becomes `log2(N) + log2(K̂_t)` (varies per + table by K); the batch already handles mixed sizes. `instance_eval_point` + yields the full (κ, ρ) point; ρ = its last `log2(N)` coordinates feeds the + kernel/bridge exactly where the whole point used to. + +## Fact 3: the verifier check becomes exact linear reconstruction + +At the input layer the claims `(n̂, d̂)` are MLE evaluations of the leaf +vectors at `(κ, ρ)`. Both vectors are multilinear in (k-bits, columns): + +``` +d̂ = Σ_{k1, n_layers>0 and + expect rejection); cross-mode oracle unchanged; bridge parity unchanged; + all e2e suites (monolithic, recursion, continuation) re-run — wire format + changes only in the transcript (more sumcheck layers), so GKR-mode proofs + from before this change will NOT verify (fine: experimental mode, no + compatibility promise). + +## Costs (estimates to validate on the box) + +- Proof size: +`log2(K̂)` layers per table ≈ +6 layers × (rounds × 4 evals + + 4 child claims) for the big tables — hundreds of KB total, ≈ noise + against the −12 % we currently have. +- Guest cycles: the extra layers add sumcheck replay + the K-term + reconstruction per table — small vs the −3.3 % headroom; measure. +- Prover: Stage 1 expected somewhat slower + more transient memory than + today; Stage 2 expected ≈ today or better (leaf cross-multiplication was + the single largest GKR phase and it dissolves into cheaper on-the-fly + evals). From 6fe035d6e4e8fe1e2b8ff1f506894e4343f524b8 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:17:28 -0300 Subject: [PATCH 15/20] =?UTF-8?q?feat(gkr):=20linear=20input=20layer=20?= =?UTF-8?q?=E2=80=94=20the=20leaf-binding=20soundness=20fix=20(stage=201)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements thoughts/logup-gkr/input-layer-design.md: each table's GKR instance now extends log2(K-hat) layers below the per-row fractions, to K-hat*N leaves indexed i*K-hat + k (interaction bits LOW) with leaf (i,k) = (sign_k*m_k(i), fp_k(i)) for k < K and the fraction identity (0,1) for padding. Leaves are LINEAR in the trace columns. By pair-level associativity of fraction addition, the extended tree's layer at size N equals the cross-multiplied per-row fractions bit-for-bit (regression-tested against compute_logup_leaf_fractions), so every layer from N up, the root, the batch sumcheck code, the bridge/kernel columns, the rap-challenge layout, and the wire format are all unchanged. The instance point splits (kappa, rho): rho (the row part, same length as before) feeds the column claims and kernel/bridge; kappa only weights the verifier's reconstruction. The verifier's reconstruct_and_verify_gkr_claims becomes an EXACT check for every table and size: n = sum_k eq(kappa,k) * sign_k * m_k(claims) d = sum_k eq(kappa,k) * fp_k(claims) + (1 - sum_k eq(kappa,k)) replacing the three-way branch whose multi-interaction arm was FAIL-OPEN (the known soundness gap, port-plan.md section 6): a prover can no longer run an honest GKR over fabricated leaves. n_layers on the verifier side = log2(N) + log2(K-hat) with K-hat derived from the AIR's interactions, never the proof. The forgery test is un-ignored and extended (honest claims accepted, fabricated claims rejected, wrong-length kappa rejected). Single-row multi-interaction tables now run a real GKR over interactions (the old 0-layer special case dissolves). Stage 1 materializes the deep layers (transient O(K-hat*N) on the biggest tables) — correctness first; stage 2 (virtual deep layers per the design doc) restores today's memory profile. Gates: 291 stark tests green including all GKR e2e roundtrip/tamper tests, the cross-mode oracle (root unchanged), and the Fact-1 bit-equality test; VM-level GKR e2e + library-path tests green. --- crypto/stark/src/lagrange_kernel.rs | 2 +- crypto/stark/src/logup_gkr.rs | 308 ++++++++++++++++++----- crypto/stark/src/prover.rs | 13 +- crypto/stark/src/verifier.rs | 23 +- prover/src/tests/recursion_smoke_test.rs | 9 +- 5 files changed, 282 insertions(+), 73 deletions(-) diff --git a/crypto/stark/src/lagrange_kernel.rs b/crypto/stark/src/lagrange_kernel.rs index 099b10d9a..ca7d51261 100644 --- a/crypto/stark/src/lagrange_kernel.rs +++ b/crypto/stark/src/lagrange_kernel.rs @@ -216,7 +216,7 @@ mod tests { let mut expected = FE::one(); for (l, r_l) in r.iter().enumerate() { let bit = (i >> l) & 1; - expected *= if bit == 1 { *r_l } else { &one - r_l }; + expected *= if bit == 1 { *r_l } else { one - r_l }; } assert_eq!(*v, expected, "kernel mismatch at n={n}, i={i}"); } diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index b06b9ceac..0303bb627 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -76,6 +76,39 @@ pub const fn logup_random_point_start(num_columns: usize) -> usize { LOGUP_GAMMA_POWERS_START + num_columns + 1 } +// ============================================================================= +// Input-layer geometry (the linear input layer — see +// thoughts/logup-gkr/input-layer-design.md) +// ============================================================================= +// Each table is ONE GKR instance over K̂·N leaves indexed `i·K̂ + k` +// (interaction bits LOW): leaf (i, k) = (±m_k(i), fp_k(i)) for k < K, the +// fraction identity (0, 1) for padding. Leaves are LINEAR in the trace +// columns, so the verifier reconstructs the input-layer claims exactly from +// the column claims (no fail-open branch). By pair-level associativity of +// fraction addition, the tree's layer at size N equals the cross-multiplied +// per-row fractions bit-for-bit, so every layer from N up — and the root — +// is unchanged. + +/// Number of padded interaction variables for a table: `log2(K̂)` where +/// `K̂ = K.next_power_of_two()`. Zero for single-interaction tables (the +/// instance degenerates to the per-row tree). +pub fn gkr_input_num_vars(num_interactions: usize) -> usize { + debug_assert!(num_interactions > 0); + num_interactions.next_power_of_two().trailing_zeros() as usize +} + +/// Split an instance's full evaluation point into `(κ, ρ)`: the low +/// `input_num_vars` coordinates κ index the interaction bits (they weight the +/// verifier's claim reconstruction), the remaining coordinates ρ index the +/// rows (they are THE random point for the column claims and the +/// kernel/bridge — same length `log2(N)` as a row-only point). +pub fn split_input_point( + point: &[FieldElement], + input_num_vars: usize, +) -> (&[FieldElement], &[FieldElement]) { + point.split_at(input_num_vars) +} + // ============================================================================= // Column extraction // ============================================================================= @@ -194,24 +227,100 @@ where (numerators, denominators) } -/// Compute the full GKR layer tree for one table's interactions: leaf -/// fractions, then pairwise fraction-summation layers up to the root. -pub fn compute_logup_layers( +/// Build the LINEAR input layer: `K̂·N` leaves indexed `i·K̂ + k`, leaf +/// `(i, k) = (sign_k·m_k(i), fp_k(i))` for `k < K` and the fraction identity +/// `(0, 1)` for padding. Row-parallel; rows are read in place from the +/// row-major main table. +fn compute_gkr_input_layer( interactions: &[BusInteraction], main: &Table, trace_len: usize, challenges: &[FieldElement], -) -> Vec> +) -> Layer where F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { - let (numerators, denominators) = - compute_logup_leaf_fractions(interactions, main, trace_len, challenges); - gen_layers(Layer::LogUpGeneric { + assert!( + !interactions.is_empty(), + "the input layer requires at least one interaction" + ); + + let z = &challenges[LOGUP_CHALLENGE_Z]; + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let max_bus_elements = interactions + .iter() + .map(|inter| inter.num_bus_elements()) + .max() + .unwrap(); + let alpha_powers = compute_alpha_powers(alpha, max_bus_elements); + let shifts = PackingShifts::::new(); + let k_hat = interactions.len().next_power_of_two(); + + let row_leaves = |row_idx: usize| -> Vec<(FieldElement, FieldElement)> { + let row = main.get_row(row_idx); + let mut leaves = Vec::with_capacity(k_hat); + for inter in interactions { + let fp = compute_fingerprint_at_row(inter, row, z, &alpha_powers, &shifts); + let m = inter.multiplicity.evaluate_from(|col| row[col].clone()); + let n = if inter.is_sender { + m.to_extension() + } else { + (-m).to_extension() + }; + leaves.push((n, fp)); + } + // Padding: the fraction identity 0/1 (contributes nothing to any sum). + leaves.resize_with(k_hat, || { + (FieldElement::::zero(), FieldElement::::one()) + }); + leaves + }; + + #[cfg(feature = "parallel")] + let per_row: Vec, FieldElement)>> = + (0..trace_len).into_par_iter().map(row_leaves).collect(); + #[cfg(not(feature = "parallel"))] + let per_row: Vec, FieldElement)>> = + (0..trace_len).map(row_leaves).collect(); + + let mut numerators = Vec::with_capacity(k_hat * trace_len); + let mut denominators = Vec::with_capacity(k_hat * trace_len); + for row in per_row { + for (n, d) in row { + numerators.push(n); + denominators.push(d); + } + } + + Layer::LogUpGeneric { numerators, denominators, - }) + } +} + +/// Compute the full GKR layer tree for one table's interactions, from the +/// LINEAR input layer (`K̂·N` per-interaction leaves) up to the root. The +/// first `log2(K̂)` summation layers absorb the interaction sum; by pair +/// associativity the layer at size N equals the cross-multiplied per-row +/// fractions ([`compute_logup_leaf_fractions`]) bit-for-bit, and everything +/// above — including the root — is unchanged from the row-only tree. +pub fn compute_logup_layers( + interactions: &[BusInteraction], + main: &Table, + trace_len: usize, + challenges: &[FieldElement], +) -> Vec> +where + F: IsFFTField + IsSubFieldOf + Send + Sync, + E: IsField + Send + Sync, +{ + gen_layers(compute_gkr_input_layer( + interactions, + main, + trace_len, + challenges, + )) } // ============================================================================= @@ -471,36 +580,34 @@ impl ClaimLookup<'_, E> { } } -/// Verify a table's `column_claims` against the batch-GKR leaf claims -/// `(n_claim, d_claim)`. +/// Verify a table's `column_claims` against the batch-GKR input-layer claims +/// `(n_claim, d_claim)` — the EXACT reconstruction the linear input layer +/// affords (thoughts/logup-gkr/input-layer-design.md). /// /// Always enforced: the claim index set must EQUAL the canonical sorted /// distinct column set of the interactions (same indices, same order). /// -/// For single-interaction tables and 0-layer (single-row) instances the leaf -/// fraction is reconstructible from the column MLEs — single-interaction -/// leaves are linear in the columns, and at the empty point the MLE is the row -/// value itself — so the rational cross-check -/// `n_recon·d_claim == n_claim·d_recon` is exact and enforced. +/// The input-layer leaves `(±m_k(i), fp_k(i))` are LINEAR in the trace +/// columns, so the leaf-vector MLEs at the instance point `(κ, ρ)` factor +/// through the column MLEs at ρ (the `column_claims`, bound to the committed +/// trace by the bridge constraint): /// -/// # KNOWN SOUNDNESS GAP (multi-interaction tables) +/// ```text +/// n̂ = Σ_{k 1` with `n_layers > 0` this check is FAIL-OPEN: -/// the leaf fraction is a nonlinear (cross-multiplied) function of the -/// columns, MLE does not commute with products, and nothing else binds -/// `(n_claim, d_claim)` to the committed trace — the bridge constraint binds -/// only `column_claims`. A malicious prover can therefore run an honest GKR -/// over fabricated leaves. Every production table is multi-interaction. The -/// fix (a linear input layer or an input-layer sumcheck) is designed as the -/// immediate follow-up — see `thoughts/logup-gkr/port-plan.md` §6. Do NOT -/// treat GKR mode as production-sound until it lands. +/// (the trailing term is the padding leaves' `d = 1`, via the eq kernel's +/// partition of unity). Both are checked for EXACT equality against the +/// transcript-derived claims — every table, every size, no fail-open branch. +/// `kappa` is the κ part of the instance point (`split_input_point`). pub fn reconstruct_and_verify_gkr_claims( n_claim: &FieldElement, d_claim: &FieldElement, column_claims: &[(usize, FieldElement)], interactions: &[BusInteraction], challenges: &[FieldElement], - n_layers: usize, + kappa: &[FieldElement], ) -> bool { // Structural binding: exact index-set (and order) equality with the // canonical column list. Subsumes presence checks and pins the transcript @@ -517,6 +624,13 @@ pub fn reconstruct_and_verify_gkr_claims( return false; } + // κ must have exactly the padded-interaction bit count (a public function + // of the AIR, never proof-derived — the caller slices it off the + // transcript-derived instance point). + if kappa.len() != gkr_input_num_vars(interactions.len()) { + return false; + } + let claims = ClaimLookup(column_claims); let z = &challenges[LOGUP_CHALLENGE_Z]; let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; @@ -527,11 +641,14 @@ pub fn reconstruct_and_verify_gkr_claims( .unwrap_or(0); let alpha_powers = compute_alpha_powers(alpha, max_bus_elements); - // Reconstruct the leaf fraction from the column claims with the SAME - // cross-multiplication recurrence as `compute_logup_leaf_fractions`. - let mut running_n = FieldElement::::zero(); - let mut running_d = FieldElement::::one(); - for inter in interactions { + // eq(κ, bits(k)) weights over the padded interaction hypercube. + let eq_kappa = compute_lagrange_kernel(kappa); + + // Per-interaction linear forms on the column claims, eq-weighted. + let mut n_recon = FieldElement::::zero(); + let mut d_recon = FieldElement::::zero(); + let mut eq_sum = FieldElement::::zero(); + for (k, inter) in interactions.iter().enumerate() { let mut lc = FieldElement::::from(inter.bus_id); let mut alpha_offset = 1; for bv in &inter.values { @@ -542,19 +659,18 @@ pub fn reconstruct_and_verify_gkr_claims( } let fp = z - &lc; let m = inter.multiplicity.evaluate_from(|col| claims.get(col)); - let cross = &m * &running_d; - let cross = if inter.is_sender { cross } else { -cross }; - running_n = &running_n * &fp + cross; - running_d = &running_d * &fp; - } + let m = if inter.is_sender { m } else { -m }; - if interactions.len() == 1 || n_layers == 0 { - // Rational cross-check: n_recon/d_recon == n_claim/d_claim. - &running_n * d_claim == n_claim * &running_d - } else { - // FAIL-OPEN — see the doc comment's KNOWN SOUNDNESS GAP. - true + let eq_k = &eq_kappa[k]; + n_recon += eq_k * &m; + d_recon += eq_k * &fp; + eq_sum += eq_k.clone(); } + // Padding leaves: n = 0 (nothing), d = 1 → Σ_{k≥K} eq(κ,k)·1, computed + // via partition of unity (Σ_all eq = 1). + d_recon += FieldElement::::one() - eq_sum; + + n_recon == *n_claim && d_recon == *d_claim } #[cfg(test)] @@ -702,7 +818,7 @@ mod tests { &column_claims, &interactions, &challenges, - 0, + &[], )); let tampered = FE::from(999u64); @@ -712,7 +828,7 @@ mod tests { &column_claims, &interactions, &challenges, - 0, + &[], )); // Missing claim → reject. @@ -722,7 +838,7 @@ mod tests { &column_claims[..1], &interactions, &challenges, - 0, + &[], )); // Extra claim → reject. let mut extra = column_claims.clone(); @@ -733,10 +849,56 @@ mod tests { &extra, &interactions, &challenges, - 0, + &[], )); } + /// Fact 1 of the input-layer design: by pair-level associativity of + /// fraction addition, the extended tree's layer at size N (after the + /// `log2(K̂)` interaction-summing layers) equals the cross-multiplied + /// per-row fractions of [`compute_logup_leaf_fractions`] BIT-FOR-BIT — + /// including with padding (K = 3 → K̂ = 4 exercises the (0,1) identity). + #[test] + fn extended_tree_n_layer_matches_cross_multiplied_fractions() { + let trace_len = 8usize; + let col0: Vec = (1..=8).map(|v| FE::from(v as u64)).collect(); + let col1: Vec = (21..=28).map(|v| FE::from(v as u64)).collect(); + let col2: Vec = [1u64, 0, 2, 1, 0, 3, 1, 1] + .iter() + .map(|&v| FE::from(v)) + .collect(); + let main = Table::from_columns(vec![col0, col1, col2]); + + let interactions = vec![ + BusInteraction::sender(3u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver(3u64, Multiplicity::Column(2), Packing::Direct.columns(&[1])), + BusInteraction::sender(5u64, Multiplicity::Column(2), Packing::Word2L.columns(&[0])), + ]; + let challenges = vec![FE::from(0xDEAD_BEEFu64), FE::from(0x1234_5678u64)]; + + let (expected_n, expected_d) = + compute_logup_leaf_fractions::(&interactions, &main, trace_len, &challenges); + + let layers = compute_logup_layers::(&interactions, &main, trace_len, &challenges); + let input_vars = gkr_input_num_vars(interactions.len()); + assert_eq!(input_vars, 2, "K = 3 pads to K̂ = 4"); + let n_layer = &layers[input_vars]; + match n_layer { + Layer::LogUpGeneric { + numerators, + denominators, + } => { + assert_eq!(numerators.len(), trace_len); + assert_eq!(numerators, &expected_n, "numerators diverge at the N layer"); + assert_eq!( + denominators, &expected_d, + "denominators diverge at the N layer" + ); + } + other => panic!("unexpected layer variant: {other:?}"), + } + } + /// The cross-mode oracle: the GKR summation-tree ROOT must equal the /// standard-mode table contribution `L = Σ_rows Σ_k ±m_k/fp_k` (computed /// via the standard aux path's per-interaction term columns). This ties @@ -795,24 +957,43 @@ mod tests { assert_eq!(root_value, expected_l, "GKR root != standard-mode L"); } - /// KNOWN SOUNDNESS GAP (`thoughts/logup-gkr/port-plan.md` §6): for - /// multi-interaction tables nothing binds the batch-GKR leaf claims - /// `(n̂, d̂)` to the committed columns — the leaf fraction is nonlinear in - /// the columns, so the reconstruction check is fail-open there. This test - /// asserts the DESIRED fail-closed behavior with fabricated leaf claims - /// and honest column claims; it is `#[ignore]`d because it FAILS today. - /// The input-layer binding fix un-ignores it. + /// The leaf-binding check (formerly the KNOWN SOUNDNESS GAP of + /// port-plan.md §6, closed by the linear input layer): fabricated + /// multi-interaction leaf claims are REJECTED, and the honest + /// reconstruction is accepted — exact equality, both n̂ and d̂. #[test] - #[ignore = "documents the multi-interaction leaf-binding gap (fail-open); the input-layer fix makes this pass"] fn reconstruct_multi_interaction_rejects_fabricated_leaf_claims() { let interactions = vec![ BusInteraction::sender(1u64, Multiplicity::One, Packing::Direct.columns(&[0])), BusInteraction::receiver(2u64, Multiplicity::One, Packing::Direct.columns(&[1])), ]; let challenges = vec![FE::from(1000u64), FE::from(7u64)]; - // Honest-looking column claims, fabricated leaf claims that no leaf - // vector consistent with the columns could produce. + let alpha_powers = compute_alpha_powers(&challenges[1], 2); let column_claims = vec![(0usize, FE::from(5u64)), (1usize, FE::from(9u64))]; + // K = 2 → K̂ = 2 → one κ coordinate. + let kappa = vec![FE::from(13u64)]; + + // Honest input-layer claims: eq-weighted linear forms on the claims. + let eq = compute_lagrange_kernel(&kappa); + let fp0 = challenges[0] - (FE::from(1u64) + FE::from(5u64) * alpha_powers[1]); + let fp1 = challenges[0] - (FE::from(2u64) + FE::from(9u64) * alpha_powers[1]); + let honest_n = eq[0] * FE::one() + eq[1] * (-FE::one()); + let honest_d = eq[0] * fp0 + eq[1] * fp1; + assert!( + reconstruct_and_verify_gkr_claims( + &honest_n, + &honest_d, + &column_claims, + &interactions, + &challenges, + &kappa, + ), + "honest multi-interaction leaf claims must be accepted" + ); + + // Fabricated leaf claims that no leaf vector consistent with the + // columns could produce — must be rejected (this was the fail-open + // hole before the linear input layer). let fabricated_n = FE::from(0xBADu64); let fabricated_d = FE::from(0xC0DEu64); assert!( @@ -822,10 +1003,23 @@ mod tests { &column_claims, &interactions, &challenges, - 3, + &kappa, ), "fabricated multi-interaction leaf claims must be rejected" ); + + // A wrong-length κ (proof-shape confusion) must be rejected. + assert!( + !reconstruct_and_verify_gkr_claims( + &honest_n, + &honest_d, + &column_claims, + &interactions, + &challenges, + &[], + ), + "wrong κ length must be rejected" + ); } /// Bridge parameters: Δ·N == Σ γʲ·cⱼ + γᴷ·norm_claim, and the kernel norm diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index bc9ee36e1..b36fb21c7 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2689,8 +2689,15 @@ pub trait IsStarkProver< let results: Vec> = crate::par::par_map_collect(0..gkr_indices.len(), |k| { let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; - let n_vars = trace.num_rows().trailing_zeros() as usize; - let random_point = instance_eval_point(&shared_point, n_vars); + // Instance variables = interaction bits (low) + row bits; + // ρ (the row part) is THE random point for column claims + // and the kernel/bridge. + let input_vars = + crate::logup_gkr::gkr_input_num_vars(air.bus_interactions().len()); + let n_vars = input_vars + trace.num_rows().trailing_zeros() as usize; + let full_point = instance_eval_point(&shared_point, n_vars); + let (_kappa, rho) = + crate::logup_gkr::split_input_point(&full_point, input_vars); let (n_claim, d_claim) = final_claims[k]; let (root_n, root_d) = batch_gkr_proof.root_claims[k]; let table_contribution = root_n @@ -2700,7 +2707,7 @@ pub trait IsStarkProver< finalize_logup_gkr_result( air.bus_interactions(), &trace.main_table, - random_point, + rho.to_vec(), n_claim, d_claim, table_contribution, diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index 1b438ceb6..a52077f63 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -1347,8 +1347,10 @@ pub trait IsStarkVerifier< .map(|(idx, _)| idx) .collect(); - // Instance sizes come from the proofs' trace lengths (nonzero was - // checked in Phase A); reject non-power-of-two outright. + // Instance sizes: interaction bits (a public function of the AIR, + // never the proof) + row bits from the proofs' trace lengths + // (nonzero was checked in Phase A); reject non-power-of-two. + let mut input_vars_by_instance = Vec::with_capacity(gkr_indices.len()); let mut n_layers_by_instance = Vec::with_capacity(gkr_indices.len()); for &idx in &gkr_indices { let trace_length = trace_lengths[idx]; @@ -1356,7 +1358,10 @@ pub trait IsStarkVerifier< error!("Table {idx}: trace length {trace_length} is not a power of two"); return false; } - n_layers_by_instance.push(trace_length.trailing_zeros() as usize); + let input_vars = + crate::logup_gkr::gkr_input_num_vars(airs[idx].bus_interactions().len()); + input_vars_by_instance.push(input_vars); + n_layers_by_instance.push(input_vars + trace_length.trailing_zeros() as usize); } let (shared_point, per_instance_claims) = @@ -1377,13 +1382,16 @@ pub trait IsStarkVerifier< return false; }; let (n_claim, d_claim) = &per_instance_claims[k]; + let full_point = instance_eval_point(&shared_point, n_layers_by_instance[k]); + let (kappa, _rho) = + crate::logup_gkr::split_input_point(&full_point, input_vars_by_instance[k]); if !reconstruct_and_verify_gkr_claims( n_claim, d_claim, claims, airs[idx].bus_interactions(), &lookup_challenges, - n_layers_by_instance[k], + kappa, ) { error!("Table {idx}: GKR column-claim reconstruction failed"); return false; @@ -1406,14 +1414,17 @@ pub trait IsStarkVerifier< .as_ref() .expect("presence checked above"); let trace_length = trace_lengths[idx]; - let random_point = instance_eval_point(&shared_point, n_layers_by_instance[k]); + let full_point = instance_eval_point(&shared_point, n_layers_by_instance[k]); + // ρ (the row part) is the random point for the kernel/bridge. + let (_kappa, rho) = + crate::logup_gkr::split_input_point(&full_point, input_vars_by_instance[k]); let mut challenges = lookup_challenges.clone(); extend_rap_challenges_with_bridge( &mut challenges, claims, &gamma, trace_length, - &random_point, + rho, ); table_challenges[idx] = challenges; } diff --git a/prover/src/tests/recursion_smoke_test.rs b/prover/src/tests/recursion_smoke_test.rs index defec8ef4..77c92e3c9 100644 --- a/prover/src/tests/recursion_smoke_test.rs +++ b/prover/src/tests/recursion_smoke_test.rs @@ -629,12 +629,9 @@ fn test_gkr_continuation_blob_decodes_and_verifies_on_host() { .expect("GKR continuation bundle must verify"); // Guest path: blob → archived verify+attest → consumer check. - let blob = recursion::encode_gkr_continuation_guest_input( - bundle, - &fib_elf_bytes, - &MIN_PROOF_OPTIONS, - ) - .expect("encode_gkr_continuation_guest_input failed"); + let blob = + recursion::encode_gkr_continuation_guest_input(bundle, &fib_elf_bytes, &MIN_PROOF_OPTIONS) + .expect("encode_gkr_continuation_guest_input failed"); assert!( blob.len() <= executor::vm::memory::MAX_PRIVATE_INPUT_SIZE as usize, "GKR continuation input exceeds MAX_PRIVATE_INPUT_SIZE" From 69d24cfdf43ec00cd6e17ae109cf4bf037d4650f Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:23:59 -0300 Subject: [PATCH 16/20] =?UTF-8?q?docs(gkr):=20stage-2=20implementation=20s?= =?UTF-8?q?pec=20=E2=80=94=20streamed=20deep=20layers=20(weight-vector=20f?= =?UTF-8?q?olds)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- thoughts/logup-gkr/input-layer-design.md | 70 ++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/thoughts/logup-gkr/input-layer-design.md b/thoughts/logup-gkr/input-layer-design.md index c609cda39..5601c9c76 100644 --- a/thoughts/logup-gkr/input-layer-design.md +++ b/thoughts/logup-gkr/input-layer-design.md @@ -145,3 +145,73 @@ the box (time and peak heap vs Stage 1 and vs the pre-fix branch). today; Stage 2 expected ≈ today or better (leaf cross-multiplication was the single largest GKR phase and it dissolves into cheaper on-the-fly evals). + +## Stage 2 implementation spec (refined after Stage 1's OOM on ethrex-10tx) + +Stage 1 measured: simple-tx 32.8 s / 33.8 GB peak (pre-fix GKR: 27.3 GB on +the BIGGER 10tx workload); 10tx OOMs a 62 GiB box. The deep layers must +never be materialized — the deepest alone is K̂N/2 pairs (~13 GB for CPU). + +**Seam:** `gkr_prove_batch` takes `Vec>` instead of +`Vec>>`: + +```rust +pub struct GkrInstance { + /// Layers from size N up to the root (tree indices input_vars..), built + /// by gen_layers from the N-sized cross-multiplied fractions — EXACTLY + /// today's tree. + pub upper_layers: Vec>, + /// log2(K̂); 0 = fully materialized instance (old behavior). + pub input_num_vars: usize, + /// Streams the K̂ leaf pairs of a row range: out[r][k] = (±m_k, fp_k), + /// padding (0,1). Boxed closure built in logup_gkr.rs over + /// (interactions, main table, challenges). Send + Sync for chunked + /// parallel streaming. + pub input_oracle: Option>>, +} +``` + +**Deep layer j (child tree index input_vars − j, j = 1..=input_vars) — +per-instance state:** the k-weight vector `w: Vec<(E, E)>` (numerator/ +denominator weights... NO — weights are scalars): after r bound challenges +c_1..c_r of this layer, the partially-folded child entry at (row i, slot t) += Σ_{k in block(t)} w[k mod 2^?]·leaf — maintain `w: Vec` of length +2^(k_bits_remaining_at_round_r)… concretely: at layer j the child has +K_j = K̂/2^(j−1) k-slots per row. Sumcheck round r < log2(K_j) − ... the +k-bit rounds fold w: w'[t] = (1−c)·w[2t] + c·w[2t+1] starting from +w = [1; K_j]-shaped identity per slot pair… implement as: fold_weights +matrix W (K_j columns → shrinking), where folded_child(row, t) = +Σ_k W[t][k]·leaf(row, k) with W starting as identity blocks and each +challenge combining adjacent column groups. Store W as one Vec of len +K_j (per current slot: the run of leaf weights) — because folds only ever +combine ADJACENT full blocks, W is fully described by one weight per leaf: +w[k], with slot t owning the contiguous range [t·2^r, (t+1)·2^r). Update +per round: for k in block: new_w[k] = w[k] · ((1−c) if k's bit r == 0 else +c). So: ONE Vec of length K_j, O(K_j) update per round. Folded child +entry (row, t) = Σ_{k in slot t's range} w[k]·leaf(row, k). + +**k-bit rounds (streamed):** h(0)/h(2) accumulate over parent entries = +(row, parent slot). Per row chunk: oracle streams the K̂ leaves once per +round; for each parent slot pair compute the two folded-child values via w; +gate as usual; eq factor = eq_k(remaining k-bits) ⊗ eq_row(all row bits) — +eq_row is the standard table (size N, halved only during ROW rounds, so +during k rounds it stays full and multiplies per row); eq_k is tiny. Fold +after challenge: update w only (no table exists). Parallel over row chunks +(rayon fold/reduce as in the current PAR_SUM paths). + +**Row-bit rounds:** once k-bits of the layer are exhausted (r = log2 slots), +materialize the folded child as FOUR N-sized tables (nl/nr/dl/dr where +l/r split is now on the LOWEST ROW BIT — one more stream applying w) and +fall through to the EXISTING materialized code path for the remaining +rounds (eq pre-halving, SVO if applicable, fold_table). Free at layer end. + +**Costs:** streams per deep layer = (k-bit rounds + 1 materialization) +≈ log2(K̂)−j+2; total ≈ 21 streams for K̂ = 64, each O(N·(K fingerprints + +K̂ muls)) — chunk-parallel. Peak memory: eq_row (N) + the 4 N-sized tables +at handoff = today's footprint. + +**The gate:** Stage 2 is prover-internal only — same transcript, same +values. Same seed ⇒ BYTE-IDENTICAL proofs vs Stage 1. Differential-test on +fixtures small enough for Stage 1 (unit AIRs + the gkr_mode_tests suite), +plus the ethrex-10tx run that Stage 1 OOMs (memory gate: peak ≤ pre-fix ++ ~1 GB; correctness gate: verifies). From 57ad8cff8631f2b94ebddc6dd8249189bf9ac2b7 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:36:28 -0300 Subject: [PATCH 17/20] =?UTF-8?q?feat(gkr):=20stage=202=20=E2=80=94=20stre?= =?UTF-8?q?amed=20deep=20layers=20via=20a=20DeepLayerOracle?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch prover's instances become GkrInstance { upper_layers (the materialized >=N tree, exactly the pre-extension shape), input_num_vars, deep_oracle }. For the below-N layers of an extended-input instance, the four split tables (nl/nr/dl/dr) are materialized ONE LAYER AT A TIME directly from the oracle — per-row streaming that rebuilds the K-hat leaf pairs and folds them c levels — and dropped when that batch layer's sumcheck completes. The K-hat*N input layer and the intermediate deep layers are never resident: peak transient is the deepest layer's split tables (2x child size) instead of the whole extended tree, which OOM'd a 62 GiB box on ethrex-10tx with stage 1. Layer values are the balanced partial fraction sums either way (pair-level associativity), so the transcript is byte-identical to a fully materialized extended tree. Gated by a differential test: same seed over mixed-size padded instances => identical proof bytes, points, claims, and transcript states between GkrInstance::materialized(full extended tree) and the oracle-backed instance. A fully materialized tree is the degenerate case (input_num_vars = 0), which keeps all in-module tests and single-interaction tables on the old path unchanged. compute_logup_layers (the stage-1 builder) stays as the differential/test oracle; the prover now uses build_gkr_instance. 292 stark tests green (incl. the differential gate and all GKR e2e roundtrips, which now exercise the streamed path end-to-end). --- crypto/stark/src/gkr.rs | 257 ++++++++++++++++++++++++++------- crypto/stark/src/logup_gkr.rs | 258 +++++++++++++++++++++++++++++++++- crypto/stark/src/prover.rs | 10 +- 3 files changed, 473 insertions(+), 52 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index 5ae11d179..7168cdf7e 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -1222,6 +1222,126 @@ pub fn gkr_verify( /// Run the batch GKR prover on multiple instances (each a Vec). /// +/// Streams the deep (below-N) layers of an extended-input instance +/// (thoughts/logup-gkr/input-layer-design.md). The batch prover materializes +/// ONE deep layer's four split tables at a time — directly from this oracle, +/// freed when the layer's sumcheck completes — so the K̂·N input layer and the +/// intermediate deep layers are never resident. Layer values are the balanced +/// partial fraction sums of the per-row interaction leaves, identical to what +/// a materialized extended tree would hold (pair-level associativity), so the +/// transcript is byte-identical either way. +pub trait DeepLayerOracle: Send + Sync { + /// Padded interaction count K̂ (a power of two). + fn k_hat(&self) -> usize; + /// Trace rows N (a power of two). + fn num_rows(&self) -> usize; + /// Materialize the child layer at extended-tree index `c` + /// (`0 <= c < log2(K̂)`, size `K̂·N/2^c`) split by slot parity: + /// `(nl, nr, dl, dr)`, each of `K̂·N/2^(c+1)` entries, where entry + /// `i·S/2 + s/2` (S = per-row slots) is the balanced fraction sum of the + /// row's leaf range for slot `s` (even → l, odd → r). + #[allow(clippy::type_complexity)] + fn materialize_split( + &self, + c: usize, + ) -> ( + Vec>, + Vec>, + Vec>, + Vec>, + ); +} + +/// One batch-GKR instance: the materialized layers from size N up to the root +/// plus (for extended-input instances) the deep-layer oracle. A plain +/// materialized tree is the degenerate case (`input_num_vars = 0`). +pub struct GkrInstance<'a, E: IsField> { + /// Layers from size N (index 0) to the root (last), as from [`gen_layers`]. + pub upper_layers: Vec>, + /// Number of extended-tree layers BELOW `upper_layers[0]` (= log2(K̂)). + pub input_num_vars: usize, + /// Streams the deep layers; required iff `input_num_vars > 0`. + pub deep_oracle: Option + 'a>>, +} + +impl GkrInstance<'_, E> { + /// A fully materialized instance (no deep layers) — the pre-extension + /// shape, used by tests and as the differential oracle for the streamed + /// path. + pub fn materialized(upper_layers: Vec>) -> Self { + Self { + upper_layers, + input_num_vars: 0, + deep_oracle: None, + } + } + + /// Number of reduction steps (batch layers) this instance spans. + fn n_layers(&self) -> usize { + self.upper_layers.len() - 1 + self.input_num_vars + } + + /// The child layer's variable count at extended-tree index `t`. + fn child_n_vars(&self, t: usize) -> usize { + if t >= self.input_num_vars { + self.upper_layers[t - self.input_num_vars].n_variables() + } else { + self.upper_layers[0].n_variables() + self.input_num_vars - t + } + } +} + +/// Assemble a batch layer's [`PerInstanceTables`] from its four split tables: +/// the instance evaluation point, the eq tables (SVO split above the +/// threshold), and the fold bookkeeping. Shared by the materialized path +/// (tables split from a resident tree layer) and the deep path (tables +/// streamed from a [`DeepLayerOracle`]). +fn build_tables_from_split( + nl_table: Vec>, + nr_table: Vec>, + dl_table: Vec>, + dr_table: Vec>, + is_singles: bool, + my_parent_num_vars: usize, + current_point: &[FieldElement], +) -> PerInstanceTables { + // current_point must have at least parent_num_vars coordinates + // (it grows by max_parent_vars+1 each layer, matching the tree doubling). + debug_assert!( + current_point.len() >= my_parent_num_vars, + "current_point.len()={} < parent_num_vars={}", + current_point.len(), + my_parent_num_vars + ); + let inst_point = instance_eval_point(current_point, my_parent_num_vars); + let use_svo = my_parent_num_vars >= SVO_THRESHOLD; + let svo_suffix_len = if use_svo { my_parent_num_vars / 2 } else { 0 }; + + let (eq_table, eq_prefix, eq_suffix) = if use_svo { + let suffix = compute_eq_evals(&inst_point[..svo_suffix_len]); + let prefix = compute_eq_evals(&inst_point[svo_suffix_len..my_parent_num_vars]); + (Vec::new(), prefix, suffix) + } else { + (compute_eq_evals(&inst_point), Vec::new(), Vec::new()) + }; + + PerInstanceTables { + nl_table, + nr_table, + dl_table, + dr_table, + eq_table, + eq_correction: FieldElement::one(), + is_singles, + parent_num_vars: my_parent_num_vars, + instance_point: inst_point, + use_svo, + svo_suffix_len, + eq_prefix, + eq_suffix, + } +} + /// All instances share one sumcheck per layer via random linear combination /// (sumcheck_alpha). Instances can have different numbers of layers (different /// trace lengths). Smaller instances start participating at later layers. @@ -1235,7 +1355,7 @@ pub fn gkr_verify( /// A `BatchGkrProof`, the shared `random_point`, and per-instance `(n_claim, d_claim)`. #[allow(clippy::type_complexity)] pub fn gkr_prove_batch( - layers_per_instance: Vec>>, + layers_per_instance: Vec>, transcript: &mut impl IsTranscript, ) -> ( BatchGkrProof, @@ -1259,18 +1379,19 @@ pub fn gkr_prove_batch( ); } - // n_layers_by_instance[i] = number of tree layers - 1 = number of reduction steps + // n_layers_by_instance[i] = reduction steps: materialized upper layers + // plus the streamed deep layers. let n_layers_by_instance: Vec = layers_per_instance .iter() - .map(|layers| layers.len() - 1) + .map(|inst| inst.n_layers()) .collect(); let max_layers = *n_layers_by_instance.iter().max().unwrap(); // Extract root (numerator, denominator) for each instance let root_claims: Vec<(FieldElement, FieldElement)> = layers_per_instance .iter() - .map(|layers| { - let root = layers.last().unwrap(); + .map(|inst| { + let root = inst.upper_layers.last().unwrap(); root.try_into_output_values() .expect("root layer must be output") }) @@ -1348,7 +1469,7 @@ pub fn gkr_prove_batch( let tree_layer_idx = n_remaining - 1; // tree_layer_idx is the child layer; the parent has half the size // child has 2^k elements → parent has 2^{k-1} → parent_num_vars = k-1 - let child_n_vars = layers_per_instance[i][tree_layer_idx].n_variables(); + let child_n_vars = layers_per_instance[i].child_n_vars(tree_layer_idx); debug_assert!( child_n_vars >= 1, "child of a non-output layer must have >= 2 elements" @@ -1365,7 +1486,25 @@ pub fn gkr_prove_batch( for (idx, &i) in active_instances.iter().enumerate() { let tree_layer_idx = n_remaining - 1; - let child = &layers_per_instance[i][tree_layer_idx]; + let inst = &layers_per_instance[i]; + if tree_layer_idx < inst.input_num_vars { + // Deep trivial layer (single-row extended instance): + // materialize the 2-entry child from the oracle. + let oracle = inst + .deep_oracle + .as_ref() + .expect("extended instance requires a deep oracle"); + let (nl, nr, dl, dr) = oracle.materialize_split(tree_layer_idx); + debug_assert_eq!(nl.len(), 1, "trivial deep child has 2 entries"); + child_claims_by_instance.push([ + nl[0].clone(), + nr[0].clone(), + dl[0].clone(), + dr[0].clone(), + ]); + continue; + } + let child = &inst.upper_layers[tree_layer_idx - inst.input_num_vars]; let (child_n, child_d) = match child { Layer::LogUpGeneric { @@ -1456,7 +1595,31 @@ pub fn gkr_prove_batch( let build_instance_tables = |&i: &usize| { let tree_layer_idx = n_remaining - 1; - let child = &layers_per_instance[i][tree_layer_idx]; + let inst = &layers_per_instance[i]; + + // Deep (below-N) layer of an extended instance: materialize + // THIS layer's four split tables straight from the oracle — + // the child layer itself is never resident, and these tables + // are dropped when the batch layer's sumcheck completes. + if tree_layer_idx < inst.input_num_vars { + let oracle = inst + .deep_oracle + .as_ref() + .expect("extended instance requires a deep oracle"); + let (nl_table, nr_table, dl_table, dr_table) = + oracle.materialize_split(tree_layer_idx); + return build_tables_from_split( + nl_table, + nr_table, + dl_table, + dr_table, + false, + parent_num_vars_by_instance + [active_instances.iter().position(|&x| x == i).unwrap()], + ¤t_point, + ); + } + let child = &inst.upper_layers[tree_layer_idx - inst.input_num_vars]; let parent_size = match child { Layer::LogUpGeneric { denominators, .. } @@ -1483,41 +1646,15 @@ pub fn gkr_prove_batch( let my_parent_num_vars = parent_num_vars_by_instance [active_instances.iter().position(|&x| x == i).unwrap()]; - // current_point must have at least parent_num_vars coordinates - // (it grows by max_parent_vars+1 each layer, matching the tree doubling). - debug_assert!( - current_point.len() >= my_parent_num_vars, - "current_point.len()={} < parent_num_vars={}", - current_point.len(), - my_parent_num_vars - ); - let inst_point = instance_eval_point(¤t_point, my_parent_num_vars); - let use_svo = my_parent_num_vars >= SVO_THRESHOLD; - let svo_suffix_len = if use_svo { my_parent_num_vars / 2 } else { 0 }; - - let (eq_table, eq_prefix, eq_suffix) = if use_svo { - let suffix = compute_eq_evals(&inst_point[..svo_suffix_len]); - let prefix = compute_eq_evals(&inst_point[svo_suffix_len..my_parent_num_vars]); - (Vec::new(), prefix, suffix) - } else { - (compute_eq_evals(&inst_point), Vec::new(), Vec::new()) - }; - - PerInstanceTables { + build_tables_from_split( nl_table, nr_table, dl_table, dr_table, - eq_table, - eq_correction: FieldElement::one(), is_singles, - parent_num_vars: my_parent_num_vars, - instance_point: inst_point, - use_svo, - svo_suffix_len, - eq_prefix, - eq_suffix, - } + my_parent_num_vars, + ¤t_point, + ) }; #[cfg(feature = "parallel")] @@ -3179,8 +3316,13 @@ mod tests { (0..3).map(|_| gen_layers(make_generic_leaf(2))).collect(); let mut prover_transcript = DefaultTranscript::::new(&[]); - let (proof, shared_point, final_claims) = - gkr_prove_batch(instances, &mut prover_transcript); + let (proof, shared_point, final_claims) = gkr_prove_batch( + instances + .into_iter() + .map(GkrInstance::materialized) + .collect(), + &mut prover_transcript, + ); assert_eq!(proof.root_claims.len(), 3); assert_eq!(final_claims.len(), 3); @@ -3208,7 +3350,13 @@ mod tests { let instances: Vec>> = (0..2).map(|_| gen_layers(make_generic_leaf(2))).collect(); let mut prover_transcript = DefaultTranscript::::new(&[]); - let (proof, _, _) = gkr_prove_batch(instances, &mut prover_transcript); + let (proof, _, _) = gkr_prove_batch( + instances + .into_iter() + .map(GkrInstance::materialized) + .collect(), + &mut prover_transcript, + ); let n_layers = vec![2usize, 2]; // rkyv roundtrip → verifies. @@ -3264,8 +3412,13 @@ mod tests { assert_eq!(n_layers, vec![2, 4, 6]); let mut prover_transcript = DefaultTranscript::::new(&[]); - let (proof, shared_point, final_claims) = - gkr_prove_batch(instances, &mut prover_transcript); + let (proof, shared_point, final_claims) = gkr_prove_batch( + instances + .into_iter() + .map(GkrInstance::materialized) + .collect(), + &mut prover_transcript, + ); assert_eq!(proof.root_claims.len(), 3); @@ -3295,8 +3448,13 @@ mod tests { let n_layers: Vec = instances.iter().map(|l| l.len() - 1).collect(); let mut prover_transcript = DefaultTranscript::::new(&[]); - let (proof, shared_point, final_claims) = - gkr_prove_batch(instances, &mut prover_transcript); + let (proof, shared_point, final_claims) = gkr_prove_batch( + instances + .into_iter() + .map(GkrInstance::materialized) + .collect(), + &mut prover_transcript, + ); let mut verifier_transcript = DefaultTranscript::::new(&[]); let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); @@ -3325,8 +3483,13 @@ mod tests { let n_layers: Vec = instances.iter().map(|l| l.len() - 1).collect(); let mut prover_transcript = DefaultTranscript::::new(&[]); - let (proof, shared_point, final_claims) = - gkr_prove_batch(instances, &mut prover_transcript); + let (proof, shared_point, final_claims) = gkr_prove_batch( + instances + .into_iter() + .map(GkrInstance::materialized) + .collect(), + &mut prover_transcript, + ); let mut verifier_transcript = DefaultTranscript::::new(&[]); let result = gkr_verify_batch(&proof, &n_layers, &mut verifier_transcript); diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index 0303bb627..eafc3f0b2 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -24,7 +24,7 @@ use math::field::{ #[cfg(feature = "parallel")] use rayon::prelude::*; -use crate::gkr::{Layer, gen_layers}; +use crate::gkr::{DeepLayerOracle, GkrInstance, Layer, gen_layers}; use crate::lagrange_kernel::compute_lagrange_kernel; use crate::lookup::{BusInteraction, LOGUP_CHALLENGE_ALPHA, PackingShifts, compute_alpha_powers}; use crate::table::Table; @@ -299,6 +299,204 @@ where } } +/// The K̂ leaf pairs of one row: `(sign_k·m_k, fp_k)` for `k < K`, the +/// fraction identity `(0, 1)` for padding. Shared by the Stage-1 materialized +/// input layer and the deep-layer oracle. +fn row_leaf_pairs( + interactions: &[BusInteraction], + row: &[FieldElement], + z: &FieldElement, + alpha_powers: &[FieldElement], + shifts: &PackingShifts, + k_hat: usize, +) -> Vec<(FieldElement, FieldElement)> +where + F: IsField + IsSubFieldOf, + E: IsField, +{ + let mut leaves = Vec::with_capacity(k_hat); + for inter in interactions { + let fp = compute_fingerprint_at_row(inter, row, z, alpha_powers, shifts); + let m = inter.multiplicity.evaluate_from(|col| row[col].clone()); + let n = if inter.is_sender { + m.to_extension() + } else { + (-m).to_extension() + }; + leaves.push((n, fp)); + } + leaves.resize_with(k_hat, || { + (FieldElement::::zero(), FieldElement::::one()) + }); + leaves +} + +/// [`DeepLayerOracle`] over a table's interactions and (borrowed) row-major +/// main trace: materializes one deep layer's four split tables at a time by +/// streaming rows — the K̂·N input layer is never resident. +struct LogUpDeepOracle<'a, F: IsField, E: IsField> { + interactions: &'a [BusInteraction], + main: &'a Table, + z: FieldElement, + alpha_powers: Vec>, + shifts: PackingShifts, + k_hat: usize, + trace_len: usize, +} + +impl DeepLayerOracle for LogUpDeepOracle<'_, F, E> +where + F: IsFFTField + IsSubFieldOf + Send + Sync, + E: IsField + Send + Sync, +{ + fn k_hat(&self) -> usize { + self.k_hat + } + + fn num_rows(&self) -> usize { + self.trace_len + } + + #[allow(clippy::type_complexity)] + fn materialize_split( + &self, + c: usize, + ) -> ( + Vec>, + Vec>, + Vec>, + Vec>, + ) { + debug_assert!(c < self.k_hat.trailing_zeros() as usize); + // Per-row slots at child level c, and per-row entries per split table. + let slots = self.k_hat >> c; + let half = slots / 2; + let out_len = self.trace_len * half; + + let mut nl = vec![FieldElement::::zero(); out_len]; + let mut nr = vec![FieldElement::::zero(); out_len]; + let mut dl = vec![FieldElement::::zero(); out_len]; + let mut dr = vec![FieldElement::::zero(); out_len]; + + // Each row owns the contiguous span [row·half, (row+1)·half) in every + // table: zip per-row chunks so the fill is embarrassingly parallel. + let fill_row = |row_idx: usize, + nl_c: &mut [FieldElement], + nr_c: &mut [FieldElement], + dl_c: &mut [FieldElement], + dr_c: &mut [FieldElement]| { + let row = self.main.get_row(row_idx); + let mut leaves = row_leaf_pairs( + self.interactions, + row, + &self.z, + &self.alpha_powers, + &self.shifts, + self.k_hat, + ); + // Fold c levels of pairwise fraction addition (balanced, matching + // the extended tree's deep layers bit-for-bit). + for _ in 0..c { + let m = leaves.len() / 2; + for t in 0..m { + let (n0, d0) = leaves[2 * t].clone(); + let (n1, d1) = leaves[2 * t + 1].clone(); + leaves[t] = (&(&n0 * &d1) + &(&n1 * &d0), &d0 * &d1); + } + leaves.truncate(m); + } + for t in 0..half { + let (n_even, d_even) = leaves[2 * t].clone(); + let (n_odd, d_odd) = leaves[2 * t + 1].clone(); + nl_c[t] = n_even; + dl_c[t] = d_even; + nr_c[t] = n_odd; + dr_c[t] = d_odd; + } + }; + + #[cfg(feature = "parallel")] + nl.par_chunks_mut(half) + .zip(nr.par_chunks_mut(half)) + .zip(dl.par_chunks_mut(half)) + .zip(dr.par_chunks_mut(half)) + .enumerate() + .for_each(|(row_idx, (((nl_c, nr_c), dl_c), dr_c))| { + fill_row(row_idx, nl_c, nr_c, dl_c, dr_c) + }); + #[cfg(not(feature = "parallel"))] + for row_idx in 0..self.trace_len { + let span = row_idx * half..(row_idx + 1) * half; + let (nl_c, nr_c, dl_c, dr_c) = ( + &mut nl[span.clone()], + &mut nr[span.clone()], + &mut dl[span.clone()], + &mut dr[span.clone()], + ); + // Split borrows: fall back to per-row temporary buffers. + let mut nl_t = nl_c.to_vec(); + let mut nr_t = nr_c.to_vec(); + let mut dl_t = dl_c.to_vec(); + let mut dr_t = dr_c.to_vec(); + fill_row(row_idx, &mut nl_t, &mut nr_t, &mut dl_t, &mut dr_t); + nl[span.clone()].clone_from_slice(&nl_t); + nr[span.clone()].clone_from_slice(&nr_t); + dl[span.clone()].clone_from_slice(&dl_t); + dr[span].clone_from_slice(&dr_t); + } + + (nl, nr, dl, dr) + } +} + +/// Build a table's batch-GKR instance (Stage 2 of the input-layer design): +/// the materialized layers from the N-sized cross-multiplied fractions up to +/// the root, plus the deep-layer oracle streaming the below-N layers. Peak +/// memory is one deep layer's split tables at a time — the K̂·N input layer is +/// never resident. Transcript-identical to a fully materialized extended tree. +pub fn build_gkr_instance<'a, F, E>( + interactions: &'a [BusInteraction], + main: &'a Table, + trace_len: usize, + challenges: &[FieldElement], +) -> GkrInstance<'a, E> +where + F: IsFFTField + IsSubFieldOf + Send + Sync + 'a, + E: IsField + Send + Sync + 'a, +{ + let (numerators, denominators) = + compute_logup_leaf_fractions(interactions, main, trace_len, challenges); + let upper_layers = gen_layers(Layer::LogUpGeneric { + numerators, + denominators, + }); + let input_num_vars = gkr_input_num_vars(interactions.len()); + let deep_oracle: Option + 'a>> = if input_num_vars > 0 { + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let max_bus_elements = interactions + .iter() + .map(|inter| inter.num_bus_elements()) + .max() + .unwrap(); + Some(Box::new(LogUpDeepOracle { + interactions, + main, + z: challenges[LOGUP_CHALLENGE_Z].clone(), + alpha_powers: compute_alpha_powers(alpha, max_bus_elements), + shifts: PackingShifts::::new(), + k_hat: interactions.len().next_power_of_two(), + trace_len, + })) + } else { + None + }; + GkrInstance { + upper_layers, + input_num_vars, + deep_oracle, + } +} + /// Compute the full GKR layer tree for one table's interactions, from the /// LINEAR input layer (`K̂·N` per-interaction leaves) up to the root. The /// first `log2(K̂)` summation layers absorb the interaction sum; by pair @@ -899,6 +1097,64 @@ mod tests { } } + /// THE Stage-2 gate: the streamed deep-layer path + /// ([`build_gkr_instance`], oracle-backed) must be TRANSCRIPT-IDENTICAL + /// to a fully materialized extended tree ([`compute_logup_layers`]) — + /// same seed, byte-identical proof, same point, same claims. Mixed sizes + /// and padding included (K = 3 → K̂ = 4, two trace lengths). + #[test] + fn streamed_deep_layers_match_materialized_extended_tree() { + use crate::gkr::{GkrInstance, gkr_prove_batch}; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use crypto::fiat_shamir::is_transcript::IsTranscript; + + let mk_table = |rows: usize, seed: u64| { + let col0: Vec = (0..rows).map(|v| FE::from(seed + v as u64)).collect(); + let col1: Vec = (0..rows).map(|v| FE::from(seed + 100 + v as u64)).collect(); + let col2: Vec = (0..rows).map(|v| FE::from((v as u64) % 3)).collect(); + Table::from_columns(vec![col0, col1, col2]) + }; + let interactions = vec![ + BusInteraction::sender(3u64, Multiplicity::One, Packing::Direct.columns(&[0])), + BusInteraction::receiver(3u64, Multiplicity::Column(2), Packing::Direct.columns(&[1])), + BusInteraction::sender(5u64, Multiplicity::Column(2), Packing::Word2L.columns(&[0])), + ]; + let challenges = vec![FE::from(0xDEAD_BEEFu64), FE::from(0x1234_5678u64)]; + + let table_a = mk_table(16, 7); + let table_b = mk_table(8, 1_000); + + // Materialized extended trees (the Stage-1 shape). + let mat_a = compute_logup_layers::(&interactions, &table_a, 16, &challenges); + let mat_b = compute_logup_layers::(&interactions, &table_b, 8, &challenges); + let mut t1 = DefaultTranscript::::new(&[42]); + let (proof_m, point_m, claims_m) = gkr_prove_batch( + vec![ + GkrInstance::materialized(mat_a), + GkrInstance::materialized(mat_b), + ], + &mut t1, + ); + + // Streamed instances (the Stage-2 shape). + let inst_a = build_gkr_instance::(&interactions, &table_a, 16, &challenges); + let inst_b = build_gkr_instance::(&interactions, &table_b, 8, &challenges); + let mut t2 = DefaultTranscript::::new(&[42]); + let (proof_s, point_s, claims_s) = gkr_prove_batch(vec![inst_a, inst_b], &mut t2); + + assert_eq!(point_m, point_s, "shared random points diverge"); + assert_eq!(claims_m, claims_s, "instance claims diverge"); + let bytes_m = rkyv::to_bytes::(&proof_m).unwrap(); + let bytes_s = rkyv::to_bytes::(&proof_s).unwrap(); + assert_eq!( + bytes_m.as_slice(), + bytes_s.as_slice(), + "streamed and materialized proofs are not byte-identical" + ); + // The transcripts must have absorbed identical data too. + assert_eq!(t1.state(), t2.state(), "transcript states diverge"); + } + /// The cross-mode oracle: the GKR summation-tree ROOT must equal the /// standard-mode table contribution `L = Σ_rows Σ_k ±m_k/fp_k` (computed /// via the standard aux path's per-interaction term columns). This ties diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b36fb21c7..cd4ef4748 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -29,8 +29,7 @@ use crate::debug::validate_trace; use crate::fri; use crate::gkr::{BatchGkrProof, gkr_prove_batch, instance_eval_point}; use crate::logup_gkr::{ - LogUpGkrResult, compute_logup_layers, extend_rap_challenges_with_bridge, - finalize_logup_gkr_result, + LogUpGkrResult, extend_rap_challenges_with_bridge, finalize_logup_gkr_result, }; use crate::lookup::{LOGUP_NUM_CHALLENGES, LogUpMode}; use crate::proof::stark::{DeepPolynomialOpenings, GkrMultiProof, PolynomialOpenings}; @@ -2661,10 +2660,13 @@ pub trait IsStarkProver< // committed term columns lives. #[cfg(feature = "instruments")] let __sp_gkr = crate::instruments::span("r1_gkr_trees"); - let layers_per_instance: Vec>> = + // Stage 2 of the input-layer design: instances carry only the + // >=N tree plus a deep-layer oracle — the K-hat*N input layer is + // never resident (one deep layer's split tables at a time). + let layers_per_instance: Vec> = crate::par::par_map_collect(0..gkr_indices.len(), |k| { let (air, trace, _) = &air_trace_pairs[gkr_indices[k]]; - compute_logup_layers( + crate::logup_gkr::build_gkr_instance( air.bus_interactions(), &trace.main_table, trace.num_rows(), From 837b458e3442892c256f83e12b9f3dccdb5a1e04 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:44:32 -0300 Subject: [PATCH 18/20] =?UTF-8?q?perf(gkr):=20stream=20the=20deep=20k-bit?= =?UTF-8?q?=20rounds=20=E2=80=94=20nothing=20below=20N=20is=20materialized?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-layer materialization of the previous commit still held one deep layer's split tables (2x child size — up to ~13 GB per big table) and cost ~30 s of huge-table sumcheck on ethrex-10tx. Now the k-bit rounds of a deep layer materialize NOTHING: gate sums stream per row through the oracle (rebuild the K-hat leaves, fold c tree levels + the bound challenges, apply gate_generic's exact formulas), with eq factored as a tiny k-part (pre-halved per round) times the untouched row-part. Once every k-bit is bound, the four tables materialize at N size — absorbing the folded k-eq into eq_correction and handing the row rounds to the unchanged fold path. kp == 0 layers materialize directly (already N-sized). Peak transient per deep layer drops from O(K-hat*N) to O(N) (the handoff tables + eq_row), the same order as the resident >=N tree. Transcript-identical by construction; the stage-2 differential test (byte-identical proofs vs the materialized extended tree) now pins the streamed rounds, the kp == 0 path, and the handoff. 292 stark tests green. --- crypto/stark/src/gkr.rs | 220 +++++++++++++++++++++++++++++----- crypto/stark/src/logup_gkr.rs | 168 ++++++++++++++++++++++++++ 2 files changed, 359 insertions(+), 29 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index 7168cdf7e..a07826cff 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -1250,6 +1250,37 @@ pub trait DeepLayerOracle: Send + Sync { Vec>, Vec>, ); + + /// Streamed gate sums for a k-bit round of the deep layer at child index + /// `c`: `(h0, h2)` = Σ over gate pairs of `eq_k[pair] · eq_row[row] · + /// gate_t`, where the four per-row split arrays are folded by `bound` + /// (this layer's challenges so far) and the gate is EXACTLY + /// `gate_generic`'s: `g0 = nl_l·dr_l + dl_l·(nr_l + λ·dr_l)` at the pair's + /// left entries, `g2` on the `2·r − l` extrapolations. `eq_k` is the + /// pre-halved k-part eq table (len = pairs per row). + fn deep_gate_sums( + &self, + c: usize, + bound: &[FieldElement], + eq_k: &[FieldElement], + eq_row: &[FieldElement], + lambda: &FieldElement, + ) -> (FieldElement, FieldElement); + + /// Materialize the deep layer's four split tables AFTER all k-bit rounds: + /// each per-row array folded by the full `bound` vector down to one entry + /// → four N-sized tables, handed to the ordinary row-round fold path. + #[allow(clippy::type_complexity)] + fn materialize_folded_rows( + &self, + c: usize, + bound: &[FieldElement], + ) -> ( + Vec>, + Vec>, + Vec>, + Vec>, + ); } /// One batch-GKR instance: the materialized layers from size N up to the root @@ -1296,7 +1327,7 @@ impl GkrInstance<'_, E> { /// threshold), and the fold bookkeeping. Shared by the materialized path /// (tables split from a resident tree layer) and the deep path (tables /// streamed from a [`DeepLayerOracle`]). -fn build_tables_from_split( +fn build_tables_from_split<'a, E: IsField>( nl_table: Vec>, nr_table: Vec>, dl_table: Vec>, @@ -1304,7 +1335,7 @@ fn build_tables_from_split( is_singles: bool, my_parent_num_vars: usize, current_point: &[FieldElement], -) -> PerInstanceTables { +) -> PerInstanceTables<'a, E> { // current_point must have at least parent_num_vars coordinates // (it grows by max_parent_vars+1 each layer, matching the tree doubling). debug_assert!( @@ -1326,6 +1357,7 @@ fn build_tables_from_split( }; PerInstanceTables { + deep: None, nl_table, nr_table, dl_table, @@ -1597,27 +1629,68 @@ pub fn gkr_prove_batch( let tree_layer_idx = n_remaining - 1; let inst = &layers_per_instance[i]; - // Deep (below-N) layer of an extended instance: materialize - // THIS layer's four split tables straight from the oracle — - // the child layer itself is never resident, and these tables - // are dropped when the batch layer's sumcheck completes. + // Deep (below-N) layer of an extended instance: nothing is + // materialized during the k-bit rounds — gate sums stream + // through the oracle; the four split tables appear only at + // N size once every k-bit is bound (or immediately when the + // layer has no k-bit rounds left, kp == 0). The K̂·N input + // layer and the intermediate deep layers are never resident. if tree_layer_idx < inst.input_num_vars { - let oracle = inst + let oracle: &dyn DeepLayerOracle = inst .deep_oracle .as_ref() - .expect("extended instance requires a deep oracle"); - let (nl_table, nr_table, dl_table, dr_table) = - oracle.materialize_split(tree_layer_idx); - return build_tables_from_split( - nl_table, - nr_table, - dl_table, - dr_table, - false, - parent_num_vars_by_instance - [active_instances.iter().position(|&x| x == i).unwrap()], - ¤t_point, + .expect("extended instance requires a deep oracle") + .as_ref(); + let my_parent_num_vars = parent_num_vars_by_instance + [active_instances.iter().position(|&x| x == i).unwrap()]; + let inst_point = instance_eval_point(¤t_point, my_parent_num_vars); + // Per-row slots at this child level; k-bit rounds kp. + let slots = oracle.k_hat() >> tree_layer_idx; + let kp = (slots / 2).trailing_zeros() as usize; + debug_assert_eq!( + my_parent_num_vars, + kp + oracle.num_rows().trailing_zeros() as usize, ); + if kp == 0 { + // No k-bit rounds: the split tables are already + // N-sized — materialize now, ordinary path after. + let (nl_table, nr_table, dl_table, dr_table) = + oracle.materialize_split(tree_layer_idx); + return build_tables_from_split( + nl_table, + nr_table, + dl_table, + dr_table, + false, + my_parent_num_vars, + ¤t_point, + ); + } + let eq_k = compute_eq_evals(&inst_point[..kp]); + let eq_row = compute_eq_evals(&inst_point[kp..my_parent_num_vars]); + return PerInstanceTables { + deep: Some(DeepPhase { + oracle, + c: tree_layer_idx, + kp, + bound: Vec::with_capacity(kp), + eq_k, + eq_row, + }), + nl_table: Vec::new(), + nr_table: Vec::new(), + dl_table: Vec::new(), + dr_table: Vec::new(), + eq_table: Vec::new(), + eq_correction: FieldElement::one(), + is_singles: false, + parent_num_vars: my_parent_num_vars, + instance_point: inst_point, + use_svo: false, + svo_suffix_len: 0, + eq_prefix: Vec::new(), + eq_suffix: Vec::new(), + }; } let child = &inst.upper_layers[tree_layer_idx - inst.input_num_vars]; @@ -1658,12 +1731,12 @@ pub fn gkr_prove_batch( }; #[cfg(feature = "parallel")] - let mut per_instance_tables: Vec> = active_instances + let mut per_instance_tables: Vec> = active_instances .par_iter() .map(build_instance_tables) .collect(); #[cfg(not(feature = "parallel"))] - let mut per_instance_tables: Vec> = + let mut per_instance_tables: Vec> = active_instances.iter().map(build_instance_tables).collect(); let mut round_polys = Vec::with_capacity(max_parent_vars); @@ -1697,7 +1770,7 @@ pub fn gkr_prove_batch( // parallel across instances, with parallel inner sums for the // large early rounds. let compute_instance = |idx: usize, - tables: &mut PerInstanceTables| + tables: &mut PerInstanceTables<'_, E>| -> [FieldElement; 4] { let n_unused = max_parent_vars - tables.parent_num_vars; @@ -1713,6 +1786,44 @@ pub fn gkr_prove_batch( } let instance_round = round_idx - n_unused; + + // Streamed deep k-bit round: pre-halve the k-part eq and + // let the oracle compute the raw gate sums; the S(t) + // recovery below is shared with the materialized path. + if let Some(deep) = tables.deep.as_mut() { + let eq_half = deep.eq_k.len() / 2; + for j in 0..eq_half { + deep.eq_k[j] = &deep.eq_k[2 * j] + &deep.eq_k[2 * j + 1]; + } + deep.eq_k.truncate(eq_half); + + let (raw_h0, raw_h2) = deep.oracle.deep_gate_sums( + deep.c, + &deep.bound, + &deep.eq_k, + &deep.eq_row, + &lambda, + ); + + let r_round = tables.instance_point[instance_round].clone(); + let one = FieldElement::::one(); + let total_h0 = &tables.eq_correction * &raw_h0; + let total_h2 = &tables.eq_correction * &raw_h2; + let one_minus_r = &one - &r_round; + let s0 = &one_minus_r * &total_h0; + let s1 = &combined_claims[idx] - &s0; + let r_inv = r_round.inv().expect("r_round = 0 is probability 2^{-64}"); + let h1 = &s1 * &r_inv; + let three = FieldElement::::from(3u64); + let h3 = &(&(&three * &total_h2) - &(&three * &h1)) + &total_h0; + let eq_at_2 = &(&three * &r_round) - &one; + let s2 = &eq_at_2 * &total_h2; + let eq_at_3 = &(&FieldElement::::from(5u64) * &r_round) + - &FieldElement::::from(2u64); + let s3 = &eq_at_3 * &h3; + return [s0, s1, s2, s3]; + } + let half = tables.nl_table.len().max(tables.dl_table.len()) / 2; if half == 0 { @@ -1731,7 +1842,7 @@ pub fn gkr_prove_batch( let one = FieldElement::::one(); // Helper: compute gate h(0) and h(2) for a generic pair at index j. - let gate_generic = |tables: &PerInstanceTables, + let gate_generic = |tables: &PerInstanceTables<'_, E>, j: usize| -> [FieldElement; 2] { let nl_l = &tables.nl_table[2 * j]; @@ -1753,7 +1864,7 @@ pub fn gkr_prove_batch( }; let gate_singles = - |tables: &PerInstanceTables, j: usize| -> [FieldElement; 2] { + |tables: &PerInstanceTables<'_, E>, j: usize| -> [FieldElement; 2] { let dl_l = &tables.dl_table[2 * j]; let dl_r = &tables.dl_table[2 * j + 1]; let dr_l = &tables.dr_table[2 * j]; @@ -1783,7 +1894,7 @@ pub fn gkr_prove_batch( // Eq mutations are done; sum over an immutable view so the // gate closures can run in parallel. - let view: &PerInstanceTables = tables; + let view: &PerInstanceTables<'_, E> = tables; let suffix_sum = |suffix_idx: usize| -> (FieldElement, FieldElement) { let eq_s = &view.eq_suffix[suffix_idx]; let mut ch0 = FieldElement::::zero(); @@ -1849,7 +1960,7 @@ pub fn gkr_prove_batch( tables.eq_table.truncate(half); // Eq mutations are done; sum over an immutable view. - let view: &PerInstanceTables = tables; + let view: &PerInstanceTables<'_, E> = tables; let pair_sum = |j: usize| -> (FieldElement, FieldElement) { let eq_rem = &view.eq_table[j]; let [g0, g2] = if view.is_singles { @@ -1956,7 +2067,7 @@ pub fn gkr_prove_batch( // update combined_claims via O(1) polynomial evaluation (not // O(N) table scan). Instances are independent — parallel. let update_instance = - |tables: &mut PerInstanceTables, + |tables: &mut PerInstanceTables<'_, E>, claim: &mut FieldElement, evals: &[FieldElement; 4]| { // S_i(challenge) via degree-3 Lagrange interpolation at {0,1,2,3}. @@ -1970,11 +2081,43 @@ pub fn gkr_prove_batch( *claim = instance_poly.evaluate(&challenge); let n_unused = max_parent_vars - tables.parent_num_vars; - if round_idx < n_unused || tables.dl_table.len() / 2 == 0 { + if round_idx < n_unused { return; } let instance_round = round_idx - n_unused; + + // Streamed deep k-bit round: bind the challenge, update + // eq_correction, and — once every k-bit is bound — + // materialize the four N-sized tables and hand off to + // the ordinary row-round fold path below (next round). + if let Some(deep) = tables.deep.as_mut() { + let r_round = &tables.instance_point[instance_round]; + let one = FieldElement::::one(); + let eq_update = &(r_round * &challenge) + + &(&(&one - r_round) * &(&one - &challenge)); + tables.eq_correction = &tables.eq_correction * &eq_update; + + deep.bound.push(challenge.clone()); + if deep.bound.len() == deep.kp { + debug_assert_eq!(deep.eq_k.len(), 1); + tables.eq_correction = &tables.eq_correction * &deep.eq_k[0]; + let (nl, nr, dl, dr) = + deep.oracle.materialize_folded_rows(deep.c, &deep.bound); + let eq_row = std::mem::take(&mut deep.eq_row); + tables.nl_table = nl; + tables.nr_table = nr; + tables.dl_table = dl; + tables.dr_table = dr; + tables.eq_table = eq_row; + tables.deep = None; + } + return; + } + + if tables.dl_table.len() / 2 == 0 { + return; + } let half = tables.dl_table.len() / 2; // Update eq_correction using instance-specific eval point @@ -2127,7 +2270,26 @@ pub(crate) fn instance_eval_point( } /// Per-instance bookkeeping tables for the batch sumcheck inner loop. -struct PerInstanceTables { +/// In-progress k-bit-round state of a streamed deep layer: nothing is +/// materialized; gate sums and folds are answered by the oracle from `bound`. +struct DeepPhase<'a, E: IsField> { + oracle: &'a dyn DeepLayerOracle, + /// Extended-tree child index of this layer. + c: usize, + /// Number of k-bit rounds (`log2(per-row slots) − 1`). + kp: usize, + /// This layer's bound challenges so far (k-bit rounds only). + bound: Vec>, + /// eq over the parent's k-bits, pre-halved per k-round. + eq_k: Vec>, + /// eq over the row bits — untouched during k-rounds, becomes the layer's + /// eq_table at materialization. + eq_row: Vec>, +} + +struct PerInstanceTables<'a, E: IsField> { + /// `Some` while a streamed deep layer is in its k-bit rounds. + deep: Option>, nl_table: Vec>, nr_table: Vec>, dl_table: Vec>, diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index eafc3f0b2..2f58aec60 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -344,6 +344,77 @@ struct LogUpDeepOracle<'a, F: IsField, E: IsField> { trace_len: usize, } +impl LogUpDeepOracle<'_, F, E> +where + F: IsFFTField + IsSubFieldOf + Send + Sync, + E: IsField + Send + Sync, +{ + /// One row's four split arrays at child level `c`, folded by this layer's + /// bound challenges: build the K̂ leaves, fold `c` levels of balanced + /// fraction addition (the deep-tree values), split by slot parity into + /// (nl, nr, dl, dr) row arrays, then fold each `bound.len()` times with + /// the sumcheck's multilinear rule `v' = l + ch·(r − l)` — exactly what + /// `fold_table` does to materialized tables, per row. + #[allow(clippy::type_complexity)] + fn folded_row_arrays( + &self, + c: usize, + bound: &[FieldElement], + row_idx: usize, + ) -> ( + Vec>, + Vec>, + Vec>, + Vec>, + ) { + let row = self.main.get_row(row_idx); + let mut leaves = row_leaf_pairs( + self.interactions, + row, + &self.z, + &self.alpha_powers, + &self.shifts, + self.k_hat, + ); + for _ in 0..c { + let m = leaves.len() / 2; + for t in 0..m { + let (n0, d0) = leaves[2 * t].clone(); + let (n1, d1) = leaves[2 * t + 1].clone(); + leaves[t] = (&(&n0 * &d1) + &(&n1 * &d0), &d0 * &d1); + } + leaves.truncate(m); + } + + let half = leaves.len() / 2; + let mut nl = Vec::with_capacity(half); + let mut nr = Vec::with_capacity(half); + let mut dl = Vec::with_capacity(half); + let mut dr = Vec::with_capacity(half); + for t in 0..half { + let (n_even, d_even) = leaves[2 * t].clone(); + let (n_odd, d_odd) = leaves[2 * t + 1].clone(); + nl.push(n_even); + dl.push(d_even); + nr.push(n_odd); + dr.push(d_odd); + } + + for ch in bound { + for table in [&mut nl, &mut nr, &mut dl, &mut dr] { + let m = table.len() / 2; + for j in 0..m { + let left = table[2 * j].clone(); + let right = table[2 * j + 1].clone(); + table[j] = &left + &(ch * &(&right - &left)); + } + table.truncate(m); + } + } + (nl, nr, dl, dr) + } +} + impl DeepLayerOracle for LogUpDeepOracle<'_, F, E> where F: IsFFTField + IsSubFieldOf + Send + Sync, @@ -447,6 +518,103 @@ where (nl, nr, dl, dr) } + + fn deep_gate_sums( + &self, + c: usize, + bound: &[FieldElement], + eq_k: &[FieldElement], + eq_row: &[FieldElement], + lambda: &FieldElement, + ) -> (FieldElement, FieldElement) { + let row_sums = |row_idx: usize| -> (FieldElement, FieldElement) { + let (fnl, fnr, fdl, fdr) = self.folded_row_arrays(c, bound, row_idx); + let pairs = fnl.len() / 2; + debug_assert_eq!(eq_k.len(), pairs); + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for j in 0..pairs { + // EXACTLY gate_generic's formulas (crypto/stark/src/gkr.rs): + // g0 on the pair's left entries, g2 on the 2·r − l + // extrapolations. The stage-2 differential test pins parity. + let (nl_l, nl_r) = (&fnl[2 * j], &fnl[2 * j + 1]); + let (nr_l, nr_r) = (&fnr[2 * j], &fnr[2 * j + 1]); + let (dl_l, dl_r) = (&fdl[2 * j], &fdl[2 * j + 1]); + let (dr_l, dr_r) = (&fdr[2 * j], &fdr[2 * j + 1]); + + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(lambda * dr_l))); + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(lambda * &dr_2))); + + h0 = &h0 + &(&eq_k[j] * &gate_0); + h2 = &h2 + &(&eq_k[j] * &gate_2); + } + (&eq_row[row_idx] * &h0, &eq_row[row_idx] * &h2) + }; + + #[cfg(feature = "parallel")] + { + (0..self.trace_len).into_par_iter().map(row_sums).reduce( + || (FieldElement::::zero(), FieldElement::::zero()), + |(a0, a2), (b0, b2)| (&a0 + &b0, &a2 + &b2), + ) + } + #[cfg(not(feature = "parallel"))] + { + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for row_idx in 0..self.trace_len { + let (r0, r2) = row_sums(row_idx); + h0 = &h0 + &r0; + h2 = &h2 + &r2; + } + (h0, h2) + } + } + + #[allow(clippy::type_complexity)] + fn materialize_folded_rows( + &self, + c: usize, + bound: &[FieldElement], + ) -> ( + Vec>, + Vec>, + Vec>, + Vec>, + ) { + let n = self.trace_len; + let row_entry = |row_idx: usize| { + let (fnl, fnr, fdl, fdr) = self.folded_row_arrays(c, bound, row_idx); + debug_assert_eq!(fnl.len(), 1, "all k-bits must be bound"); + ( + fnl.into_iter().next().unwrap(), + fnr.into_iter().next().unwrap(), + fdl.into_iter().next().unwrap(), + fdr.into_iter().next().unwrap(), + ) + }; + + #[cfg(feature = "parallel")] + let entries: Vec<_> = (0..n).into_par_iter().map(row_entry).collect(); + #[cfg(not(feature = "parallel"))] + let entries: Vec<_> = (0..n).map(row_entry).collect(); + + let mut nl = Vec::with_capacity(n); + let mut nr = Vec::with_capacity(n); + let mut dl = Vec::with_capacity(n); + let mut dr = Vec::with_capacity(n); + for (a, b, c_, d) in entries { + nl.push(a); + nr.push(b); + dl.push(c_); + dr.push(d); + } + (nl, nr, dl, dr) + } } /// Build a table's batch-GKR instance (Stage 2 of the input-layer design): From 36164f6cc357ac749a65dfa852bc3818c9c66cb3 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:51:20 -0300 Subject: [PATCH 19/20] perf(gkr): cache per-instance fingerprints for the streamed deep rounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streamed k-bit rounds recomputed every leaf fingerprint on every round (21 full passes on the biggest table), pushing the batch prove to 62s on ethrex-10tx. Now each oracle caches its K*N fingerprints once (row-major, built in parallel on first deep round, freed with the instance; the padded tail is the constant 1 and never stored) — the one expensive leaf component. Numerators stay recomputed (1-2 column reads). Rounds also reuse per-thread scratch buffers (rayon for_each_init / fold) instead of allocating six vectors per row, and the c-level fraction folds are padding-aware (identity tail costs nothing). Same values, same transcript; the byte-identical differential test still pins the streamed path against the materialized extended tree. 292 stark tests green. --- crypto/stark/src/logup_gkr.rs | 273 +++++++++++++++++++++++----------- 1 file changed, 190 insertions(+), 83 deletions(-) diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index 2f58aec60..816d92557 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -342,6 +342,12 @@ struct LogUpDeepOracle<'a, F: IsField, E: IsField> { shifts: PackingShifts, k_hat: usize, trace_len: usize, + /// Row-major `[row·K + k]` cache of the K per-interaction fingerprints — + /// the one expensive leaf component (numerators are 1–2 column reads, + /// recomputed on the fly). Built once on first deep round, freed with the + /// instance; K·N ext elements (the padded K̂−K tail is the constant 1 and + /// never stored). + fp_cache: std::sync::OnceLock>>, } impl LogUpDeepOracle<'_, F, E> @@ -349,48 +355,101 @@ where F: IsFFTField + IsSubFieldOf + Send + Sync, E: IsField + Send + Sync, { + /// The fingerprint cache, built once (parallel) on first use. + fn fps(&self) -> &[FieldElement] { + self.fp_cache.get_or_init(|| { + let k = self.interactions.len(); + let fill_row = |row_idx: usize, out: &mut [FieldElement]| { + let row = self.main.get_row(row_idx); + for (slot, inter) in out.iter_mut().zip(self.interactions) { + *slot = compute_fingerprint_at_row( + inter, + row, + &self.z, + &self.alpha_powers, + &self.shifts, + ); + } + }; + let mut cache = vec![FieldElement::::zero(); k * self.trace_len]; + #[cfg(feature = "parallel")] + cache + .par_chunks_mut(k) + .enumerate() + .for_each(|(row_idx, out)| fill_row(row_idx, out)); + #[cfg(not(feature = "parallel"))] + for (row_idx, out) in cache.chunks_mut(k).enumerate() { + fill_row(row_idx, out); + } + cache + }) + } + /// One row's four split arrays at child level `c`, folded by this layer's - /// bound challenges: build the K̂ leaves, fold `c` levels of balanced - /// fraction addition (the deep-tree values), split by slot parity into - /// (nl, nr, dl, dr) row arrays, then fold each `bound.len()` times with - /// the sumcheck's multilinear rule `v' = l + ch·(r − l)` — exactly what - /// `fold_table` does to materialized tables, per row. - #[allow(clippy::type_complexity)] - fn folded_row_arrays( + /// bound challenges, written into `scratch` (reused across a thread's + /// rows): build the K leaf pairs from the fingerprint cache (numerators + /// recomputed — 1–2 column reads each), fold `c` levels of balanced + /// fraction addition padding-aware (`x + (0,1) = x`, and slots entirely + /// in the padded tail are `(0,1)` with no work), split by slot parity, + /// then fold each array `bound.len()` times with the sumcheck rule + /// `v' = l + ch·(r − l)` — exactly what `fold_table` does per row. + fn folded_row_arrays_into( &self, c: usize, bound: &[FieldElement], row_idx: usize, - ) -> ( - Vec>, - Vec>, - Vec>, - Vec>, + scratch: &mut RowScratch, ) { + let k = self.interactions.len(); + let fps = &self.fps()[row_idx * k..(row_idx + 1) * k]; let row = self.main.get_row(row_idx); - let mut leaves = row_leaf_pairs( - self.interactions, - row, - &self.z, - &self.alpha_powers, - &self.shifts, - self.k_hat, - ); + + let leaves = &mut scratch.leaves; + leaves.clear(); + for (inter, fp) in self.interactions.iter().zip(fps) { + let m = inter.multiplicity.evaluate_from(|col| row[col].clone()); + let n = if inter.is_sender { + m.to_extension() + } else { + (-m).to_extension() + }; + leaves.push((n, fp.clone())); + } + + // Fold c levels, padding-aware: the tail beyond `active` is the + // fraction identity (0, 1); an odd trailing real entry passes through + // unchanged (x + identity = x). + let mut active = k; for _ in 0..c { - let m = leaves.len() / 2; + let m = active / 2; for t in 0..m { let (n0, d0) = leaves[2 * t].clone(); let (n1, d1) = leaves[2 * t + 1].clone(); leaves[t] = (&(&n0 * &d1) + &(&n1 * &d0), &d0 * &d1); } - leaves.truncate(m); + if active % 2 == 1 { + leaves[m] = leaves[active - 1].clone(); + } + active = active.div_ceil(2); + leaves.truncate(active); } + // Pad the slot array to the full K̂/2^c width with identities. + let slots = self.k_hat >> c; + leaves.resize_with(slots, || { + (FieldElement::::zero(), FieldElement::::one()) + }); - let half = leaves.len() / 2; - let mut nl = Vec::with_capacity(half); - let mut nr = Vec::with_capacity(half); - let mut dl = Vec::with_capacity(half); - let mut dr = Vec::with_capacity(half); + let half = slots / 2; + let (nl, nr, dl, dr) = ( + &mut scratch.nl, + &mut scratch.nr, + &mut scratch.dl, + &mut scratch.dr, + ); + nl.clear(); + nr.clear(); + dl.clear(); + dr.clear(); for t in 0..half { let (n_even, d_even) = leaves[2 * t].clone(); let (n_odd, d_odd) = leaves[2 * t + 1].clone(); @@ -401,7 +460,7 @@ where } for ch in bound { - for table in [&mut nl, &mut nr, &mut dl, &mut dr] { + for table in [&mut *nl, &mut *nr, &mut *dl, &mut *dr] { let m = table.len() / 2; for j in 0..m { let left = table[2 * j].clone(); @@ -411,7 +470,27 @@ where table.truncate(m); } } - (nl, nr, dl, dr) + } +} + +/// Per-thread scratch for the streamed deep rounds (reused across rows). +struct RowScratch { + leaves: Vec<(FieldElement, FieldElement)>, + nl: Vec>, + nr: Vec>, + dl: Vec>, + dr: Vec>, +} + +impl RowScratch { + fn new(k_hat: usize) -> Self { + Self { + leaves: Vec::with_capacity(k_hat), + nl: Vec::with_capacity(k_hat / 2), + nr: Vec::with_capacity(k_hat / 2), + dl: Vec::with_capacity(k_hat / 2), + dr: Vec::with_capacity(k_hat / 2), + } } } @@ -527,47 +606,66 @@ where eq_row: &[FieldElement], lambda: &FieldElement, ) -> (FieldElement, FieldElement) { - let row_sums = |row_idx: usize| -> (FieldElement, FieldElement) { - let (fnl, fnr, fdl, fdr) = self.folded_row_arrays(c, bound, row_idx); - let pairs = fnl.len() / 2; - debug_assert_eq!(eq_k.len(), pairs); - let mut h0 = FieldElement::::zero(); - let mut h2 = FieldElement::::zero(); - for j in 0..pairs { - // EXACTLY gate_generic's formulas (crypto/stark/src/gkr.rs): - // g0 on the pair's left entries, g2 on the 2·r − l - // extrapolations. The stage-2 differential test pins parity. - let (nl_l, nl_r) = (&fnl[2 * j], &fnl[2 * j + 1]); - let (nr_l, nr_r) = (&fnr[2 * j], &fnr[2 * j + 1]); - let (dl_l, dl_r) = (&fdl[2 * j], &fdl[2 * j + 1]); - let (dr_l, dr_r) = (&fdr[2 * j], &fdr[2 * j + 1]); - - let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(lambda * dr_l))); - let nl_2 = &(nl_r + nl_r) - nl_l; - let nr_2 = &(nr_r + nr_r) - nr_l; - let dl_2 = &(dl_r + dl_r) - dl_l; - let dr_2 = &(dr_r + dr_r) - dr_l; - let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(lambda * &dr_2))); - - h0 = &h0 + &(&eq_k[j] * &gate_0); - h2 = &h2 + &(&eq_k[j] * &gate_2); - } - (&eq_row[row_idx] * &h0, &eq_row[row_idx] * &h2) - }; + let row_sums = + |scratch: &mut RowScratch, row_idx: usize| -> (FieldElement, FieldElement) { + self.folded_row_arrays_into(c, bound, row_idx, scratch); + let (fnl, fnr, fdl, fdr) = (&scratch.nl, &scratch.nr, &scratch.dl, &scratch.dr); + let pairs = fnl.len() / 2; + debug_assert_eq!(eq_k.len(), pairs); + let mut h0 = FieldElement::::zero(); + let mut h2 = FieldElement::::zero(); + for j in 0..pairs { + // EXACTLY gate_generic's formulas (crypto/stark/src/gkr.rs): + // g0 on the pair's left entries, g2 on the 2·r − l + // extrapolations. The stage-2 differential test pins parity. + let (nl_l, nl_r) = (&fnl[2 * j], &fnl[2 * j + 1]); + let (nr_l, nr_r) = (&fnr[2 * j], &fnr[2 * j + 1]); + let (dl_l, dl_r) = (&fdl[2 * j], &fdl[2 * j + 1]); + let (dr_l, dr_r) = (&fdr[2 * j], &fdr[2 * j + 1]); + + let gate_0 = &(nl_l * dr_l) + &(dl_l * &(nr_l + &(lambda * dr_l))); + let nl_2 = &(nl_r + nl_r) - nl_l; + let nr_2 = &(nr_r + nr_r) - nr_l; + let dl_2 = &(dl_r + dl_r) - dl_l; + let dr_2 = &(dr_r + dr_r) - dr_l; + let gate_2 = &(&nl_2 * &dr_2) + &(&dl_2 * &(&nr_2 + &(lambda * &dr_2))); + + h0 = &h0 + &(&eq_k[j] * &gate_0); + h2 = &h2 + &(&eq_k[j] * &gate_2); + } + (&eq_row[row_idx] * &h0, &eq_row[row_idx] * &h2) + }; #[cfg(feature = "parallel")] { - (0..self.trace_len).into_par_iter().map(row_sums).reduce( - || (FieldElement::::zero(), FieldElement::::zero()), - |(a0, a2), (b0, b2)| (&a0 + &b0, &a2 + &b2), - ) + (0..self.trace_len) + .into_par_iter() + .fold( + || { + ( + RowScratch::new(self.k_hat), + FieldElement::::zero(), + FieldElement::::zero(), + ) + }, + |(mut scratch, h0, h2), row_idx| { + let (r0, r2) = row_sums(&mut scratch, row_idx); + (scratch, &h0 + &r0, &h2 + &r2) + }, + ) + .map(|(_, h0, h2)| (h0, h2)) + .reduce( + || (FieldElement::::zero(), FieldElement::::zero()), + |(a0, a2), (b0, b2)| (&a0 + &b0, &a2 + &b2), + ) } #[cfg(not(feature = "parallel"))] { + let mut scratch = RowScratch::new(self.k_hat); let mut h0 = FieldElement::::zero(); let mut h2 = FieldElement::::zero(); for row_idx in 0..self.trace_len { - let (r0, r2) = row_sums(row_idx); + let (r0, r2) = row_sums(&mut scratch, row_idx); h0 = &h0 + &r0; h2 = &h2 + &r2; } @@ -587,31 +685,39 @@ where Vec>, ) { let n = self.trace_len; - let row_entry = |row_idx: usize| { - let (fnl, fnr, fdl, fdr) = self.folded_row_arrays(c, bound, row_idx); - debug_assert_eq!(fnl.len(), 1, "all k-bits must be bound"); - ( - fnl.into_iter().next().unwrap(), - fnr.into_iter().next().unwrap(), - fdl.into_iter().next().unwrap(), - fdr.into_iter().next().unwrap(), - ) - }; + let mut nl = vec![FieldElement::::zero(); n]; + let mut nr = vec![FieldElement::::zero(); n]; + let mut dl = vec![FieldElement::::zero(); n]; + let mut dr = vec![FieldElement::::zero(); n]; #[cfg(feature = "parallel")] - let entries: Vec<_> = (0..n).into_par_iter().map(row_entry).collect(); + nl.par_iter_mut() + .zip(nr.par_iter_mut()) + .zip(dl.par_iter_mut()) + .zip(dr.par_iter_mut()) + .enumerate() + .for_each_init( + || RowScratch::new(self.k_hat), + |scratch, (row_idx, (((nl_e, nr_e), dl_e), dr_e))| { + self.folded_row_arrays_into(c, bound, row_idx, scratch); + debug_assert_eq!(scratch.nl.len(), 1, "all k-bits must be bound"); + *nl_e = scratch.nl[0].clone(); + *nr_e = scratch.nr[0].clone(); + *dl_e = scratch.dl[0].clone(); + *dr_e = scratch.dr[0].clone(); + }, + ); #[cfg(not(feature = "parallel"))] - let entries: Vec<_> = (0..n).map(row_entry).collect(); - - let mut nl = Vec::with_capacity(n); - let mut nr = Vec::with_capacity(n); - let mut dl = Vec::with_capacity(n); - let mut dr = Vec::with_capacity(n); - for (a, b, c_, d) in entries { - nl.push(a); - nr.push(b); - dl.push(c_); - dr.push(d); + { + let mut scratch = RowScratch::new(self.k_hat); + for row_idx in 0..n { + self.folded_row_arrays_into(c, bound, row_idx, &mut scratch); + debug_assert_eq!(scratch.nl.len(), 1, "all k-bits must be bound"); + nl[row_idx] = scratch.nl[0].clone(); + nr[row_idx] = scratch.nr[0].clone(); + dl[row_idx] = scratch.dl[0].clone(); + dr[row_idx] = scratch.dr[0].clone(); + } } (nl, nr, dl, dr) } @@ -654,6 +760,7 @@ where shifts: PackingShifts::::new(), k_hat: interactions.len().next_power_of_two(), trace_len, + fp_cache: std::sync::OnceLock::new(), })) } else { None From d4d4d3bc1a4c592b66d51abb97faf97a564245c0 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 21 Jul 2026 17:57:07 -0300 Subject: [PATCH 20/20] perf(gkr): stream/materialize threshold for deep layers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming every deep layer rebuilds a per-row sub-tree of ~K-hat fraction-adds on EVERY k-round — dominant for the high-c layers whose tables are small anyway. Split by slot count: layers with > 8 per-row slots stream (shallow per-row rebuild, big tables that must not materialize); layers with <= 8 slots materialize whole (bounded transient <= 8N pairs, decaying per round) and run the ordinary fold path. The differential test gains a K = 20 (K-hat = 32) instance so the streamed branch stays pinned byte-identical alongside the materialized ones. --- crypto/stark/src/gkr.rs | 12 +++++++++--- crypto/stark/src/logup_gkr.rs | 25 ++++++++++++++++++++++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/crypto/stark/src/gkr.rs b/crypto/stark/src/gkr.rs index a07826cff..2ddc318dd 100644 --- a/crypto/stark/src/gkr.rs +++ b/crypto/stark/src/gkr.rs @@ -1651,9 +1651,15 @@ pub fn gkr_prove_batch( my_parent_num_vars, kp + oracle.num_rows().trailing_zeros() as usize, ); - if kp == 0 { - // No k-bit rounds: the split tables are already - // N-sized — materialize now, ordinary path after. + // Small deep layers (child ≤ 8·N) are materialized whole + // and run the ordinary fold path — their tables are a + // bounded transient (≤ 8N pairs, decaying per round) while + // streaming them would rebuild a per-row sub-tree of + // ~K̂ fraction-adds EVERY round (the dominant cost of the + // all-streamed variant). The big layers (slots > 8) stream: + // their per-row rebuild is shallow (≤ 2 fold levels) and + // their tables are the ones that don't fit. + if kp == 0 || slots <= 8 { let (nl_table, nr_table, dl_table, dr_table) = oracle.materialize_split(tree_layer_idx); return build_tables_from_split( diff --git a/crypto/stark/src/logup_gkr.rs b/crypto/stark/src/logup_gkr.rs index 816d92557..aa2b83919 100644 --- a/crypto/stark/src/logup_gkr.rs +++ b/crypto/stark/src/logup_gkr.rs @@ -1398,15 +1398,37 @@ mod tests { let table_a = mk_table(16, 7); let table_b = mk_table(8, 1_000); + // A wide instance (K = 20 → K̂ = 32) exercises the STREAMED deep + // rounds (slots 32 and 16 stream; slots ≤ 8 materialize). + let wide_interactions: Vec = (0..20) + .map(|k| { + if k % 2 == 0 { + BusInteraction::sender( + k as u64, + Multiplicity::One, + Packing::Direct.columns(&[k % 3]), + ) + } else { + BusInteraction::receiver( + k as u64, + Multiplicity::Column(2), + Packing::Direct.columns(&[(k + 1) % 3]), + ) + } + }) + .collect(); + let table_c = mk_table(16, 55); // Materialized extended trees (the Stage-1 shape). let mat_a = compute_logup_layers::(&interactions, &table_a, 16, &challenges); let mat_b = compute_logup_layers::(&interactions, &table_b, 8, &challenges); + let mat_c = compute_logup_layers::(&wide_interactions, &table_c, 16, &challenges); let mut t1 = DefaultTranscript::::new(&[42]); let (proof_m, point_m, claims_m) = gkr_prove_batch( vec![ GkrInstance::materialized(mat_a), GkrInstance::materialized(mat_b), + GkrInstance::materialized(mat_c), ], &mut t1, ); @@ -1414,8 +1436,9 @@ mod tests { // Streamed instances (the Stage-2 shape). let inst_a = build_gkr_instance::(&interactions, &table_a, 16, &challenges); let inst_b = build_gkr_instance::(&interactions, &table_b, 8, &challenges); + let inst_c = build_gkr_instance::(&wide_interactions, &table_c, 16, &challenges); let mut t2 = DefaultTranscript::::new(&[42]); - let (proof_s, point_s, claims_s) = gkr_prove_batch(vec![inst_a, inst_b], &mut t2); + let (proof_s, point_s, claims_s) = gkr_prove_batch(vec![inst_a, inst_b, inst_c], &mut t2); assert_eq!(point_m, point_s, "shared random points diverge"); assert_eq!(claims_m, claims_s, "instance claims diverge");