diff --git a/crypto/stark/src/fri/batched.rs b/crypto/stark/src/fri/batched.rs new file mode 100644 index 000000000..7fdb4385b --- /dev/null +++ b/crypto/stark/src/fri/batched.rs @@ -0,0 +1,499 @@ +use crypto::fiat_shamir::is_transcript::{IsStarkTranscript, IsTranscript}; +use math::field::element::FieldElement; +use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; +use math::traits::AsBytes; + +use crate::config::{FriLayerMerkleTree, FriLayerMerkleTreeBackend}; +use crate::fri::fri_commitment::FriLayer; +use crate::fri::fri_functions::{ + compute_coset_twiddles_inv, fold_evaluations_in_place, update_twiddles_in_place, +}; + +/// Combine DEEP polynomial codewords by their FRI height for batched FRI. +/// +/// Each element of `inputs` is a pair `(codeword, height)` where `height` is +/// the log₂ of the codeword length (i.e. `codeword.len() == 2^height`). +/// The global index `i` into `inputs` is used to derive the mixing power +/// `alpha^i` (index 0 → alpha^0 = 1, index 1 → alpha^1, …). +/// +/// Returns a `Vec` of length `max_height + 1`. Index `h` contains +/// `Some(combined)` where `combined[j] = Σ_{i : height_i == h} alpha^i * codeword_i[j]`, +/// or `None` when no input has height `h`. +pub fn combine_by_height( + inputs: &[(Vec>, usize)], + alpha: &FieldElement, +) -> Vec>>> +where + E: IsField, + FieldElement: Clone, +{ + if inputs.is_empty() { + return vec![]; + } + + let max_height = inputs + .iter() + .map(|(_, h)| *h) + .max() + .expect("inputs is non-empty so max height exists"); + + let mut out: Vec>>> = vec![None; max_height + 1]; + + // Precompute alpha^0, alpha^1, …, alpha^(n-1) via repeated multiplication. + let mut alpha_pows: Vec> = Vec::with_capacity(inputs.len()); + let mut cur = FieldElement::one(); + for _ in 0..inputs.len() { + alpha_pows.push(cur.clone()); + cur = &cur * alpha; + } + + for (i, (codeword, height)) in inputs.iter().enumerate() { + let h = *height; + let expected_len = 1usize << h; + assert_eq!( + codeword.len(), + expected_len, + "codeword at index {i} has length {} but height {h} expects {expected_len}", + codeword.len() + ); + + let a_i = &alpha_pows[i]; + + match &mut out[h] { + None => { + let combined: Vec> = codeword.iter().map(|x| a_i * x).collect(); + out[h] = Some(combined); + } + Some(acc) => { + for (j, x) in codeword.iter().enumerate() { + acc[j] = &acc[j] + &(a_i * x); + } + } + } + } + + out +} + +/// FRI commit phase that operates on the bucketed output of [`combine_by_height`]. +/// +/// `combined[h]` is `Some(codeword)` when there are DEEP polynomial contributions +/// at height `h` (i.e., with codeword length `2^h`), or `None` otherwise. +/// +/// The function starts from the tallest bucket (index `h_max`) and folds +/// downward. After each fold to height `h`, any bucket at `combined[h]` is +/// *injected* into the running codeword with coefficient `β²` (where `β` is +/// the current folding challenge), before the layer is committed to the +/// transcript. Termination mirrors [`crate::fri::commit_phase_from_evaluations`]: +/// the last fold produces a single scalar (`last_value`) that is appended to +/// the transcript rather than committed as a Merkle tree layer. +pub fn batched_commit_phase( + mut combined: Vec>>>, + transcript: &mut T, + coset_offset: &FieldElement, +) -> ( + FieldElement, + Vec>>, +) +where + F: IsFFTField + IsSubFieldOf + 'static, + E: IsField + 'static, + T: IsStarkTranscript + Clone, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, +{ + // Find h_max: the largest index that has a Some(codeword). + let h_max = combined + .iter() + .enumerate() + .rev() + .find_map(|(h, slot)| if slot.is_some() { Some(h) } else { None }) + .expect("batched_commit_phase: combined must have at least one Some entry"); + + // Take the starting codeword — NOT committed; it plays the role of layer 0. + let mut running = combined[h_max] + .take() + .expect("combined[h_max] is Some by construction"); + + let domain_size = 1usize << h_max; + debug_assert_eq!( + running.len(), + domain_size, + "starting codeword length must equal 2^h_max" + ); + + // Inverse twiddle factors for the initial domain size. + let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size); + + // Commit (h_max − 1) layers; one final fold yields the last_value scalar. + let num_committed_layers = h_max.saturating_sub(1); + let mut fri_layer_list = Vec::with_capacity(num_committed_layers); + + for _ in 0..num_committed_layers { + // <<<< Receive challenge β + let beta = transcript.sample_field_element(); + + // Fold evaluations in-place; running halves in length. + fold_evaluations_in_place(&mut running, &beta, &inv_twiddles); + + // Height h of the folded codeword. + let h = running.len().trailing_zeros() as usize; + + // Inject oracle at height h (if any): running[j] += β² · ro[j] + if let Some(ro) = combined[h].take() { + let beta_sq = beta.square(); + for (j, val) in ro.iter().enumerate() { + running[j] = &running[j] + &(&beta_sq * val); + } + } + + // Build the row-pair Merkle tree over the current running codeword. + let leaves: Vec<[FieldElement; 2]> = running + .chunks_exact(2) + .map(|chunk| [chunk[0].clone(), chunk[1].clone()]) + .collect(); + let merkle_tree = FriLayerMerkleTree::build(&leaves) + .expect("FRI batched commit: Merkle tree construction must succeed"); + let root = merkle_tree.root; + fri_layer_list.push(FriLayer::new(&running, merkle_tree)); + + // >>>> Send commitment: append root to transcript. + transcript.append_bytes(&root); + + // Update twiddles for the next (halved) level. + update_twiddles_in_place(&mut inv_twiddles); + } + + // <<<< Receive the final folding challenge. + let beta = transcript.sample_field_element(); + + // Final fold: running goes from length 2 → length 1. + fold_evaluations_in_place(&mut running, &beta, &inv_twiddles); + + let last_value = running + .first() + .expect("FRI evals are non-empty after final fold") + .clone(); + + // >>>> Send value: append the last scalar to the transcript. + transcript.append_field_element(&last_value); + + (last_value, fri_layer_list) +} + +/// Canonical, order-deterministic absorption of an epoch's table-height histogram +/// into the transcript. Single source of truth for the structural binding: the +/// multiset of `lde_log_height`s across an epoch's tables fully determines the fold +/// order and injection points of the batched FRI (arity is uniformly 2 — #729 is not +/// on this branch, see the task-3 brief override). Binding this histogram is +/// therefore equivalent to binding the whole arity/injection schedule. +/// +/// Encoding (fixed-width, length-prefixed, order-preserving): +/// `u64::to_le_bytes(heights.len())` followed by `u64::to_le_bytes(h)` for each `h` +/// in `heights`, in the exact order given. Caller (prover and verifier alike) must +/// pass `heights` in the same canonical per-epoch table order — this function does +/// not sort or deduplicate. +pub fn absorb_height_histogram(transcript: &mut T, heights: &[usize]) +where + E: IsField, + T: IsTranscript, +{ + transcript.append_bytes(&(heights.len() as u64).to_le_bytes()); + for h in heights { + transcript.append_bytes(&(*h as u64).to_le_bytes()); + } +} + +/// Challenges derived from replaying the shared batched round-4 transcript sequence. +/// See [`derive_batched_fri_challenges`]. +#[derive(Debug, Clone)] +pub struct BatchedFriChallenges { + /// Sampled once after the height histogram (and, at the call site, after all + /// per-table OOD evaluations have been absorbed). + pub alpha: FieldElement, + /// One per committed layer plus one final challenge for the last fold. + /// `betas.len() == layer_roots.len() + 1`. + pub betas: Vec>, + /// Transcript state right before the grinding nonce bytes are appended. + /// All-zero when `grinding_factor == 0` or `nonce` is `None` (mirrors + /// the per-table convention in `verifier.rs`). + pub grinding_seed: [u8; 32], + /// One `sample_u64(domain_size >> 1)` draw per query. + pub iotas: Vec, +} + +/// Replays the shared batched round-4 transcript sequence (histogram, alpha, +/// per-layer beta/root, final beta/last_value, grinding, query iotas) and returns +/// the derived challenges. The single routine the prover (T4) and verifier (T5) +/// both call so they provably derive identical challenges; calls +/// `absorb_height_histogram` internally instead of duplicating the encoding. +#[allow(clippy::too_many_arguments)] +pub fn derive_batched_fri_challenges( + transcript: &mut T, + heights: &[usize], + layer_roots: &[[u8; 32]], + last_value: &FieldElement, + grinding_factor: u8, + nonce: Option, + num_queries: usize, + domain_size: usize, +) -> BatchedFriChallenges +where + E: IsField, + T: IsTranscript, +{ + absorb_height_histogram(transcript, heights); + + let alpha = transcript.sample_field_element(); + + let mut betas = Vec::with_capacity(layer_roots.len() + 1); + for root in layer_roots { + let beta = transcript.sample_field_element(); + transcript.append_bytes(root); + betas.push(beta); + } + + let final_beta = transcript.sample_field_element(); + transcript.append_field_element(last_value); + betas.push(final_beta); + + let mut grinding_seed = [0u8; 32]; + if grinding_factor > 0 + && let Some(nonce_value) = nonce + { + grinding_seed = transcript.state(); + transcript.append_bytes(&nonce_value.to_be_bytes()); + } + + let iotas = (0..num_queries) + .map(|_| transcript.sample_u64((domain_size as u64) >> 1) as usize) + .collect(); + + BatchedFriChallenges { + alpha, + betas, + grinding_seed, + iotas, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fri::fri_functions::{compute_coset_twiddles_inv, fold_evaluations_in_place}; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use crypto::fiat_shamir::is_transcript::IsTranscript; + use math::field::element::FieldElement; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + type Transcript = DefaultTranscript; + + #[test] + fn combine_by_height_two_height3_one_height2() { + // Three codewords: indices 0, 1 have height 3 (length 8); + // index 2 has height 2 (length 4). + let cw0: Vec = (1u64..=8).map(FE::from).collect(); + let cw1: Vec = (10u64..=17).map(FE::from).collect(); + let cw2: Vec = (100u64..=103).map(FE::from).collect(); + + let alpha = FE::from(7u64); + + let inputs: Vec<(Vec, usize)> = + vec![(cw0.clone(), 3), (cw1.clone(), 3), (cw2.clone(), 2)]; + + let out = combine_by_height(&inputs, &alpha); + + // Output vec length = max_height + 1 = 4 (indices 0..=3 only). + assert_eq!(out.len(), 4, "output length should be max_height+1 = 4"); + + // Heights 0 and 1 have no inputs. + assert!(out[0].is_none(), "height 0 should be None"); + assert!(out[1].is_none(), "height 1 should be None"); + + // Height 3: combined[j] = alpha^0 * cw0[j] + alpha^1 * cw1[j] + let alpha0 = FE::one(); + let alpha1 = alpha; + let expected3: Vec = cw0 + .iter() + .zip(cw1.iter()) + .map(|(a, b)| &(&alpha0 * a) + &(&alpha1 * b)) + .collect(); + + let got3 = out[3].as_ref().expect("height 3 should be Some"); + assert_eq!( + got3.len(), + 8, + "height-3 combined codeword should have length 8" + ); + assert_eq!(got3, &expected3, "height-3 combined values mismatch"); + + // Height 2: combined[j] = alpha^2 * cw2[j] + let alpha2 = &alpha * α + let expected2: Vec = cw2.iter().map(|x| &alpha2 * x).collect(); + + let got2 = out[2].as_ref().expect("height 2 should be Some"); + assert_eq!( + got2.len(), + 4, + "height-2 combined codeword should have length 4" + ); + assert_eq!(got2, &expected2, "height-2 combined values mismatch"); + } + + /// Verify that after the first fold in `batched_commit_phase`: + /// - The committed layer[0] evaluation equals fold(combined[4], β₀) + β₀² · combined[3] + /// - The total number of committed layers is h_max − 1 = 3 + #[test] + fn batched_commit_phase_first_layer_matches_manual_fold_and_inject() { + // Build synthetic codewords for h=4 (len 16) and h=3 (len 8). + let data_h4: Vec = (1u64..=16).map(FE::from).collect(); + let data_h3: Vec = (101u64..=108).map(FE::from).collect(); + + // combined = [None, None, None, Some(data_h3), Some(data_h4)] + let combined: Vec>> = vec![ + None, + None, + None, + Some(data_h3.clone()), + Some(data_h4.clone()), + ]; + + let coset_offset = FE::from(3u64); + + // Create transcript; clone before mutating so we can replay independently. + let mut transcript = Transcript::new(b"batched_fri_test"); + let mut transcript_check = transcript.clone(); + + // Run the commit phase. + let (_last_val, layers) = batched_commit_phase(combined, &mut transcript, &coset_offset); + + // h_max = 4, so we expect h_max − 1 = 3 committed layers. + assert_eq!( + layers.len(), + 3, + "expected 3 committed FRI layers for h_max=4" + ); + + // --- Independent recomputation of layer[0] --- + + // The first beta is the first thing sampled from the transcript. + let beta_0 = transcript_check.sample_field_element(); + + // Fold data_h4 with beta_0 using the same twiddles as the function. + let inv_twiddles_h4 = compute_coset_twiddles_inv::(&coset_offset, 16); + let mut expected = data_h4.clone(); + fold_evaluations_in_place(&mut expected, &beta_0, &inv_twiddles_h4); + // expected now has length 8 (height 3) + + // Inject combined[3]: expected[j] += beta_0² · data_h3[j] + let beta_0_sq = beta_0.square(); + for (j, val) in data_h3.iter().enumerate() { + expected[j] = &expected[j] + &(&beta_0_sq * val); + } + + // The first committed layer's evaluation vector must match. + assert_eq!( + layers[0].evaluation, expected, + "layer[0] evaluation does not match manual fold+inject" + ); + } + + /// The prover, by hand, runs exactly the sequence documented in the design spec's + /// "Transcript binding" section (restricted to round 4, batched-arity-2, no + /// final_poly / no log_arities — see task-3 override): absorb the height + /// histogram, sample α, then per committed layer sample β and append its root, + /// then a final β and append `last_value`, then grinding, then query indices. + /// `derive_batched_fri_challenges` must reproduce byte-identical outputs when + /// fed the same inputs, starting from a transcript in the same state. + #[test] + fn batched_round4_prover_inline_matches_verifier_replay() { + // Canonical per-epoch table heights (the multiset bound into the transcript). + let heights: Vec = vec![10, 10, 8, 8, 8, 5]; + + // Fake committed layer roots (K = 4) — only their bytes/order matter here. + let layer_roots: Vec<[u8; 32]> = (0u8..4).map(|i| [i; 32]).collect(); + let last_value = FE::from(999u64); + + let grinding_factor: u8 = 4; + let num_queries = 3; + let domain_size = 1usize << 10; + + let seed_transcript = Transcript::new(b"batched_round4_test"); + let mut transcript_a = seed_transcript.clone(); + let mut transcript_b = seed_transcript.clone(); + + // --- Clone A: prover-inline sequence, by hand --- + absorb_height_histogram(&mut transcript_a, &heights); + let alpha_a = transcript_a.sample_field_element(); + + let mut betas_a = Vec::with_capacity(layer_roots.len() + 1); + for root in &layer_roots { + let beta = transcript_a.sample_field_element(); + transcript_a.append_bytes(root); + betas_a.push(beta); + } + let final_beta_a = transcript_a.sample_field_element(); + transcript_a.append_field_element(&last_value); + betas_a.push(final_beta_a); + assert_eq!( + betas_a.len(), + layer_roots.len() + 1, + "beta count must be K+1 to match batched_commit_phase" + ); + + let grinding_seed_a = transcript_a.state(); + // Test-only: derive a real PoW nonce so the grinding step is exercised + // identically by both sides (the nonce search itself is not under test). + let nonce = crate::grinding::generate_nonce(&grinding_seed_a, grinding_factor) + .expect("a valid grinding nonce exists for this small grinding_factor"); + transcript_a.append_bytes(&nonce.to_be_bytes()); + + let iotas_a: Vec = (0..num_queries) + .map(|_| transcript_a.sample_u64((domain_size as u64) >> 1) as usize) + .collect(); + + // --- Clone B: shared replay routine --- + let result = derive_batched_fri_challenges( + &mut transcript_b, + &heights, + &layer_roots, + &last_value, + grinding_factor, + Some(nonce), + num_queries, + domain_size, + ); + + assert_eq!(result.alpha, alpha_a, "alpha mismatch"); + assert_eq!(result.betas, betas_a, "beta vector mismatch"); + assert_eq!( + result.grinding_seed, grinding_seed_a, + "grinding seed mismatch" + ); + assert_eq!(result.iotas, iotas_a, "iotas mismatch"); + } + + /// Tampering with the height histogram (without changing anything else) must + /// change the derived batching challenge α — this is the structural binding + /// that protects the fold/injection schedule (spec trap #4). + #[test] + fn absorb_height_histogram_binds_heights_into_alpha() { + let heights_a: Vec = vec![10, 10, 8, 8, 8, 5]; + let heights_b: Vec = vec![10, 10, 8, 8, 8, 6]; // one height differs + + let mut transcript_a = Transcript::new(b"histogram_binding_test"); + let mut transcript_b = Transcript::new(b"histogram_binding_test"); + + absorb_height_histogram(&mut transcript_a, &heights_a); + absorb_height_histogram(&mut transcript_b, &heights_b); + + let alpha_a = transcript_a.sample_field_element(); + let alpha_b = transcript_b.sample_field_element(); + + assert_ne!( + alpha_a, alpha_b, + "different height histograms must yield different alpha" + ); + } +} diff --git a/crypto/stark/src/fri/mmcs.rs b/crypto/stark/src/fri/mmcs.rs new file mode 100644 index 000000000..f4be66880 --- /dev/null +++ b/crypto/stark/src/fri/mmcs.rs @@ -0,0 +1,722 @@ +//! Mixed-height, row-pair MMCS (Merkle Mixed Commitment Scheme). +//! +//! Commits ALL of an epoch's matrices (one per table, of possibly different +//! heights) into ONE mixed-height Merkle tree, so a single query opens ONE +//! authentication path that covers every table's row at that query — the +//! proof-size / opening-path win of the unified-shard design (SP1 / OpenVM / +//! Plonky3). Mirrors Plonky3's `MerkleTreeMmcs`, adapted to our Keccak backends +//! and to the row-pair `(x, -x)` leaf layout (#735). +//! +//! NOTE: this is a STANDALONE primitive (Task 1). It is not yet wired into the +//! prover/verifier — Tasks 2/7 consume the leaf/injection layout documented here +//! as the single source of truth. +//! +//! # Inputs +//! +//! `commit` takes matrices as `(row_major_bit_reversed_lde, log_height, width)`: +//! - `row_major_bit_reversed_lde`: the matrix's LDE evaluations, already +//! bit-reversed and laid out ROW-MAJOR. Row `j` is the `width` contiguous +//! elements `data[j*width .. (j+1)*width]`; row `j` corresponds to bit-reversed +//! LDE position `j` (same layout the per-table trace commit produces internally). +//! - `log_height`: `log2` of the number of rows; the matrix has `2^log_height` +//! rows and `data.len() == width << log_height`. +//! - `width`: number of columns. +//! +//! # Row-pair leaves +//! +//! Leaf `k` of a matrix groups LDE positions `2k` and `2k+1` (the FRI fold pair +//! `x` and `-x`), all `width` columns batched. A matrix of `log_height h` has +//! `2^(h-1)` leaves. In `open_batch`/`PolynomialOpenings`: `evaluations` = row +//! `2k`, `evaluations_sym` = row `2k+1`. +//! +//! # Tree layout (the soundness-relevant contract — verbatim for Task 7) +//! +//! Let `h_max = max(log_height)`. The base digest layer (layer 0) has +//! `N0 = 2^(h_max-1)` nodes. Layer `i` has `N0 >> i` nodes; the root is the sole +//! node of layer `h_max-1`. A matrix of `log_height h` is *injected* at layer +//! index `i = h_max - h` (so the tallest matrices, `h == h_max`, populate the +//! base layer; shorter matrices enter where the layer width matches their leaf +//! count `2^(h-1)`). +//! +//! Hashing (`H = BatchedMerkleTreeBackend::hash_data` over a `Vec` of field +//! elements; `C = BatchedMerkleTreeBackend::hash_new_parent`, the 2-input Keccak +//! compression — identical to the existing per-table tree): +//! +//! - **Base layer** node `k` (`k in [0, N0)`): +//! `layer0[k] = H( CONCAT_{m : h_m == h_max} (row_m(2k) || row_m(2k+1)) )` +//! where matrices of height `h_max` are concatenated in INPUT order. +//! - **Climb** from layer `i` to layer `i+1` (`j in [0, N_{i+1})`): +//! `parent = C(layer_i[2j], layer_i[2j+1])`. Let `inject_h = h_max - 1 - i`. If +//! any matrix has `h_m == inject_h`, then +//! `layer_{i+1}[j] = C( parent, H( CONCAT_{m : h_m == inject_h} (row_m(2j) || row_m(2j+1)) ) )` +//! (injecting matrices concatenated in INPUT order); otherwise +//! `layer_{i+1}[j] = parent`. +//! - `root = layer_{h_max-1}[0]`. +//! +//! # Query opening +//! +//! For query `iota in [0, N0)`, matrix `m` is opened at leaf +//! `k_m = iota >> (h_max - h_m)` (`= iota >> i_m`). The shared authentication +//! path holds, for each level `level in [0, h_max-1)`, the sibling +//! `layer_level[(iota >> level) ^ 1]`. ONE path authenticates all matrices. +//! The per-matrix `PolynomialOpenings.proof` fields are empty; the single +//! `MixedOpening.proof` is the authenticator. +//! +//! # Width binding (soundness) +//! +//! `verify_batch` takes per-matrix `widths` alongside `heights`. Within a height +//! group the leaf hash is over the FLAT concatenation of every matrix's opened +//! row pair (`A.eval ‖ A.eval_sym ‖ B.eval ‖ B.eval_sym ‖ …`), which does NOT by +//! itself record where each matrix's columns end. Fixing `widths[m]` (matrix +//! `m`'s column count) makes those boundaries unambiguous: without it a prover +//! could shift a boundary — e.g. lengthen one matrix's `evaluations` by one +//! element and shorten its `evaluations_sym` by one — leaving the flat bytes (and +//! therefore the group hash) identical while feeding a corrupted row downstream. +//! Consumers (the Task 7 verifier) MUST pass the committed public per-table +//! column counts, in the same INPUT order as `heights`. +//! +//! NOTE: `widths` and `heights` should ALSO be bound into the Fiat-Shamir +//! transcript by the consumer. Scope A's `absorb_height_histogram` currently +//! binds heights only; extending it to `(height, width)` pairs is a Task 4 / +//! verifier concern (out of scope for this primitive) and is flagged here. +//! +//! # Determinism +//! +//! The tree is a pure function of `(matrices, input order)`. Grouping within a +//! height (base batching and injection) follows INPUT order; the prover and +//! verifier MUST pass matrices and `heights` in the same per-epoch order. The +//! single-matrix case is byte-identical to the existing per-table row-pair tree. + +use crypto::merkle_tree::proof::Proof; +use crypto::merkle_tree::traits::IsMerkleTreeBackend; +use math::field::element::FieldElement; +use math::field::traits::IsField; +use math::traits::AsBytes; + +use crate::config::{BatchedMerkleTreeBackend, Commitment}; +use crate::proof::stark::PolynomialOpenings; + +/// One committed matrix, kept so `open_batch` can serve its row pairs. +struct MatrixEntry { + /// Row-major, bit-reversed LDE evaluations. `data[r*width .. (r+1)*width]` + /// is row `r`. + data: Vec>, + log_height: usize, + width: usize, +} + +impl MatrixEntry { + #[inline] + fn row(&self, r: usize) -> &[FieldElement] { + &self.data[r * self.width..(r + 1) * self.width] + } +} + +/// A committed mixed-height, row-pair MMCS. Owns the matrices (to serve openings) +/// and every digest layer (to serve the shared authentication path). +pub struct MixedMmcs { + root: Commitment, + /// `layers[0]` is the base digest layer; `layers[h_max-1] == [root]`. + layers: Vec>, + matrices: Vec>, + h_max: usize, +} + +/// The opening of ALL matrices at one query index, authenticated by a single +/// shared Merkle path. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct MixedOpening { + /// The one authentication path covering every matrix's row at the query. + pub proof: Proof, + /// Per-matrix row pair (in the same INPUT order as `commit`). Each entry's + /// own `proof` is empty — `MixedOpening::proof` is the authenticator. + pub per_matrix: Vec>, +} + +/// Hash the row pair `(row(2*leaf), row(2*leaf+1))` of every matrix in `group` +/// (in the given order), all columns batched, into one digest. +fn hash_group_leaf(group: &[&MatrixEntry], leaf: usize) -> Commitment +where + E: IsField, + FieldElement: AsBytes + Sync + Send, +{ + let mut buf: Vec> = Vec::new(); + for m in group { + buf.extend_from_slice(m.row(2 * leaf)); + buf.extend_from_slice(m.row(2 * leaf + 1)); + } + as IsMerkleTreeBackend>::hash_data(&buf) +} + +/// Verifier-side analogue of [`hash_group_leaf`]: hash the opened row pairs of a +/// group of openings (in the given order) into one digest. +fn hash_group_openings(group: &[&PolynomialOpenings]) -> Commitment +where + E: IsField, + FieldElement: AsBytes + Sync + Send, +{ + let mut buf: Vec> = Vec::new(); + for o in group { + buf.extend_from_slice(&o.evaluations); + buf.extend_from_slice(&o.evaluations_sym); + } + as IsMerkleTreeBackend>::hash_data(&buf) +} + +#[inline] +fn compress(left: &Commitment, right: &Commitment) -> Commitment +where + E: IsField, + FieldElement: AsBytes + Sync + Send, +{ + as IsMerkleTreeBackend>::hash_new_parent(left, right) +} + +impl MixedMmcs +where + E: IsField, + FieldElement: AsBytes + Sync + Send, +{ + /// Commit the given matrices into one mixed-height row-pair tree. See the + /// module docs for the exact leaf/injection layout. + pub fn commit(matrices: &[(Vec>, usize, usize)]) -> Self { + assert!( + !matrices.is_empty(), + "MixedMmcs::commit requires at least one matrix" + ); + + let entries: Vec> = matrices + .iter() + .map(|(data, log_height, width)| { + assert!( + *log_height >= 1, + "log_height must be >= 1 (row-pair leaves need at least 2 rows)" + ); + assert_eq!( + data.len(), + width << log_height, + "matrix data length must equal width * 2^log_height" + ); + MatrixEntry { + data: data.clone(), + log_height: *log_height, + width: *width, + } + }) + .collect(); + + let h_max = entries + .iter() + .map(|m| m.log_height) + .max() + .expect("entries is non-empty"); + let n0 = 1usize << (h_max - 1); + + // Base digest layer: batch all tallest matrices' row pairs (input order). + let base_group: Vec<&MatrixEntry> = + entries.iter().filter(|m| m.log_height == h_max).collect(); + + let mut layers: Vec> = Vec::with_capacity(h_max); + let base: Vec = (0..n0).map(|k| hash_group_leaf(&base_group, k)).collect(); + layers.push(base); + + // Climb, compressing pairs and injecting shorter matrices where the layer + // width matches their leaf count. + let mut i = 0usize; + while layers[i].len() > 1 { + let next_len = layers[i].len() / 2; + let inject_h = h_max - 1 - i; + let inject_group: Vec<&MatrixEntry> = entries + .iter() + .filter(|m| m.log_height == inject_h) + .collect(); + + let mut next: Vec = Vec::with_capacity(next_len); + for j in 0..next_len { + let mut parent = compress::(&layers[i][2 * j], &layers[i][2 * j + 1]); + if !inject_group.is_empty() { + let inj = hash_group_leaf(&inject_group, j); + parent = compress::(&parent, &inj); + } + next.push(parent); + } + layers.push(next); + i += 1; + } + + let root = layers.last().expect("at least the base layer exists")[0]; + + MixedMmcs { + root, + layers, + matrices: entries, + h_max, + } + } + + /// The committed root. + pub fn root(&self) -> Commitment { + self.root + } + + /// Open all matrices at query `iota in [0, 2^(h_max-1))`, returning each + /// matrix's row pair plus one shared authentication path. + pub fn open_batch(&self, iota: usize) -> MixedOpening { + let n0 = 1usize << (self.h_max - 1); + assert!(iota < n0, "iota {iota} out of range (n0 = {n0})"); + + let per_matrix: Vec> = self + .matrices + .iter() + .map(|m| { + let k = iota >> (self.h_max - m.log_height); + PolynomialOpenings { + proof: Proof { + merkle_path: Vec::new(), + }, + evaluations: m.row(2 * k).to_vec(), + evaluations_sym: m.row(2 * k + 1).to_vec(), + } + }) + .collect(); + + let mut merkle_path = Vec::with_capacity(self.h_max - 1); + for level in 0..(self.h_max - 1) { + let sibling = (iota >> level) ^ 1; + merkle_path.push(self.layers[level][sibling]); + } + + MixedOpening { + proof: Proof { merkle_path }, + per_matrix, + } + } + + /// Verify a batched opening at `iota` against `root`. `heights[m]` is the + /// `log_height` of matrix `m` and `widths[m]` its column count, both in the + /// SAME order as `opening.per_matrix`. + /// + /// `widths` binds each matrix's boundary inside the per-height-group leaf + /// hash (see the module `# Width binding` section): the group leaf hashes the + /// FLAT concatenation of every matrix's `evaluations ‖ evaluations_sym`, so + /// without fixed widths a prover could shift a matrix boundary while keeping + /// the flat bytes — and thus the hash — identical. Pinning `widths` makes the + /// boundaries unambiguous and closes that forgery. + pub fn verify_batch( + root: &Commitment, + iota: usize, + opening: &MixedOpening, + heights: &[usize], + widths: &[usize], + ) -> bool { + if opening.per_matrix.len() != heights.len() + || heights.len() != widths.len() + || heights.is_empty() + { + return false; + } + // Bind per-matrix boundaries: every opened matrix must present exactly + // `widths[m]` columns in BOTH rows of its pair. A boundary shift keeps the + // flat per-group concatenation identical but changes these lengths. + for (o, w) in opening.per_matrix.iter().zip(widths.iter()) { + if o.evaluations.len() != *w || o.evaluations_sym.len() != *w { + return false; + } + } + let h_max = *heights.iter().max().expect("heights is non-empty"); + // Defensive: honest heights are >= 1 (row-pair leaves need >= 2 rows), so + // h_max >= 1; guard the `h_max - 1` underflow regardless. + if h_max == 0 { + return false; + } + // Only the low `h_max - 1` bits of `iota` are consumed (one per level). + if opening.proof.merkle_path.len() != h_max - 1 { + return false; + } + + // Base node: batch all tallest matrices' opened row pairs (input order). + let base_group: Vec<&PolynomialOpenings> = opening + .per_matrix + .iter() + .zip(heights.iter()) + .filter(|(_, h)| **h == h_max) + .map(|(o, _)| o) + .collect(); + let mut acc = hash_group_openings(&base_group); + + for level in 0..(h_max - 1) { + let sibling = &opening.proof.merkle_path[level]; + let bit = (iota >> level) & 1; + let mut parent = if bit == 0 { + compress::(&acc, sibling) + } else { + compress::(sibling, &acc) + }; + + // Inject matrices whose leaf count matches this (halved) layer, in + // INPUT order — mirroring `commit`'s climb exactly. + let inject_h = h_max - 1 - level; + let inject_group: Vec<&PolynomialOpenings> = opening + .per_matrix + .iter() + .zip(heights.iter()) + .filter(|(_, h)| **h == inject_h) + .map(|(o, _)| o) + .collect(); + if !inject_group.is_empty() { + let inj = hash_group_openings(&inject_group); + parent = compress::(&parent, &inj); + } + acc = parent; + } + + &acc == root + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::commitment::commit_bit_reversed; + use math::fft::bit_reversing::reverse_index; + use math::field::element::FieldElement; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + + /// Build a row-major, bit-reversed flat vec from column-major natural-order + /// `columns`, matching the layout the existing trace commit consumes: row `j` + /// of the output = `[col_0[br(j)], ..., col_{w-1}[br(j)]]` with + /// `br = reverse_index(., num_rows)`. + fn row_major_bit_reversed(columns: &[Vec], num_rows: usize) -> Vec { + let width = columns.len(); + let mut out = vec![FE::from(0u64); num_rows * width]; + for (r, chunk) in out.chunks_exact_mut(width).enumerate() { + let br = reverse_index(r, num_rows as u64); + for (c, col) in columns.iter().enumerate() { + chunk[c] = col[br]; + } + } + out + } + + fn make_columns(width: usize, num_rows: usize, seed: u64) -> Vec> { + (0..width) + .map(|c| { + (0..num_rows) + .map(|r| { + FE::from(seed.wrapping_mul(31) + (c as u64) * 1009 + (r as u64) * 7 + 1) + }) + .collect() + }) + .collect() + } + + #[test] + fn single_matrix_commit_open_verify_and_tamper() { + let log_height = 2usize; + let num_rows = 1usize << log_height; + let width = 3usize; + let columns = make_columns(width, num_rows, 5); + let data = row_major_bit_reversed(&columns, num_rows); + + let mmcs = MixedMmcs::commit(&[(data.clone(), log_height, width)]); + let heights = [log_height]; + let widths = [width]; + let n0 = 1usize << (log_height - 1); + + for iota in 0..n0 { + let opening = mmcs.open_batch(iota); + assert_eq!(opening.per_matrix.len(), 1); + let k = iota; + let row_2k = data[(2 * k) * width..(2 * k + 1) * width].to_vec(); + let row_2k1 = data[(2 * k + 1) * width..(2 * k + 2) * width].to_vec(); + assert_eq!(opening.per_matrix[0].evaluations, row_2k); + assert_eq!(opening.per_matrix[0].evaluations_sym, row_2k1); + assert!(MixedMmcs::verify_batch( + &mmcs.root(), + iota, + &opening, + &heights, + &widths + )); + } + + let mut opening = mmcs.open_batch(0); + opening.per_matrix[0].evaluations[0] = + &opening.per_matrix[0].evaluations[0] + &FE::from(1u64); + assert!(!MixedMmcs::verify_batch( + &mmcs.root(), + 0, + &opening, + &heights, + &widths + )); + } + + #[test] + fn single_matrix_root_matches_existing_row_pair_tree() { + let log_height = 3usize; + let num_rows = 1usize << log_height; + let width = 4usize; + let columns = make_columns(width, num_rows, 9); + + let (_, existing_root) = + commit_bit_reversed(&columns, 2).expect("non-empty columns build a tree"); + + let data = row_major_bit_reversed(&columns, num_rows); + let mmcs = MixedMmcs::commit(&[(data, log_height, width)]); + + assert_eq!(mmcs.root(), existing_root); + } + + #[test] + fn mixed_height_open_positions_verify_and_tamper() { + // Three matrices, log_heights {5, 5, 3}, widths {2, 1, 4}. + let (ha, hb, hc) = (5usize, 5usize, 3usize); + let (wa, wb, wc) = (2usize, 1usize, 4usize); + let a = row_major_bit_reversed(&make_columns(wa, 1 << ha, 1), 1 << ha); + let b = row_major_bit_reversed(&make_columns(wb, 1 << hb, 2), 1 << hb); + let c = row_major_bit_reversed(&make_columns(wc, 1 << hc, 3), 1 << hc); + + let mmcs = MixedMmcs::commit(&[ + (a.clone(), ha, wa), + (b.clone(), hb, wb), + (c.clone(), hc, wc), + ]); + let heights = [ha, hb, hc]; + let widths = [wa, wb, wc]; + let h_max = 5usize; + let n0 = 1usize << (h_max - 1); // 16 + + let row = |data: &[FE], w: usize, r: usize| data[r * w..(r + 1) * w].to_vec(); + + for iota in [0usize, 1, 2, 3, 7, 8, 13, n0 - 1] { + let opening = mmcs.open_batch(iota); + assert_eq!(opening.per_matrix.len(), 3); + + // Tall matrices open at k = iota >> 0 = iota. + assert_eq!(opening.per_matrix[0].evaluations, row(&a, wa, 2 * iota)); + assert_eq!( + opening.per_matrix[0].evaluations_sym, + row(&a, wa, 2 * iota + 1) + ); + assert_eq!(opening.per_matrix[1].evaluations, row(&b, wb, 2 * iota)); + + // Height-3 matrix opens at k = iota >> (5 - 3) = iota >> 2. + let kc = iota >> (h_max - hc); + assert_eq!(opening.per_matrix[2].evaluations, row(&c, wc, 2 * kc)); + assert_eq!( + opening.per_matrix[2].evaluations_sym, + row(&c, wc, 2 * kc + 1) + ); + + assert!( + MixedMmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "honest opening at iota={iota} must verify" + ); + } + + // Tamper the height-3 matrix's opened row -> rejection (proves the short + // matrix is bound by the shared path via injection). + let iota = 6usize; + let mut opening = mmcs.open_batch(iota); + opening.per_matrix[2].evaluations[0] = + &opening.per_matrix[2].evaluations[0] + &FE::from(1u64); + assert!( + !MixedMmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "tampered height-3 row must be rejected" + ); + + // Tamper a tall-matrix row too -> rejection. + let mut opening2 = mmcs.open_batch(iota); + opening2.per_matrix[0].evaluations[0] = + &opening2.per_matrix[0].evaluations[0] + &FE::from(1u64); + assert!( + !MixedMmcs::verify_batch(&mmcs.root(), iota, &opening2, &heights, &widths), + "tampered tall-matrix row must be rejected" + ); + } + + /// Vector test: hand-compute the root for `{log_height 2, log_height 1}` + /// matrices per the documented layout and assert equality. Pins the + /// leaf/injection contract consumed verbatim by Task 7, plus determinism. + #[test] + fn vector_root_layout_contract_and_determinism() { + // A: log_height 2 (4 rows), width 2 ; B: log_height 1 (2 rows), width 3. + let a_data = row_major_bit_reversed(&make_columns(2, 4, 3), 4); + let b_data = row_major_bit_reversed(&make_columns(3, 2, 8), 2); + + let mmcs = MixedMmcs::commit(&[(a_data.clone(), 2, 2), (b_data.clone(), 1, 3)]); + + // Hand recomputation via the backend primitives, in the documented order. + let arow = |r: usize| a_data[r * 2..(r + 1) * 2].to_vec(); + let brow = |r: usize| b_data[r * 3..(r + 1) * 3].to_vec(); + let h = |v: Vec| { + as IsMerkleTreeBackend>::hash_data(&v) + }; + + // Base layer (matrix A only): leaf k = H(A.row(2k) || A.row(2k+1)). + let mut leaf0 = arow(0); + leaf0.extend(arow(1)); + let mut leaf1 = arow(2); + leaf1.extend(arow(3)); + let l00 = h(leaf0); + let l01 = h(leaf1); + + // Climb to layer 1 (root): compress the base pair, then inject B (h=1). + let parent = compress::(&l00, &l01); + let mut binj = brow(0); + binj.extend(brow(1)); + let inj = h(binj); + let expected_root = compress::(&parent, &inj); + + assert_eq!( + mmcs.root(), + expected_root, + "root must match the hand-computed mixed-height layout" + ); + + // Determinism: a second commit over the same inputs yields the same root. + let mmcs2 = MixedMmcs::commit(&[(a_data, 2, 2), (b_data, 1, 3)]); + assert_eq!(mmcs.root(), mmcs2.root(), "commit must be deterministic"); + + for iota in 0..2usize { + let opening = mmcs.open_batch(iota); + // heights {2, 1}, widths {2, 3}. + assert!(MixedMmcs::verify_batch( + &mmcs.root(), + iota, + &opening, + &[2, 1], + &[2, 3] + )); + } + } + + /// Two SAME-HEIGHT matrices share one base-group leaf, whose hash is over the + /// FLAT concatenation `A.eval ‖ A.eval_sym ‖ B.eval ‖ B.eval_sym`. A malicious + /// prover can shift the A|A_sym boundary (move one element from A's + /// `evaluations_sym` into A's `evaluations`) leaving that flat concatenation — + /// and hence the leaf hash — byte-identical, so a width-blind `verify_batch` + /// would accept it. The per-matrix width binding rejects the shift. + #[test] + fn boundary_shift_forgery_rejected() { + let h = 2usize; + let num_rows = 1usize << h; + let (wa, wb) = (2usize, 1usize); // wA >= 2 so we can steal one column. + let a = row_major_bit_reversed(&make_columns(wa, num_rows, 11), num_rows); + let b = row_major_bit_reversed(&make_columns(wb, num_rows, 22), num_rows); + + let mmcs = MixedMmcs::commit(&[(a, h, wa), (b, h, wb)]); + let heights = [h, h]; + let widths = [wa, wb]; + + let iota = 0usize; + let opening = mmcs.open_batch(iota); + assert!( + MixedMmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "honest opening must verify" + ); + + // Forge: lengthen A.evaluations by one element taken from A.evaluations_sym. + let mut forged = mmcs.open_batch(iota); + let moved = forged.per_matrix[0].evaluations_sym.remove(0); + forged.per_matrix[0].evaluations.push(moved); + + // The FLAT per-group concatenation is byte-identical to the honest one, so + // the group leaf hash is UNCHANGED — the rejection must come from the width + // check, not from a differing hash. + let flat = |o: &MixedOpening| -> Vec { + let mut v = Vec::new(); + for m in &o.per_matrix { + v.extend_from_slice(&m.evaluations); + v.extend_from_slice(&m.evaluations_sym); + } + v + }; + assert_eq!( + flat(&opening), + flat(&forged), + "the flat concatenation must be byte-identical (boundary-only shift)" + ); + + assert!( + !MixedMmcs::verify_batch(&mmcs.root(), iota, &forged, &heights, &widths), + "boundary-shift forgery must be rejected by the width binding" + ); + } + + /// Extension-field (Fp3) coverage: the aux/composition matrices Tasks 2/3 feed + /// the MMCS are cubic-extension. Byte-parity cross-check of a single Fp3 matrix + /// against the existing per-table row-pair tree, plus an open/verify/tamper + /// roundtrip over the extension path. + #[test] + fn single_matrix_fp3_root_matches_existing_row_pair_tree() { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Fp3; + type F3 = FieldElement; + + let log_height = 3usize; + let num_rows = 1usize << log_height; + let width = 3usize; + + // Populate ALL three components so the 24-byte extension serialization is + // exercised (not just the embedded-base subset). + let columns: Vec> = (0..width) + .map(|c| { + (0..num_rows) + .map(|r| { + F3::new([ + FE::from((c as u64) * 7 + r as u64 + 1), + FE::from((r as u64) * 13 + 2), + FE::from((c as u64) * 5 + (r as u64) * 3 + 4), + ]) + }) + .collect() + }) + .collect(); + + let (_, existing_root) = + commit_bit_reversed(&columns, 2).expect("non-empty columns build a tree"); + + // Row-major bit-reversed equivalent of the same column-major data. + let mut data = vec![F3::zero(); num_rows * width]; + for (r, chunk) in data.chunks_exact_mut(width).enumerate() { + let br = reverse_index(r, num_rows as u64); + for (c, col) in columns.iter().enumerate() { + chunk[c] = col[br]; + } + } + + let mmcs = MixedMmcs::commit(&[(data, log_height, width)]); + assert_eq!( + mmcs.root(), + existing_root, + "Fp3 single-matrix root must match the existing row-pair tree" + ); + + let heights = [log_height]; + let widths = [width]; + for iota in 0..(1usize << (log_height - 1)) { + let opening = mmcs.open_batch(iota); + assert!(MixedMmcs::verify_batch( + &mmcs.root(), + iota, + &opening, + &heights, + &widths + )); + } + + let mut opening = mmcs.open_batch(0); + opening.per_matrix[0].evaluations[0] = &opening.per_matrix[0].evaluations[0] + &F3::one(); + assert!(!MixedMmcs::verify_batch( + &mmcs.root(), + 0, + &opening, + &heights, + &widths + )); + } +} diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 8f1172524..0d5f08a2d 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -1,6 +1,8 @@ +pub mod batched; pub mod fri_commitment; pub mod fri_decommit; pub(crate) mod fri_functions; +pub mod mmcs; pub(crate) mod terminal; use crypto::fiat_shamir::is_transcript::IsStarkTranscript; diff --git a/crypto/stark/src/proof/stark.rs b/crypto/stark/src/proof/stark.rs index 675160837..24a314903 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, fri::mmcs::MixedOpening, + lookup::BusPublicInputs, table::Table, }; #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -83,3 +84,66 @@ pub struct StarkProof, E: IsField, PI> { pub struct MultiProof, E: IsField, PI> { pub proofs: Vec>, } + +/// Opening of all tables at ONE FRI query index, read from the per-phase +/// mixed-height MMCS trees. `main`, `aux` and `composition` each carry ONE +/// shared authentication path covering every table's row-pair at the query — +/// the unified-shard opening-path win (N×Q auth paths collapse to ~Q per phase). +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "")] +pub struct BatchedQueryOpening, E: IsField> { + pub main: MixedOpening, + pub aux: Option>, + pub composition: MixedOpening, + /// Per preprocessed table (in preprocessed-table order): the precomputed + /// columns opened against that table's hardcoded precomputed tree — those + /// columns are NOT part of the shared main MMCS. + pub precomputed: Vec>, +} + +/// Per-table data carried by a [`BatchedMultiProof`] (canonical epoch order). +/// The three commitment roots and the OOD point `z` are SHARED across the epoch +/// (roots live on `BatchedMultiProof`; `z` is re-derived by the verifier), so +/// only genuinely per-table values live here. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(bound = "PI: serde::Serialize + serde::de::DeserializeOwned")] +pub struct BatchedTableData { + pub trace_length: usize, + /// tⱼ(z gᵏ) + pub trace_ood_evaluations: Table, + /// Hᵢ(z^N) + pub composition_poly_parts_ood_evaluation: Vec>, + /// Hardcoded precomputed-columns commitment (preprocessed tables); the + /// verifier checks it against the AIR's known value. + pub precomputed_root: Option, + pub bus_public_inputs: Option>, + pub public_inputs: PI, +} + +/// A batched STARK proof for an epoch of tables sharing ONE linear transcript, +/// ONE OOD point `z`, and ONE FRI over the height-combined DEEP codewords +/// (unified-shard / Plonky3-style). Produced by `Prover::multi_prove_batched` +/// and verified by `Verifier::batched_multi_verify`. Eventually replaces +/// [`MultiProof`]. +#[derive(Debug, serde::Serialize, serde::Deserialize)] +#[serde(bound = "PI: serde::Serialize + serde::de::DeserializeOwned")] +pub struct BatchedMultiProof, E: IsField, PI> { + /// Shared mixed-height MMCS root over all tables' main-split matrices. + pub main_root: Commitment, + /// Shared mixed-height MMCS root over all aux-carrying tables' matrices. + pub aux_root: Option, + /// Shared mixed-height MMCS root over all tables' composition matrices. + pub composition_root: Commitment, + /// Merkle roots of the batched-FRI fold layers. + pub fri_layers_merkle_roots: Vec, + /// Final FRI folding value. + pub fri_last_value: FieldElement, + /// Per-query openings of the FRI fold layers (shared across tables). + pub query_list: Vec>, + /// Proof-of-work grinding nonce. + pub nonce: Option, + /// Per-query openings of the three shared trees (+ per-table precomputed). + pub deep_poly_openings: Vec>, + /// Per-table data in canonical epoch order. + pub per_table: Vec>, +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5c03292da..2d786e51b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1,3 +1,9 @@ +// The IsStarkProver trait is a crate-internal abstraction whose methods trade in +// crate-internal round types (Round1/2/3/4, TableCommit, LdeTwiddles). Exposing the +// rounds-1-3 bundle through prove_rounds_1_to_3 surfaces the private_interfaces lint +// across the whole trait; these types are never nameable/used across the crate +// boundary, so silence it module-wide rather than bump every helper type to pub. +#![allow(private_interfaces)] use std::marker::PhantomData; use std::sync::{Arc, OnceLock}; #[cfg(feature = "instruments")] @@ -27,6 +33,7 @@ use rayon::prelude::{ #[cfg(feature = "debug-checks")] use crate::debug::validate_trace; use crate::fri; +use crate::fri::mmcs::MixedMmcs; use crate::lookup::LOGUP_NUM_CHALLENGES; use crate::proof::stark::{DeepPolynomialOpenings, PolynomialOpenings}; #[cfg(feature = "disk-spill")] @@ -40,11 +47,12 @@ use super::domain::{Domain, DomainConstants}; use super::fri::fri_decommit::FriDecommitment; use super::grinding; use super::lookup::BusPublicInputs; -use super::proof::stark::{DeepPolynomialOpening, MultiProof, StarkProof}; +use super::proof::stark::{ + BatchedMultiProof, BatchedQueryOpening, BatchedTableData, DeepPolynomialOpening, MultiProof, + StarkProof, +}; use super::trace::TraceTable; use super::traits::AIR; -#[cfg(feature = "cuda")] -use crypto::merkle_tree::proof::Proof; pub use crate::commitment::{keccak_leaves_bit_reversed, keccak_leaves_row_pair_bit_reversed}; @@ -204,6 +212,10 @@ type MainCommitTuple = ( #[cfg(not(feature = "cuda"))] type MainCommitTuple = (TableCommit, (Vec>, usize)); +/// CPU aux-commit result: the committed aux `TableCommit` plus its row-major +/// extension-field LDE `(data, width)`. No GPU handle — the aux lane is CPU-only. +type AuxCommitTuple = (TableCommit, (Vec>, usize)); + /// Round 1 commitment artifacts — Merkle trees, roots, challenges, and bus inputs. /// Borrowed (not consumed) when building `Round1` in Phase D. pub(crate) struct Round1Commitments @@ -424,53 +436,6 @@ pub fn table_parallelism() -> usize { } } -/// Heuristic peak device bytes for one table: co-resident LDE columns plus the -/// resident Merkle trees, with a scratch factor for NTT and leaf transients. A -/// deliberate over estimate for a safety ceiling, not a precise allocator. Pass -/// aux_cols == 0 when the aux LDE is not yet resident (R1 main commit). -fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) -> u64 { - const BYTES_PER_BASE: u64 = 8; - const EXT3_BYTES: u64 = 24; - const SCRATCH_FACTOR: u64 = 2; - const RESIDENT_TREE_BYTES_PER_LDE: u64 = 256; - let lde = lde_size as u64; - let per_row = (main_cols as u64).saturating_mul(BYTES_PER_BASE) - + (aux_cols as u64).saturating_mul(EXT3_BYTES); - let lde_term = lde.saturating_mul(per_row).saturating_mul(SCRATCH_FACTOR); - let tree_term = lde.saturating_mul(RESIDENT_TREE_BYTES_PER_LDE); - lde_term.saturating_add(tree_term) -} - -/// Plan contiguous table chunks for parallel proving. A chunk grows until it -/// hits `k` tables or its summed VRAM estimate would exceed `budget`; a single -/// table larger than `budget` runs solo. With `budget == u64::MAX` (non-cuda, -/// or VRAM not binding) chunks fall back to fixed size `k`, identical to the -/// old `step_by(k)`, so scheduling and the proof are unchanged. Returns -/// `(start, end)` half open ranges covering `0..estimates.len()` in order. -fn plan_table_chunks(estimates: &[u64], k: usize, budget: u64) -> Vec<(usize, usize)> { - let n = estimates.len(); - let k = k.max(1); - let budget = budget as u128; - let mut chunks = Vec::new(); - let mut start = 0; - while start < n { - let mut end = start; - let mut acc: u128 = 0; - while end < n { - let next = estimates[end] as u128; - // Always admit at least one table per chunk (oversized → solo). - if end > start && (end - start >= k || acc + next > budget) { - break; - } - acc += next; - end += 1; - } - chunks.push((start, end)); - start = end; - } - chunks -} - /// A container for the results of the second round of the STARK Prove protocol. pub(crate) struct Round2 where @@ -483,12 +448,13 @@ where pub(crate) composition_poly_merkle_tree: BatchedMerkleTree, /// The commitment to the composition polynomial parts. pub(crate) composition_poly_root: Commitment, - /// The composition Merkle tree kept resident on device (when the R2 GPU tree - /// path ran), so R4 openings gather paths on device instead of walking a host - /// tree. When set, `composition_poly_merkle_tree` is a root only placeholder. - /// `None` on the CPU path. + /// Device-resident de-interleaved LDE handle from the R2 fused GPU path + /// (`try_evaluate_parts_on_lde_gpu_keep`). When present, R4 DEEP skips + /// the `num_parts * 3 * lde_size * 8` byte H2D and reads parts on + /// device. `None` when the GPU R2 path didn't run (number_of_parts <= 2, + /// below threshold, or any CPU fallback). #[cfg(feature = "cuda")] - pub(crate) gpu_composition_tree: Option, + pub(crate) gpu_composition_parts: Option, } /// A container for the results of the third round of the STARK Prove protocol. @@ -547,6 +513,28 @@ where /// helpers — only `prove`, `multi_prove` are meant for callers. The /// `private_interfaces` allow is removed once these helpers move off the trait. #[allow(private_interfaces)] +/// Owned artifacts of rounds 1-3 (the linearized unified-shard front half), +/// shared by the reference per-table path (`multi_prove`) and the batched path +/// (`multi_prove_batched`). The three MMCS are KEPT (not transient) so the +/// batched round 4 can open every table at each query from ONE tree per phase. +pub(crate) struct RoundsOneToThree +where + Field: IsSubFieldOf + IsFFTField, + FieldExtension: IsField, + FieldElement: AsBytes, + FieldElement: AsBytes, +{ + pub(crate) round1s: Vec>, + pub(crate) round2s: Vec>, + pub(crate) round3s: Vec>, + pub(crate) z: FieldElement, + pub(crate) domains: Vec>>, + pub(crate) main_mmcs: MixedMmcs, + pub(crate) aux_mmcs: Option>, + pub(crate) comp_mmcs: MixedMmcs, + pub(crate) heights: Vec, +} + pub trait IsStarkProver< Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, FieldExtension: Send + Sync + IsField + 'static, @@ -752,6 +740,49 @@ pub trait IsStarkProver< }); } + /// Build ONE mixed-height `MixedMmcs` over every table's MAIN-split LDE + /// matrix (Task 2, batched-FRI unified-shard plan). `main_commits` and + /// `main_ldes` must be Phase A's vectors, same per-epoch table order + /// `MixedMmcs::commit` requires. + /// + /// "Main split" excludes the leading precomputed-column prefix for + /// preprocessed tables (`commit.num_precomputed_cols`); those stay + /// committed separately via the existing hardcoded precomputed tree. + /// + /// `main_ldes` is row-major in NATURAL order (row r = the LDE evaluation + /// at domain point r); `MixedMmcs::commit` requires row-major data + /// already permuted into BIT-REVERSED order, so this builds a temporary + /// permuted copy per table (see the Task 2 report's memory note). The + /// resulting single-matrix root is byte-identical to the existing + /// per-table row-pair tree root over the same column range. + fn build_batched_main_mmcs( + main_commits: &[TableCommit], + main_ldes: &[(Vec>, usize)], + ) -> MixedMmcs + where + FieldElement: AsBytes + Sync + Send, + { + let mats: Vec<(Vec>, usize, usize)> = main_commits + .iter() + .zip(main_ldes.iter()) + .map(|(commit, (main_data, total_cols))| { + let col_start = commit.num_precomputed_cols; + let width = total_cols - col_start; + let num_rows = main_data.len() / total_cols; + let log_height = num_rows.trailing_zeros() as usize; + let mut data: Vec> = Vec::with_capacity(num_rows * width); + for r in 0..num_rows { + let br = reverse_index(r, num_rows as u64); + let src = br * total_cols + col_start; + data.extend_from_slice(&main_data[src..src + width]); + } + (data, log_height, width) + }) + .collect(); + + MixedMmcs::::commit(&mats) + } + /// Compute the main-trace LDE and commit. Returns a `TableCommit` along /// with the owned LDE columns (consumed later in Phase D) and (under /// cuda) the optional device LDE buffer kept alive for downstream rounds @@ -904,6 +935,44 @@ pub trait IsStarkProver< Ok((commit, (main_data, total_cols))) } + /// Commit a single table's auxiliary (LogUp) trace: row-major coset LDE over + /// the extension field, then a bit-reversed Merkle tree. CPU path (mirrors the + /// CPU branch of `commit_main_trace` / Round-1 Phase-C aux commit) — used by + /// the continuation epoch driver for the standalone L2G lane, whose aux table + /// is small and always host-resident. The verifier is agnostic to how the tree + /// was built, so a CPU-only commit is valid under every build. + #[cfg_attr(not(any(test, feature = "cuda")), allow(dead_code))] + fn commit_aux_trace( + trace: &TraceTable, + domain: &Domain, + twiddles: &LdeTwiddles, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + { + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let (trace_data, total_cols) = trace.aux_data_row_major(); + let mut aux_data: Vec> = + Vec::with_capacity(lde_size * total_cols); + aux_data.extend_from_slice(trace_data); + Polynomial::>::coset_lde_full_expand_row_major::( + &mut aux_data, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + &twiddles.two_half_inv, + &twiddles.two_half_fwd, + )?; + #[allow(unused_mut)] + let (mut tree, root) = Self::commit_rows_bit_reversed(&aux_data, total_cols) + .ok_or(ProvingError::EmptyCommitment)?; + #[cfg(feature = "disk-spill")] + Self::spill_tree(&mut tree, storage_mode, "aux Merkle tree")?; + Ok((TableCommit::plain(tree, root), (aux_data, total_cols))) + } + /// Spill a committed Merkle tree to disk when `storage_mode` is `Disk`, /// tagging any I/O error with `label`. No-op otherwise. Shared by every commit /// site (main / preprocessed split / aux). @@ -1146,7 +1215,7 @@ pub trait IsStarkProver< pub_inputs: &PI, domain: &Domain, twiddles: &LdeTwiddles, - round_1_result: &mut Round1, + round_1_result: &Round1, transition_coefficients: &[FieldElement], boundary_coefficients: &[FieldElement], ) -> Result, ProvingError> @@ -1246,55 +1315,42 @@ pub trait IsStarkProver< let t_sub = Instant::now(); // GPU fast path for the comp-poly Merkle commit: row-pair Keccak // leaves + device-side inner tree, both wrapping the host eval Vecs. - // GPU path keeps the composition tree resident on device (no whole tree - // copy) and returns a root only host tree. The device tree is threaded - // to R4 in `Round2.gpu_composition_tree`. + // `try_build_comp_poly_tree_gpu` returns `(host_tree, dev_tree)`; this batched + // comp-poly path only needs the host tree (the device tree is an unused GPU + // optimization here), so drop it to keep `gpu_tree` the same `Option` + // as the non-cuda binding and share the commit `match` below. #[cfg(feature = "cuda")] - let (composition_poly_merkle_tree, composition_poly_root, gpu_composition_tree) = - match crate::gpu_lde::try_build_comp_poly_tree_gpu::< - FieldExtension, - BatchedMerkleTreeBackend, - >(&lde_composition_poly_parts_evaluations) - { - Some((host_tree, dev_tree)) => { - let root = host_tree.root; - (host_tree, root, Some(dev_tree)) - } - None => { - let (tree, root) = crate::commitment::commit_bit_reversed( - &lde_composition_poly_parts_evaluations, - crate::commitment::ROWS_PER_LEAF, - ) - .ok_or(ProvingError::EmptyCommitment)?; - (tree, root, None) - } - }; + let gpu_tree = crate::gpu_lde::try_build_comp_poly_tree_gpu::< + FieldExtension, + BatchedMerkleTreeBackend, + >(&lde_composition_poly_parts_evaluations) + .map(|(host_tree, _dev_tree)| host_tree); #[cfg(not(feature = "cuda"))] - let (composition_poly_merkle_tree, composition_poly_root) = - crate::commitment::commit_bit_reversed( + let gpu_tree: Option> = None; + + let (composition_poly_merkle_tree, composition_poly_root) = match gpu_tree { + Some(tree) => { + let root = tree.root; + (tree, root) + } + None => crate::commitment::commit_bit_reversed( &lde_composition_poly_parts_evaluations, crate::commitment::ROWS_PER_LEAF, ) - .ok_or(ProvingError::EmptyCommitment)?; + .ok_or(ProvingError::EmptyCommitment)?, + }; #[cfg(feature = "instruments")] let merkle_dur = t_sub.elapsed(); #[cfg(feature = "instruments")] crate::instruments::store_r2_sub(constraints_dur, fft_dur, merkle_dur); - // Fold the R2 device composition parts handle into the session (resident - // R2 to R4). The host evaluations stay in `Round2` for R4 openings. - #[cfg(feature = "cuda")] - if let Some(handle) = gpu_composition_parts { - round_1_result.lde_trace.set_gpu_composition_parts(handle); - } - Ok(Round2 { lde_composition_poly_evaluations: lde_composition_poly_parts_evaluations, composition_poly_merkle_tree, composition_poly_root, #[cfg(feature = "cuda")] - gpu_composition_tree, + gpu_composition_parts, }) } @@ -1560,7 +1616,7 @@ pub trait IsStarkProver< && let Some(deep_evals) = crate::gpu_lde::try_deep_composition_gpu::( lde_trace, - lde_trace.gpu_composition_parts(), + round_2_result.gpu_composition_parts.as_ref(), &round_2_result.lde_composition_poly_evaluations, h_ood, &trace_ood_columns, @@ -1596,7 +1652,7 @@ pub trait IsStarkProver< if let Some(deep_evals) = crate::gpu_lde::try_deep_composition_gpu::( lde_trace, - lde_trace.gpu_composition_parts(), + round_2_result.gpu_composition_parts.as_ref(), &round_2_result.lde_composition_poly_evaluations, h_ood, &trace_ood_columns, @@ -1717,45 +1773,6 @@ pub trait IsStarkProver< } } - /// Like [`Self::open_composition_poly`] but uses a Merkle proof already - /// gathered from the resident device composition tree - /// ([`crate::gpu_lde::gather_proofs_dev`]) instead of walking a host tree. - /// Row-pair leaf: one proof at position `index` authenticates both rows. - #[cfg(feature = "cuda")] - fn open_composition_poly_with_proof( - proof: Proof, - lde_composition_poly_evaluations: &[Vec>], - index: usize, - ) -> PolynomialOpenings - where - FieldElement: AsBytes + Sync + Send, - FieldElement: AsBytes + Sync + Send, - { - let lde_composition_poly_parts_evaluation: Vec<_> = lde_composition_poly_evaluations - .iter() - .flat_map(|part| { - vec![ - part[reverse_index(index * 2, part.len() as u64)].clone(), - part[reverse_index(index * 2 + 1, part.len() as u64)].clone(), - ] - }) - .collect(); - - PolynomialOpenings { - proof, - evaluations: lde_composition_poly_parts_evaluation - .clone() - .into_iter() - .step_by(2) - .collect(), - evaluations_sym: lde_composition_poly_parts_evaluation - .into_iter() - .skip(1) - .step_by(2) - .collect(), - } - } - /// Computes values and validity proofs of the evaluations of trace polynomials at /// the FRI query challenge `challenge` and its symmetric counterpart. The caller /// supplies a `gather` closure that pulls the row data from the column-major LDE @@ -1784,31 +1801,6 @@ pub trait IsStarkProver< } } - /// Like [`Self::open_polys_with`], but uses a Merkle proof already gathered - /// from the resident device tree (see [`crate::gpu_lde::gather_proofs_dev`]) - /// instead of walking a host tree. Row-pair leaf: one proof at position - /// `challenge` authenticates both the queried row and its symmetric - /// counterpart. Evaluations still come from the host LDE columns via `gather`. - #[cfg(feature = "cuda")] - fn open_polys_with_proofs( - domain: &Domain, - proof: Proof, - challenge: usize, - gather: G, - ) -> PolynomialOpenings - where - C: IsField, - FieldElement: AsBytes + Sync + Send, - G: Fn(usize) -> Vec>, - { - let domain_size = domain.lde_roots_of_unity_coset.len() as u64; - PolynomialOpenings { - proof, - evaluations: gather(reverse_index(challenge * 2, domain_size)), - evaluations_sym: gather(reverse_index(challenge * 2 + 1, domain_size)), - } - } - /// Open the deep composition polynomial on a list of indexes and their symmetric elements. fn open_deep_composition_poly( domain: &Domain, @@ -1828,63 +1820,7 @@ pub trait IsStarkProver< let num_precomputed_cols = main_commit.num_precomputed_cols; let total_cols = lde_trace.num_main_cols(); - // R4 trace proofs from the resident device trees, gathered in one batch - // over all query positions instead of walking the host trees (byte - // identical to the host proofs, guarded by the `merkle_gather` test). - // `*_dev_proofs` is `Some` exactly when the tree is device resident (so - // the host tree is a root only placeholder). In that case the gather - // must succeed: there is no host tree to fall back to, so a gather error - // is a hard abort. When the tree is not device resident the value is - // `None` and the openings below walk the full host tree. - #[cfg(feature = "cuda")] - let main_dev_proofs: Option>> = if is_preprocessed { - None - } else { - lde_trace - .gpu_main() - .and_then(|h| h.tree.as_ref()) - .map(|tree| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident main-tree opening"); - // Row-pair leaves: one proof per query at position `challenge`. - crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream).expect( - "device main-tree gather failed; resident tree has no host fallback", - ) - }) - }; - - // Same for the aux trace tree, when it is device resident. - #[cfg(feature = "cuda")] - let aux_dev_proofs: Option>> = round_1_result - .aux - .as_ref() - .and_then(|_aux| lde_trace.gpu_aux().and_then(|h| h.tree.as_ref())) - .map(|tree| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident aux-tree opening"); - // Row-pair leaves: one proof per query at position `challenge`. - crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream) - .expect("device aux-tree gather failed; resident tree has no host fallback") - }); - - // Composition tree: openings open a single position `index` (row pair - // leaf), so gather one proof per query challenge from the device tree. - #[cfg(feature = "cuda")] - let comp_dev_proofs: Option>> = - round_2_result.gpu_composition_tree.as_ref().map(|tree| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident composition-tree opening"); - crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream).expect( - "device composition-tree gather failed; resident tree has no host fallback", - ) - }); - - for (qi, index) in indexes_to_open.iter().enumerate() { - #[cfg(not(feature = "cuda"))] - let _ = qi; + for index in indexes_to_open.iter() { // For preprocessed tables, open the main split (multiplicities only); // for normal tables, open all main columns. let main_trace_opening = if is_preprocessed { @@ -1892,24 +1828,9 @@ pub trait IsStarkProver< lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) }) } else { - #[cfg(feature = "cuda")] - { - if let Some(proofs) = &main_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_main_row(row) - }) - } else { - Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row(row) - }) - } - } - #[cfg(not(feature = "cuda"))] - { - Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row(row) - }) - } + Self::open_polys_with(domain, &main_commit.tree, *index, |row| { + lde_trace.gather_main_row(row) + }) }; // For preprocessed tables, also open the precomputed-columns tree. @@ -1919,52 +1840,16 @@ pub trait IsStarkProver< }) }); - let composition_openings = { - #[cfg(feature = "cuda")] - { - if let Some(proofs) = &comp_dev_proofs { - Self::open_composition_poly_with_proof( - proofs[qi].clone(), - &round_2_result.lde_composition_poly_evaluations, - *index, - ) - } else { - Self::open_composition_poly( - &round_2_result.composition_poly_merkle_tree, - &round_2_result.lde_composition_poly_evaluations, - *index, - ) - } - } - #[cfg(not(feature = "cuda"))] - { - Self::open_composition_poly( - &round_2_result.composition_poly_merkle_tree, - &round_2_result.lde_composition_poly_evaluations, - *index, - ) - } - }; + let composition_openings = Self::open_composition_poly( + &round_2_result.composition_poly_merkle_tree, + &round_2_result.lde_composition_poly_evaluations, + *index, + ); let aux_trace_polys = round_1_result.aux.as_ref().map(|aux| { - #[cfg(feature = "cuda")] - { - if let Some(proofs) = &aux_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_aux_row(row) - }) - } else { - Self::open_polys_with(domain, &aux.tree, *index, |row| { - lde_trace.gather_aux_row(row) - }) - } - } - #[cfg(not(feature = "cuda"))] - { - Self::open_polys_with(domain, &aux.tree, *index, |row| { - lde_trace.gather_aux_row(row) - }) - } + Self::open_polys_with(domain, &aux.tree, *index, |row| { + lde_trace.gather_aux_row(row) + }) }); openings.push(DeepPolynomialOpening { @@ -1998,11 +1883,11 @@ pub trait IsStarkProver< /// # Warning /// /// The transcript must be safely initialized before passing it to this method. - fn multi_prove( - mut air_trace_pairs: Vec>, + fn prove_rounds_1_to_3( + air_trace_pairs: &mut Vec>, transcript: &mut (impl IsStarkTranscript + Clone + Send), #[cfg(feature = "disk-spill")] storage_mode: StorageMode, - ) -> Result, ProvingError> + ) -> Result, ProvingError> where FieldElement: AsBytes, FieldElement: AsBytes, @@ -2032,8 +1917,6 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_prepass"); // Deduplicate Domain + LdeTwiddles by (trace_length, blowup_factor, coset_offset). // Many tables share the same domain size (e.g., 7+ tables at 2^20). @@ -2076,37 +1959,10 @@ pub trait IsStarkProver< let k = table_parallelism().min(num_airs).max(1); - // VRAM budgeted admission. The budget caps the summed device working set - // of the tables proved concurrently so large blocks don't exhaust VRAM. - // It is an extra ceiling on top of `k` (it never raises concurrency). On - // non-cuda builds, or when the budget can't be queried, it is `u64::MAX` - // and chunking falls back to fixed size `k`. - #[cfg(feature = "cuda")] - let vram_budget = math_cuda::device::backend() - .map(|b| b.vram_budget_bytes()) - .unwrap_or(u64::MAX); - #[cfg(not(feature = "cuda"))] - let vram_budget = u64::MAX; - - // R1 main commit: only the main LDE and its Merkle scratch are resident, - // so the aux columns add nothing to this phase's working set. - let main_chunks = { - let estimates: Vec = air_trace_pairs - .iter() - .enumerate() - .map(|(idx, (_, trace, _))| { - let lde_size = - domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) - }) - .collect(); - plan_table_chunks(&estimates, k, vram_budget) - }; - // Spill main traces to mmap before Round 1 LDE. #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { - crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(_, trace, _)| { + crate::par::par_try_for_each_mut(air_trace_pairs, |(_, trace, _)| { trace .main_table .spill_to_disk() @@ -2114,8 +1970,6 @@ pub trait IsStarkProver< })?; } - #[cfg(feature = "instruments")] - drop(__sp); #[cfg(feature = "instruments")] let prepass_elapsed = phase_start.elapsed(); #[cfg(feature = "instruments")] @@ -2131,8 +1985,6 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_main_commit"); let mut main_commits: Vec> = Vec::with_capacity(num_airs); let mut main_ldes: Vec<(Vec>, usize)> = Vec::with_capacity(num_airs); @@ -2143,7 +1995,8 @@ pub trait IsStarkProver< let mut main_gpu_handles: Vec> = Vec::with_capacity(num_airs); - for &(chunk_start, chunk_end) in &main_chunks { + for chunk_start in (0..num_airs).step_by(k) { + let chunk_end = (chunk_start + k).min(num_airs); let chunk_range = chunk_start..chunk_end; let chunk_results: Vec> = @@ -2174,7 +2027,9 @@ pub trait IsStarkProver< if let Some(ref pre_root) = commit.precomputed_root { transcript.append_bytes(pre_root); } - transcript.append_bytes(&commit.root); + // Per-table main root is no longer absorbed here — Task 2 + // replaces it with ONE batched root below, absorbed after all + // tables' main commits are collected. main_commits.push(commit); main_ldes.push(cached_main); #[cfg(feature = "cuda")] @@ -2182,8 +2037,6 @@ pub trait IsStarkProver< } } - #[cfg(feature = "instruments")] - drop(__sp); #[cfg(feature = "instruments")] let main_commits_elapsed = phase_start.elapsed(); #[cfg(feature = "instruments")] @@ -2191,6 +2044,19 @@ pub trait IsStarkProver< heap_snaps.push(s); } + // Absorb ONE batched root over every table's main-split LDE, replacing + // the N per-table root absorptions above. This MMCS is transient: we + // only need its root for the transcript here; Round 4 still opens the + // existing per-table trees (kept in `main_commits`) unchanged. + // + // Task 5: rebuild this MMCS (or thread it through instead of dropping + // it) so Round 4 opens it directly and the per-table trees can be + // dropped. Until then this is a second full pass over every table's + // main LDE columns purely to extract a root — real, but transient, + // memory/CPU cost that the Task 10 perf gate should account for. + let main_mmcs = Self::build_batched_main_mmcs(&main_commits, &main_ldes); + transcript.append_bytes(&main_mmcs.root()); + // ===================================================================== // Round 1, Phase B: Sample shared LogUp challenges // ===================================================================== @@ -2219,8 +2085,6 @@ pub trait IsStarkProver< // but outer parallelism over 12 tables also helps on high-core-count machines. #[cfg(feature = "instruments")] let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_build"); // Disk-spill needs the aux columns in the host trace to spill them, so // disable the GPU-resident aux build (it would keep them device-only). @@ -2277,7 +2141,7 @@ pub trait IsStarkProver< // Spill all aux trace tables to mmap before any Round 1 aux LDE work. #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { - crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(air, trace, _)| { + crate::par::par_try_for_each_mut(air_trace_pairs, |(air, trace, _)| { if air.has_aux_trace() { trace .spill_aux_to_disk() @@ -2287,8 +2151,6 @@ pub trait IsStarkProver< })?; } - #[cfg(feature = "instruments")] - drop(__sp); #[cfg(feature = "instruments")] let aux_build_elapsed = phase_start.elapsed(); #[cfg(feature = "instruments")] @@ -2296,8 +2158,785 @@ pub trait IsStarkProver< heap_snaps.push(s); } - // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K. - // Each table gets its own transcript fork. + // Pass 2: Parallel extract → LDE → commit aux traces in chunks of K. + // Aux roots are batched into ONE shared MMCS root (absorbed after this + // pass), so forks are created AFTER the aux root is bound — see below. + #[cfg(feature = "instruments")] + let phase_start = Instant::now(); + + // Parallel aux commit in chunks of K. The closure returns a cfg-gated + // AuxResult. Under cuda it carries the optional ext3 GPU LDE handle as + // a third element, so Phase D's zip chain keeps it paired with its + // table without a separate handle vector. + #[cfg(feature = "cuda")] + type AuxResult = ( + Option>, + (Vec>, usize), + Option, + ); + #[cfg(not(feature = "cuda"))] + type AuxResult = (Option>, (Vec>, usize)); + #[allow(clippy::type_complexity)] + let mut aux_results: Vec> = Vec::with_capacity(num_airs); + + for chunk_start in (0..num_airs).step_by(k) { + let chunk_end = (chunk_start + k).min(num_airs); + let chunk_range = chunk_start..chunk_end; + + #[allow(clippy::type_complexity)] + let chunk_aux: Vec, ProvingError>> = + crate::par::par_map_collect(chunk_range, |idx| { + let (air, trace, _) = &air_trace_pairs[idx]; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; + + if air.has_aux_trace() { + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + + // Resident GPU path: aux columns already on device (from + // the resident LogUp aux build) — LDE straight from device + // memory, no upload, no host column extraction. When the + // resident build fired the host aux trace is empty, so a + // device LDE failure is a hard abort, not a fall through to + // the host path below (which would commit a zero aux trace). + #[cfg(feature = "cuda")] + if let Some(ra) = trace.aux_resident() { + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + let (tree, handle, aux_data) = + crate::gpu_lde::try_expand_leaf_and_tree_ext3_row_major_keep_dev::< + Field, + FieldExtension, + BatchedMerkleTreeBackend, + >( + ra, domain.blowup_factor, &twiddles.coset_weights + ) + .ok_or_else(|| { + ProvingError::Fft( + "resident aux LDE failed; host aux trace is empty" + .to_string(), + ) + })?; + let num_cols = ra.num_aux_cols; + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); + let root = tree.root; + return Ok(( + Some(TableCommit::plain(tree, root)), + (aux_data, num_cols), + Some(handle), + )); + } + + // Fused GPU path (cuda only): row-major ext3 NTT — single + // H2D, no column extraction, no CPU transpose. + #[cfg(feature = "cuda")] + { + let (trace_slice, num_cols) = trace.aux_data_row_major(); + let n = if num_cols > 0 { + trace_slice.len() / num_cols + } else { + 0 + }; + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + if let Some((tree, handle, aux_data)) = + crate::gpu_lde::try_expand_leaf_and_tree_ext3_row_major_keep::< + Field, + FieldExtension, + BatchedMerkleTreeBackend, + >( + trace_slice, + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, + ) + { + #[cfg(feature = "instruments")] + let aux_lde_dur = t_sub.elapsed(); + let root = tree.root; + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(aux_lde_dur, Duration::ZERO); + return Ok(( + Some(TableCommit::plain(tree, root)), + (aux_data, num_cols), + Some(handle), + )); + } + } + + // CPU path: copy the already-row-major aux trace directly + // (one memcpy — no transpose) and expand with the + // cache-blocked batched two-half FFT. + let (trace_data, total_cols) = trace.aux_data_row_major(); + + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + + let mut aux_data: Vec> = + Vec::with_capacity(lde_size * total_cols); + aux_data.extend_from_slice(trace_data); + + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + trace.aux_table.advise_drop_cache(); + } + + Polynomial::>::coset_lde_full_expand_row_major::( + &mut aux_data, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + &twiddles.two_half_inv, + &twiddles.two_half_fwd, + ) + .expect("row-major aux coset LDE expansion"); + + #[cfg(feature = "instruments")] + let aux_lde_dur = t_sub.elapsed(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + #[allow(unused_mut)] + let (mut tree, root) = + Self::commit_rows_bit_reversed(&aux_data, total_cols) + .ok_or(ProvingError::EmptyCommitment)?; + #[cfg(feature = "disk-spill")] + Self::spill_tree(&mut tree, storage_mode, "aux Merkle tree")?; + let commit = TableCommit::plain(tree, root); + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(aux_lde_dur, t_sub.elapsed()); + + #[cfg(feature = "cuda")] + return Ok((Some(commit), (aux_data, total_cols), None)); + #[cfg(not(feature = "cuda"))] + Ok((Some(commit), (aux_data, total_cols))) + } else { + #[cfg(feature = "cuda")] + return Ok((None, (Vec::new(), 0), None)); + #[cfg(not(feature = "cuda"))] + Ok((None, (Vec::new(), 0))) + } + }); + + // Collect aux commits/LDEs; roots are absorbed as ONE batched MMCS + // root below (not per-table into forks). + for result in chunk_aux.into_iter() { + aux_results.push(result?); + } + } + + // Absorb ONE batched root over every aux-carrying table's LDE (batched + // unified-shard plan), replacing the per-fork aux root absorptions. This + // MMCS is transient: only its root is needed for the transcript here; + // Round 4 still opens the per-table aux trees (kept in `aux_results`). + // Built AFTER Phase B LogUp challenges (aux depends on them) and BEFORE + // forking, so the shared aux root is bound into every fork. + let aux_mmcs: Option> = { + let mut aux_mats: Vec<(Vec>, usize, usize)> = + Vec::with_capacity(aux_results.len()); + for res in aux_results.iter() { + // AuxResult is cfg-gated (an extra GPU handle under cuda); bind + // the fields by name so the extraction is shape-agnostic. + #[cfg(not(feature = "cuda"))] + let (aux_commit, aux_lde) = res; + #[cfg(feature = "cuda")] + let (aux_commit, aux_lde, _aux_gpu) = res; + if aux_commit.is_none() { + continue; + } + let (aux_data, total_cols) = aux_lde; + let num_rows = aux_data.len() / *total_cols; + let log_height = num_rows.trailing_zeros() as usize; + let mut data: Vec> = + Vec::with_capacity(num_rows * *total_cols); + for r in 0..num_rows { + let br = reverse_index(r, num_rows as u64); + let src = br * *total_cols; + data.extend_from_slice(&aux_data[src..src + *total_cols]); + } + aux_mats.push((data, log_height, *total_cols)); + } + if aux_mats.is_empty() { + None + } else { + let m = MixedMmcs::::commit(&aux_mats); + transcript.append_bytes(&m.root()); + Some(m) + } + }; + + // Build commitments and cached LDEs as separate vecs: + // commitments are borrowed in Phase D, LDEs are consumed by value. + let mut commitments: Vec> = + Vec::with_capacity(num_airs); + let mut cached_ldes: Vec> = Vec::with_capacity(num_airs); + // Under cuda, fold main_gpu_handles into the zip chain so each handle + // stays paired with its table by construction. + #[cfg(feature = "cuda")] + let main_iter = main_commits + .into_iter() + .zip(main_ldes) + .zip(main_gpu_handles); + #[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) + { + #[cfg(feature = "cuda")] + let ((main_commit, main_lde), gpu_main) = main_pack; + #[cfg(not(feature = "cuda"))] + let (main_commit, main_lde) = main_pack; + #[cfg(feature = "cuda")] + let (aux_commit, cached_aux, gpu_aux) = aux_full; + #[cfg(not(feature = "cuda"))] + let (aux_commit, cached_aux) = aux_full; + commitments.push(Round1Commitments { + main: main_commit, + aux: aux_commit, + rap_challenges: lookup_challenges.clone(), + bus_public_inputs, + }); + #[cfg(feature = "cuda")] + cached_ldes.push(Lde { + main: main_lde, + aux: cached_aux, + gpu_main, + gpu_aux, + }); + #[cfg(not(feature = "cuda"))] + cached_ldes.push(Lde { + main: main_lde, + aux: cached_aux, + }); + } + + #[cfg(feature = "instruments")] + let aux_commit_elapsed = phase_start.elapsed(); + #[cfg(feature = "instruments")] + if let Some(s) = crate::instruments::snap("After aux commit") { + heap_snaps.push(s); + } + + #[cfg(feature = "debug-checks")] + Self::run_debug_checks(air_trace_pairs, &commitments, &domains, &twiddle_caches); + + // ===================================================================== + // Rounds 2-4: linear shared transcript (unified-shard; forks dropped) + // ===================================================================== + // The forked-per-table transcript is gone. Round 2 composition roots are + // batched into ONE MMCS root; round 3 (z, OOD) and round 4 (FRI) run + // per-table sequentially on the shared transcript with PER-TABLE z. + // NOTE: round 4 is still per-table here — the single batched FRI and the + // shard proof format land in the next step. + + #[cfg(feature = "instruments")] + let phase_start = Instant::now(); + #[cfg(feature = "instruments")] + let table_timings: Vec<(String, usize, Duration, crate::instruments::TableSubOps)> = + Vec::with_capacity(num_airs); + + // Bus contributions bind first (read from the commitments — the same + // value build_round1 copies into Round1), before any round-2 challenge. + for c in commitments.iter() { + if let Some(bpi) = c.bus_public_inputs.as_ref() { + transcript.append_field_element(&bpi.table_contribution); + } + } + + // <<<< Receive per-table challenge: beta_i (one per table, sequential). + let betas: Vec> = (0..num_airs) + .map(|_| transcript.sample_field_element()) + .collect(); + + // Round 2 (parallel over tables): build Round1 from the cached LDE + // (consumed) and compute the composition polynomial. Round1 is built and + // consumed inside the same task, so no &Round1 crosses threads. + #[allow(clippy::type_complexity)] + let round_12: Vec< + Result<(Round1, Round2), ProvingError>, + > = { + #[cfg(feature = "parallel")] + let iter = cached_ldes.into_par_iter().enumerate(); + #[cfg(not(feature = "parallel"))] + let iter = cached_ldes.into_iter().enumerate(); + iter.map(|(idx, lde)| { + let air = air_trace_pairs[idx].0; + let pub_inputs = air_trace_pairs[idx].2; + let domain = &domains[idx]; + let round_1_result = + commitments[idx].build_round1(lde, air.step_size(), domain.blowup_factor); + + let trace_length = domain.interpolation_domain_size; + let num_boundary_constraints = air + .boundary_constraints( + pub_inputs, + &round_1_result.rap_challenges, + round_1_result.bus_public_inputs.as_ref(), + trace_length, + ) + .constraints + .len(); + let num_transition_constraints = air.context().num_transition_constraints; + let mut coefficients: Vec<_> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * betas[idx])) + .take(num_boundary_constraints + num_transition_constraints) + .collect(); + let transition_coefficients: Vec<_> = + coefficients.drain(..num_transition_constraints).collect(); + let boundary_coefficients = coefficients; + + let round_2_result = Self::round_2_compute_composition_polynomial( + air, + pub_inputs, + domain, + &twiddle_caches[idx], + &round_1_result, + &transition_coefficients, + &boundary_coefficients, + )?; + Ok((round_1_result, round_2_result)) + }) + .collect() + }; + let mut round1s: Vec> = Vec::with_capacity(num_airs); + let mut round2s: Vec> = Vec::with_capacity(num_airs); + for r in round_12 { + let (r1, r2) = r?; + round1s.push(r1); + round2s.push(r2); + } + + // >>>> Send commitment: ONE batched composition-poly root over all + // tables' composition parts (mixed-height MMCS). Transient — round 4 + // still opens the per-table composition trees. + let comp_mmcs: MixedMmcs = { + let mut comp_mats: Vec<(Vec>, usize, usize)> = + Vec::with_capacity(num_airs); + for r2 in round2s.iter() { + let cols = &r2.lde_composition_poly_evaluations; + let width = cols.len(); + if width == 0 { + continue; + } + let num_rows = cols[0].len(); + let log_height = num_rows.trailing_zeros() as usize; + let mut data: Vec> = + Vec::with_capacity(num_rows * width); + for r in 0..num_rows { + let br = reverse_index(r, num_rows as u64); + for col in cols.iter() { + data.push(col[br]); + } + } + comp_mats.push((data, log_height, width)); + } + debug_assert!( + !comp_mats.is_empty(), + "every table produces a composition matrix" + ); + let m = MixedMmcs::::commit(&comp_mats); + transcript.append_bytes(&m.root()); + m + }; + + // ONE shared OOD point z for the whole epoch. Sampled against the + // TALLEST table's domain, which is a superset of every shorter table's + // LDE and trace domain, so z is out-of-domain for all tables at once. + // Cleaner than per-table z and simpler to mirror in the verifier; + // identical to per-table z when there is a single table. + let tallest = (0..num_airs) + .max_by_key(|&i| domains[i].lde_roots_of_unity_coset.len()) + .expect("at least one table in the epoch"); + let z = transcript.sample_z_ood( + &domains[tallest].lde_roots_of_unity_coset, + &domains[tallest].trace_roots_of_unity, + ); + + // Round 3: per-table OOD at the shared z; absorb all tables' OOD before + // round 4 (round 4 needs the full post-OOD transcript state). + let mut round3s: Vec> = Vec::with_capacity(num_airs); + for idx in 0..num_airs { + let air = air_trace_pairs[idx].0; + let domain = &domains[idx]; + let round_3_result = Self::round_3_evaluate_polynomials_in_out_of_domain_element( + air, + domain, + &round1s[idx], + &round2s[idx], + &z, + ); + // >>>> Send values: t_j(z g^k) + for col in round_3_result.trace_ood_evaluations.columns().iter() { + for elem in col.iter() { + transcript.append_field_element(elem); + } + } + // >>>> Send values: H_i(z^N) + for element in round_3_result.composition_poly_parts_ood_evaluation.iter() { + transcript.append_field_element(element); + } + round3s.push(round_3_result); + } + + // Per-table FRI heights (lde_log_height), canonical order — the batched + // FRI + histogram binding key. + let heights: Vec = domains + .iter() + .map(|d| d.lde_roots_of_unity_coset.len().trailing_zeros() as usize) + .collect(); + + #[cfg(feature = "instruments")] + { + // Coarse report: rounds_2_4 covers rounds 2-3 (round 4 now runs in the + // caller); per-table table_timings dropped by the rounds-1-3 split. + crate::instruments::store(crate::instruments::MultiProveTiming { + prepass: prepass_elapsed, + main_commits: main_commits_elapsed, + aux_build: aux_build_elapsed, + aux_commit: aux_commit_elapsed, + rounds_2_4: phase_start.elapsed(), + round1_sub: crate::instruments::take_r1_sub(), + table_timings, + heap_snapshots: heap_snaps, + }); + } + + Ok(RoundsOneToThree { + round1s, + round2s, + round3s, + z, + domains, + main_mmcs, + aux_mmcs, + comp_mmcs, + heights, + }) + } + + /// Reference (per-table) proof path: rounds 1-3 shared with the batched path + /// via `prove_rounds_1_to_3`, then a per-table FRI + `StarkProof` per table. + /// Generates STARK proofs for one or more AIRs with a shared transcript. + /// + /// # Multi-Table Proving with LogUp + /// + /// When proving multiple tables that communicate via LogUp (lookup arguments), + /// all tables must use the **same** random challenges (z, α) for the LogUp bus + /// to balance correctly. This function ensures challenge sharing by: + /// + /// 1. **Commit all main traces**: All main trace commitments go into the + /// transcript before any challenges are sampled. + /// 2. **Sample shared LogUp challenges**: The challenges (z, α) are sampled + /// once from the transcript and shared by all AIRs. + /// 3. **Build auxiliary traces**: Each AIR builds its LogUp running-sum + /// columns using the shared challenges. + /// 4. **Rounds 2-4**: Standard STARK protocol rounds for each AIR. + /// + /// # Warning + /// + /// The transcript must be safely initialized before passing it to this method. + fn multi_prove( + mut 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, + { + info!("Started proof generation..."); + + #[cfg(feature = "instruments")] + crate::instruments::reset_all(); + #[cfg(feature = "instruments")] + let mut heap_snaps: Vec = Vec::new(); + + let num_airs = air_trace_pairs.len(); + + // Check if any AIR has an auxiliary trace + let needs_lookup_challenges = air_trace_pairs + .iter() + .any(|(air, _, _)| air.has_aux_trace()); + + // ===================================================================== + // Pre-pass: compute domains and twiddles + // ===================================================================== + + #[cfg(feature = "instruments")] + let phase_start = Instant::now(); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_prepass"); + + // Deduplicate Domain + LdeTwiddles by (trace_length, blowup_factor, coset_offset). + // Many tables share the same domain size (e.g., 7+ tables at 2^20). + // Without dedup, each creates its own Domain (~24 MB) and LdeTwiddles (~32 MB). + type DomainEntry = (Arc>, Arc>); + let mut domain_cache: std::collections::HashMap<(usize, usize, u64), DomainEntry> = + std::collections::HashMap::new(); + + let mut domains = Vec::with_capacity(num_airs); + let mut twiddle_caches: Vec>> = Vec::with_capacity(num_airs); + + for (air, trace, _pub_inputs) in &*air_trace_pairs { + let trace_length = trace.num_rows(); + let blowup = air.options().blowup_factor as usize; + let coset_offset = air.options().coset_offset; + let key = (trace_length, blowup, coset_offset); + + #[cfg(test)] + let was_hit = domain_cache.contains_key(&key); + + let (domain, twiddles) = domain_cache + .entry(key) + .or_insert_with(|| { + let d = Domain::new(*air, trace_length); + let t = LdeTwiddles::new(&d); + (Arc::new(d), Arc::new(t)) + }) + .clone(); + + #[cfg(test)] + crate::tests::domain_cache_stats::record(was_hit); + + domains.push(domain); + twiddle_caches.push(twiddles); + } + // Free the HashMap (which holds extra strong Arc references) before the + // long proving rounds begin. `domains` and `twiddle_caches` already hold + // the only surviving Arcs we care about. + drop(domain_cache); + + let k = table_parallelism().min(num_airs).max(1); + + // VRAM budgeted admission. The budget caps the summed device working set + // of the tables proved concurrently so large blocks don't exhaust VRAM. + // It is an extra ceiling on top of `k` (it never raises concurrency). On + // non-cuda builds, or when the budget can't be queried, it is `u64::MAX` + // and chunking falls back to fixed size `k`. + #[cfg(feature = "cuda")] + let vram_budget = math_cuda::device::backend() + .map(|b| b.vram_budget_bytes()) + .unwrap_or(u64::MAX); + #[cfg(not(feature = "cuda"))] + let vram_budget = u64::MAX; + + // R1 main commit: only the main LDE and its Merkle scratch are resident, + // so the aux columns add nothing to this phase's working set. + let main_chunks = { + let estimates: Vec = air_trace_pairs + .iter() + .enumerate() + .map(|(idx, (_, trace, _))| { + let lde_size = + domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) + }) + .collect(); + plan_table_chunks(&estimates, k, vram_budget) + }; + + // Spill main traces to mmap before Round 1 LDE. + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(_, trace, _)| { + trace + .main_table + .spill_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("early main: {e}"))) + })?; + } + + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] + let prepass_elapsed = phase_start.elapsed(); + #[cfg(feature = "instruments")] + if let Some(s) = crate::instruments::snap("After pool alloc") { + heap_snaps.push(s); + } + + // ===================================================================== + // Round 1, Phase A: Commit all main traces (parallel in chunks of K) + // ===================================================================== + // All main trace commitments must be in the transcript before sampling + // LogUp challenges. + + #[cfg(feature = "instruments")] + let phase_start = Instant::now(); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_main_commit"); + + let mut main_commits: Vec> = Vec::with_capacity(num_airs); + let mut main_ldes: Vec<(Vec>, usize)> = Vec::with_capacity(num_airs); + // Optional device-side LDE handle per table, populated only when the + // R1 fused GPU pipeline produced one. Threaded through Phase D's zip + // chain so each handle stays paired with its table by construction. + #[cfg(feature = "cuda")] + let mut main_gpu_handles: Vec> = + Vec::with_capacity(num_airs); + + for &(chunk_start, chunk_end) in &main_chunks { + let chunk_range = chunk_start..chunk_end; + + let chunk_results: Vec> = + crate::par::par_map_collect(chunk_range, |idx| { + let (air, trace, _) = &air_trace_pairs[idx]; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; + + let precomputed = air + .is_preprocessed() + .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + Self::commit_main_trace( + *trace, + domain, + twiddles, + precomputed, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + }); + + // Sequential: append roots to shared transcript (Fiat-Shamir ordering) + for result in chunk_results { + #[cfg(feature = "cuda")] + let (commit, cached_main, gpu_main) = result?; + #[cfg(not(feature = "cuda"))] + let (commit, cached_main) = result?; + if let Some(ref pre_root) = commit.precomputed_root { + transcript.append_bytes(pre_root); + } + transcript.append_bytes(&commit.root); + main_commits.push(commit); + main_ldes.push(cached_main); + #[cfg(feature = "cuda")] + main_gpu_handles.push(gpu_main); + } + } + + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] + let main_commits_elapsed = phase_start.elapsed(); + #[cfg(feature = "instruments")] + if let Some(s) = crate::instruments::snap("After main commits") { + heap_snaps.push(s); + } + + // ===================================================================== + // Round 1, Phase B: Sample shared LogUp challenges + // ===================================================================== + + let lookup_challenges: Vec> = if needs_lookup_challenges { + (0..LOGUP_NUM_CHALLENGES) + .map(|_| transcript.sample_field_element()) + .collect() + } else { + Vec::new() + }; + + // ===================================================================== + // Phase C + Rounds 2-4: Forked per table + // ===================================================================== + // Each table gets an independent transcript fork (cloned from the shared + // state after Phase B, domain-separated by table index). This matches + // the verifier's forking and makes per-table proving independent. + // + // Split into two passes for parallelism: + // Pass 1 (parallel): Build all auxiliary traces (fingerprint + batch inversion) + // Pass 2 (parallel): Fork transcript → extract → LDE → commit + + // Pass 1: Build aux traces in parallel. + // Each build_auxiliary_trace has internal parallelism (batch_inverse, par_chunks), + // but outer parallelism over 12 tables also helps on high-core-count machines. + #[cfg(feature = "instruments")] + let phase_start = Instant::now(); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_aux_build"); + + // Disk-spill needs the aux columns in the host trace to spill them, so + // disable the GPU-resident aux build (it would keep them device-only). + #[cfg(all(feature = "cuda", feature = "disk-spill"))] + if storage_mode == StorageMode::Disk { + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.set_resident_aux_ok(false); + } + } + + // Thread each table's device-resident trace-domain main columns (kept by + // the R1 main LDE) onto its trace so the LogUp aux fingerprint kernel + // reads them in place instead of re-uploading ~3 GB. Tables without a GPU + // main handle (CPU LDE, preprocessed) fall back to the host upload path. + #[cfg(all(feature = "cuda", not(feature = "debug-checks")))] + for ((_, trace, _), gpu_main) in air_trace_pairs.iter_mut().zip(main_gpu_handles.iter()) { + if let Some(handle) = gpu_main + && let Some(td) = &handle.trace_dev + { + trace.set_main_trace_dev(std::sync::Arc::clone(td), handle.trace_rows); + } + } + + #[cfg(feature = "parallel")] + let aux_iter = air_trace_pairs.par_iter_mut(); + #[cfg(not(feature = "parallel"))] + let aux_iter = air_trace_pairs.iter_mut(); + let bus_inputs_vec: Vec>> = aux_iter + .map(|(air, trace, _)| { + if air.has_aux_trace() { + air.build_auxiliary_trace(*trace, &lookup_challenges) + } else { + None + } + }) + .collect(); + + // The trace-domain snapshots retained by the R1 main LDE (both Arcs: + // trace.main_trace_dev and GpuLdeBase.trace_dev) have exactly one + // consumer — the aux build above. Drop them now so the main-trace-sized + // device buffers are reclaimed before the aux-commit + DEEP/FRI VRAM + // peak instead of living to the end of the proof. + #[cfg(feature = "cuda")] + { + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.clear_main_trace_dev(); + } + for handle in main_gpu_handles.iter_mut().flatten() { + handle.trace_dev = None; + handle.trace_rows = 0; + } + } + + // Spill all aux trace tables to mmap before any Round 1 aux LDE work. + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(air, trace, _)| { + if air.has_aux_trace() { + trace + .spill_aux_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; + } + Ok::<(), ProvingError>(()) + })?; + } + + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] + let aux_build_elapsed = phase_start.elapsed(); + #[cfg(feature = "instruments")] + if let Some(s) = crate::instruments::snap("After aux build") { + heap_snaps.push(s); + } + + // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K. + // Each table gets its own transcript fork. #[cfg(feature = "instruments")] let phase_start = Instant::now(); #[cfg(feature = "instruments")] @@ -2679,34 +3318,6 @@ pub trait IsStarkProver< Ok(MultiProof { proofs }) } - /// Generate a STARK proof for a single AIR/trace. - /// This is equivalent to calling `multi_prove` with a single-element slice. - fn prove( - air: &dyn AIR, - trace: &mut TraceTable, - pub_inputs: &PI, - transcript: &mut (impl IsStarkTranscript + Clone + Send), - ) -> Result, ProvingError> - where - FieldElement: AsBytes, - FieldElement: AsBytes, - PI: Send + Sync + Clone, - Field: Copy + 'static, - FieldExtension: Copy + 'static, - ::BaseType: SpillSafe, - ::BaseType: SpillSafe, - { - let air_trace_pairs = vec![(air, trace, pub_inputs)]; - Self::multi_prove( - air_trace_pairs, - transcript, - #[cfg(feature = "disk-spill")] - StorageMode::Ram, - ) - .map(|mut multi_proof| multi_proof.proofs.remove(0)) - } - - // TODO: propagate errors instead of unwrap() in open_deep_composition_poly and FRI operations /// Executes rounds 2-4 and generates a STARK proof for the trace `main_trace` with public inputs `pub_inputs`. /// Warning: the transcript must be safely initialized before passing it to this method. fn prove_rounds_2_to_4( @@ -2870,6 +3481,446 @@ pub trait IsStarkProver< trace_length: domain.interpolation_domain_size, }) } + + /// One table's DEEP composition codeword (bit-reversed, ready for the batched + /// FRI) built from the SHARED gamma. Mirrors the DEEP setup inside + /// `round_4_...` but takes gamma from the shared transcript rather than + /// sampling it per table (cross-table separation is handled by `alpha` in + /// `combine_by_height`). + fn batched_table_deep_codeword( + air: &dyn AIR, + domain: &Domain, + round_1_result: &Round1, + round_2_result: &Round2, + round_3_result: &Round3, + z: &FieldElement, + gamma: &FieldElement, + ) -> Vec> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + { + let n_terms_composition_poly = round_2_result.lde_composition_poly_evaluations.len(); + let num_terms_trace = + air.context().transition_offsets.len() * air.step_size() * air.context().trace_columns; + let mut deep_composition_coefficients: Vec<_> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * gamma)) + .take(n_terms_composition_poly + num_terms_trace) + .collect(); + let trace_term_coeffs: Vec<_> = deep_composition_coefficients + .drain(..num_terms_trace) + .collect::>() + .chunks(air.context().transition_offsets.len() * air.step_size()) + .map(|chunk| chunk.to_vec()) + .collect(); + let gammas = deep_composition_coefficients; + let mut deep_evals = Self::compute_deep_composition_poly_evaluations( + &round_1_result.lde_trace, + round_2_result, + round_3_result, + z, + domain, + &domain.trace_primitive_root, + &gammas, + &trace_term_coeffs, + ); + in_place_bit_reverse_permute(&mut deep_evals); + deep_evals + } + + /// Batched (unified-shard) proof path: rounds 1-3 shared with `multi_prove`, + /// then ONE FRI over the height-combined per-table DEEP codewords, opened + /// from the three shared MMCS trees. Produces a `BatchedMultiProof`. + fn multi_prove_batched( + mut 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 rounds = Self::prove_rounds_1_to_3( + &mut air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + Self::batched_round_4(air_trace_pairs, rounds, transcript) + } + + /// Round 4 of the batched (unified-shard) path, factored out so the + /// continuation epoch driver can reuse it for the VM-table lane. Consumes + /// the shared `RoundsOneToThree`, runs ONE FRI over the height-combined + /// per-table DEEP codewords, and opens the three shared MMCS trees. + fn batched_round_4( + air_trace_pairs: Vec>, + rounds: RoundsOneToThree, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + PI: Send + Sync + Clone, + Field: Copy + 'static, + FieldExtension: Copy + 'static, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, + { + let RoundsOneToThree { + round1s, + round2s, + round3s, + z, + domains, + main_mmcs, + aux_mmcs, + comp_mmcs, + heights, + .. + } = rounds; + + let num_airs = round1s.len(); + + // ===== Round 4 (batched) ===== + // <<<< gamma: ONE shared DEEP intra-table challenge. + let gamma = transcript.sample_field_element(); + + // Per-table DEEP codewords (bit-reversed) paired with lde_log_height. + let mut deep_inputs: Vec<(Vec>, usize)> = + Vec::with_capacity(num_airs); + for idx in 0..num_airs { + let air = air_trace_pairs[idx].0; + let codeword = Self::batched_table_deep_codeword( + air, + &domains[idx], + &round1s[idx], + &round2s[idx], + &round3s[idx], + &z, + &gamma, + ); + deep_inputs.push((codeword, heights[idx])); + } + + // All tables' composition-poly LDE evaluations are now baked into the DEEP + // codewords above and are never read again. Free them before the FRI commit + // and query-opening phases, which keep every table's round1 LDE resident. + drop(round2s); + + // Bind the height histogram, then sample the cross-table batching alpha. + crate::fri::batched::absorb_height_histogram::(transcript, &heights); + let alpha = transcript.sample_field_element(); + + // Combine per-height, then run the batched (fold-and-inject) FRI. + let combined = crate::fri::batched::combine_by_height(&deep_inputs, &alpha); + // All per-table DEEP codewords are folded into `combined`; free them before FRI. + drop(deep_inputs); + let coset_offset = + FieldElement::::from(air_trace_pairs[0].0.context().proof_options.coset_offset); + let (fri_last_value, fri_layers) = crate::fri::batched::batched_commit_phase::< + Field, + FieldExtension, + _, + >(combined, transcript, &coset_offset); + + // Grinding: mirror the per-table round 4 (shared proof options). + let security_bits = air_trace_pairs[0].0.context().proof_options.grinding_factor; + let mut nonce = None; + if security_bits > 0 { + let nonce_value = grinding::generate_nonce(&transcript.state(), security_bits) + .expect("nonce not found"); + transcript.append_bytes(&nonce_value.to_be_bytes()); + nonce = Some(nonce_value); + } + + // Query indices against the tallest domain (2^h_max). + let tallest = (0..num_airs) + .max_by_key(|&i| heights[i]) + .expect("at least one table in the epoch"); + let h_max = heights[tallest]; + let number_of_queries = air_trace_pairs[0].0.options().fri_number_of_queries; + let iotas = Self::sample_query_indexes(number_of_queries, &domains[tallest], transcript); + + let query_list = fri::query_phase(&fri_layers, &iotas); + let fri_layers_merkle_roots: Vec<_> = fri_layers + .iter() + .map(|layer| layer.merkle_tree.root) + .collect(); + + // Per-query openings: one MixedOpening per phase (main/aux/composition) + // covering all tables at once, plus per-preprocessed-table precomputed + // openings (those columns are outside the shared main MMCS). + let mut deep_poly_openings = Vec::with_capacity(iotas.len()); + for &iota in iotas.iter() { + let main = main_mmcs.open_batch(iota); + let aux = aux_mmcs.as_ref().map(|m| m.open_batch(iota)); + let composition = comp_mmcs.open_batch(iota); + + let mut precomputed = Vec::new(); + for idx in 0..num_airs { + if let Some(tree) = round1s[idx].main.precomputed_tree.as_ref() { + let num_precomputed_cols = round1s[idx].main.num_precomputed_cols; + let lde_trace = &round1s[idx].lde_trace; + let local = iota >> (h_max - heights[idx]); + precomputed.push(Self::open_polys_with(&domains[idx], tree, local, |row| { + lde_trace.gather_main_row_range(row, 0, num_precomputed_cols) + })); + } + } + + deep_poly_openings.push(BatchedQueryOpening { + main, + aux, + composition, + precomputed, + }); + } + + // Per-table data (canonical epoch order). + let mut per_table = Vec::with_capacity(num_airs); + for idx in 0..num_airs { + per_table.push(BatchedTableData { + trace_length: domains[idx].interpolation_domain_size, + trace_ood_evaluations: round3s[idx].trace_ood_evaluations.clone(), + composition_poly_parts_ood_evaluation: round3s[idx] + .composition_poly_parts_ood_evaluation + .clone(), + precomputed_root: round1s[idx].main.precomputed_root, + bus_public_inputs: round1s[idx].bus_public_inputs.clone(), + public_inputs: air_trace_pairs[idx].2.clone(), + }); + } + + Ok(BatchedMultiProof { + main_root: main_mmcs.root(), + aux_root: aux_mmcs.as_ref().map(|m| m.root()), + composition_root: comp_mmcs.root(), + fri_layers_merkle_roots, + fri_last_value, + query_list, + nonce, + deep_poly_openings, + per_table, + }) + } + + /// Continuation epoch driver: prove the epoch's VM tables with the batched + /// (unified-shard) FRI while proving the single L2G sub-table as a SEPARATE + /// commitment lane (its own tree + own FRI), so the L2G<->global root binding + /// still holds. Both lanes are woven through ONE transcript: + /// + /// 1. Absorb the L2G main root FIRST (canonical order). + /// 2. `prove_rounds_1_to_3` over the VM tables (absorbs the VM roots, samples + /// the shared LogUp challenge, through OOD — ends at the round-4 seam). + /// 3. At the seam, FORK the transcript (single lane -> no idx bytes; then absorb + /// the L2G aux root, then the L2G bus `table_contribution`) and run + /// `prove_rounds_2_to_4` for L2G -> its own `StarkProof`. + /// 4. `batched_round_4` for the VM tables on the main transcript. + /// + /// The L2G lane is NOT a member of the VM shared MMCS nor the unified FRI: its + /// own tree authenticates the binding root. Mirrored in + /// `IsStarkVerifier::batched_verify_epoch`. + #[allow(clippy::type_complexity)] + fn multi_prove_batched_epoch( + mut vm_pairs: Vec>, + l2g_pair: AirTracePair<'_, Field, FieldExtension, PI>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> Result< + ( + BatchedMultiProof, + StarkProof, + ), + ProvingError, + > + where + FieldElement: AsBytes, + FieldElement: AsBytes, + PI: Send + Sync + Clone, + Field: Copy + 'static, + FieldExtension: Copy + 'static, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, + { + let (l2g_air, l2g_trace, l2g_pub) = l2g_pair; + + // L2G domain + twiddles (its own lane, never in the VM shared MMCS). + let l2g_domain = Domain::new(l2g_air, l2g_trace.num_rows()); + let l2g_tw = LdeTwiddles::new(&l2g_domain); + + // (2) Commit the L2G main trace to its own tree. L2G is never preprocessed. + #[cfg(feature = "cuda")] + let (l2g_main_commit, l2g_main_lde, l2g_gpu_main) = Self::commit_main_trace( + l2g_trace, + &l2g_domain, + &l2g_tw, + None, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + #[cfg(not(feature = "cuda"))] + let (l2g_main_commit, l2g_main_lde) = Self::commit_main_trace( + l2g_trace, + &l2g_domain, + &l2g_tw, + None, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + + // (3) Canonical transcript order: L2G main root FIRST. + transcript.append_bytes(&l2g_main_commit.root); + + // (4) VM rounds 1-3 on the main transcript (absorbs the VM roots + the + // shared LogUp challenge; ends at the round-4 seam, post-OOD). + let vm_rounds = Self::prove_rounds_1_to_3( + &mut vm_pairs, + transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + + // (5) Shared LogUp challenges (sampled inside prove_rounds_1_to_3 Phase B). + let lookup_challenges = vm_rounds + .round1s + .first() + .map(|r| r.rap_challenges.clone()) + .unwrap_or_default(); + + // (6) Build the L2G aux (LogUp) trace with the shared challenges, then + // commit it to its own tree. + let l2g_bus = l2g_air.build_auxiliary_trace(l2g_trace, &lookup_challenges); + let (l2g_aux_commit, l2g_aux_lde) = Self::commit_aux_trace( + l2g_trace, + &l2g_domain, + &l2g_tw, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + + // (7) Assemble the L2G Round1 from its own commitments + LDEs. + let l2g_r1c = Round1Commitments { + main: l2g_main_commit, + aux: Some(l2g_aux_commit), + rap_challenges: lookup_challenges, + bus_public_inputs: l2g_bus, + }; + let l2g_lde = Lde { + main: l2g_main_lde, + aux: l2g_aux_lde, + #[cfg(feature = "cuda")] + gpu_main: l2g_gpu_main, + #[cfg(feature = "cuda")] + gpu_aux: None, + }; + let mut l2g_round1 = + l2g_r1c.build_round1(l2g_lde, l2g_air.step_size(), l2g_domain.blowup_factor); + + // (8) Fork the transcript at the seam. Single lane -> no idx bytes. Absorb + // the L2G aux root, then the bus table_contribution (matches the + // non-batched per-table fork convention used by `multi_prove`). + let mut l2g_fork = transcript.clone(); + if let Some(aux) = l2g_round1.aux.as_ref() { + l2g_fork.append_bytes(&aux.root); + } + if let Some(bpi) = l2g_round1.bus_public_inputs.as_ref() { + l2g_fork.append_field_element(&bpi.table_contribution); + } + let l2g_proof = Self::prove_rounds_2_to_4( + l2g_air, + l2g_pub, + &mut l2g_round1, + &mut l2g_fork, + &l2g_domain, + &l2g_tw, + )?; + + // (9) VM batched Round 4 continues on the main (un-cloned) transcript. + let vm_proof = Self::batched_round_4(vm_pairs, vm_rounds, transcript)?; + + // (10) Two lanes: batched VM proof + standalone L2G proof. + Ok((vm_proof, l2g_proof)) + } + + /// Generate a STARK proof for a single AIR/trace. + /// This is equivalent to calling `multi_prove` with a single-element slice. + fn prove( + air: &dyn AIR, + trace: &mut TraceTable, + pub_inputs: &PI, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + PI: Send + Sync + Clone, + Field: Copy + 'static, + FieldExtension: Copy + 'static, + ::BaseType: SpillSafe, + ::BaseType: SpillSafe, + { + let air_trace_pairs = vec![(air, trace, pub_inputs)]; + Self::multi_prove( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + StorageMode::Ram, + ) + .map(|mut multi_proof| multi_proof.proofs.remove(0)) + } +} + +/// Heuristic peak device bytes for one table: co-resident LDE columns plus the +/// resident Merkle trees, with a scratch factor for NTT and leaf transients. A +/// deliberate over estimate for a safety ceiling, not a precise allocator. Pass +/// aux_cols == 0 when the aux LDE is not yet resident (R1 main commit). +fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) -> u64 { + const BYTES_PER_BASE: u64 = 8; + const EXT3_BYTES: u64 = 24; + const SCRATCH_FACTOR: u64 = 2; + const RESIDENT_TREE_BYTES_PER_LDE: u64 = 256; + let lde = lde_size as u64; + let per_row = (main_cols as u64).saturating_mul(BYTES_PER_BASE) + + (aux_cols as u64).saturating_mul(EXT3_BYTES); + let lde_term = lde.saturating_mul(per_row).saturating_mul(SCRATCH_FACTOR); + let tree_term = lde.saturating_mul(RESIDENT_TREE_BYTES_PER_LDE); + lde_term.saturating_add(tree_term) +} + +/// Plan contiguous table chunks for parallel proving. A chunk grows until it +/// hits `k` tables or its summed VRAM estimate would exceed `budget`; a single +/// table larger than `budget` runs solo. With `budget == u64::MAX` (non-cuda, +/// or VRAM not binding) chunks fall back to fixed size `k`, identical to the +/// old `step_by(k)`, so scheduling and the proof are unchanged. Returns +/// `(start, end)` half open ranges covering `0..estimates.len()` in order. +fn plan_table_chunks(estimates: &[u64], k: usize, budget: u64) -> Vec<(usize, usize)> { + let n = estimates.len(); + let k = k.max(1); + let budget = budget as u128; + let mut chunks = Vec::new(); + let mut start = 0; + while start < n { + let mut end = start; + let mut acc: u128 = 0; + while end < n { + let next = estimates[end] as u128; + // Always admit at least one table per chunk (oversized → solo). + if end > start && (end - start >= k || acc + next > budget) { + break; + } + acc += next; + end += 1; + } + chunks.push((start, end)); + start = end; + } + chunks } /// Print a global bus balance report aggregating per-bus sums across all tables. diff --git a/crypto/stark/src/test_utils.rs b/crypto/stark/src/test_utils.rs index f5cd19f80..e7f814876 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::{BatchedMultiProof, MultiProof}; use crate::prover::{IsStarkProver, Prover, ProvingError}; use crate::trace::TraceTable; use crate::traits::AIR; @@ -36,3 +36,26 @@ where crate::storage_mode::StorageMode::Ram, ) } + +/// Batched (unified-shard) analogue of [`multi_prove_ram`]: produces a +/// `BatchedMultiProof` verified by `Verifier::batched_multi_verify`. +pub fn multi_prove_batched_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_batched( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ) +} diff --git a/crypto/stark/src/tests/air_tests.rs b/crypto/stark/src/tests/air_tests.rs index b6a4108f9..0012ee598 100644 --- a/crypto/stark/src/tests/air_tests.rs +++ b/crypto/stark/src/tests/air_tests.rs @@ -29,7 +29,7 @@ type Felt = FieldElement; use crate::examples::read_only_memory_logup::{ LogReadOnlyPublicInputs, LogReadOnlyRAP, read_only_logup_trace, }; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::multi_prove_batched_ram; #[test_log::test] fn test_prove_fib() { @@ -327,7 +327,7 @@ fn test_multi_prove_fib_3_tables() { (&air_3, &mut trace_3, &pub_inputs_3), ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec< &dyn AIR< @@ -337,7 +337,7 @@ fn test_multi_prove_fib_3_tables() { >, > = vec![&air_1, &air_2, &air_3]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -427,7 +427,7 @@ fn test_multi_prove_2_tables_small_field() { (&air_2, &mut trace_2, &pub_inputs_2), ]; - let multi_proof = multi_prove_ram( + let multi_proof = multi_prove_batched_ram( air_trace_pairs, &mut DefaultTranscript::::new(&[]), ) @@ -441,7 +441,7 @@ fn test_multi_prove_2_tables_small_field() { >, > = vec![&air_1, &air_2]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/tests/bus_tests/batched_soundness_tests.rs b/crypto/stark/src/tests/bus_tests/batched_soundness_tests.rs new file mode 100644 index 000000000..e6d3a40f8 --- /dev/null +++ b/crypto/stark/src/tests/bus_tests/batched_soundness_tests.rs @@ -0,0 +1,192 @@ +//! Soundness negatives for the batched (unified-shard) verifier +//! `Verifier::batched_multi_verify`. +//! +//! Each test builds ONE valid `BatchedMultiProof` (three all-padding, bus-balanced +//! tables — the simplest valid multi-table epoch, Σ table_contribution = 0) and then +//! tampers a single component of the proof, asserting the verifier rejects. Together +//! they exercise every batched-specific verification path: the three shared +//! mixed-height MMCS openings (value + width binding), the fold-and-inject FRI query +//! check (last value, layer root, layer sym), the shared OOD/transcript replay, the +//! grinding nonce, the query-count guard, and the bus-balance check. + +use crypto::fiat_shamir::default_transcript::DefaultTranscript; +use math::field::element::FieldElement; +use math::field::{ + extensions_goldilocks::Degree3GoldilocksExtensionField, goldilocks::GoldilocksField, +}; + +use crate::examples::multi_table_lookup::{ + new_add_air_with_lookup, new_cpu_air_with_lookup, new_mul_air_with_lookup, +}; +use crate::proof::options::ProofOptions; +use crate::proof::stark::BatchedMultiProof; +use crate::test_utils::multi_prove_batched_ram; +use crate::trace::TraceTable; +use crate::traits::AIR; +use crate::verifier::{IsStarkVerifier, Verifier}; + +type F = GoldilocksField; +type E = Degree3GoldilocksExtensionField; +type FE = FieldElement; + +/// Build a valid batched proof over three all-padding (bus-balanced) tables: +/// CPU (5 main columns) + ADD + MUL (4 main columns each), all 4 rows of zeros. +fn valid_padding_proof() -> BatchedMultiProof { + let mut cpu_trace = TraceTable::from_columns_main(vec![vec![FE::zero(); 4]; 5], 1); + let mut add_trace = TraceTable::from_columns_main(vec![vec![FE::zero(); 4]; 4], 1); + let mut mul_trace = TraceTable::from_columns_main(vec![vec![FE::zero(); 4]; 4], 1); + + let proof_options = ProofOptions::default_test_options(); + let cpu_air = new_cpu_air_with_lookup(&proof_options); + let add_air = new_add_air_with_lookup(&proof_options); + let mul_air = new_mul_air_with_lookup(&proof_options); + + let air_trace_pairs: Vec<( + &dyn AIR, + _, + _, + )> = vec![ + (&cpu_air, &mut cpu_trace, &()), + (&add_air, &mut add_trace, &()), + (&mul_air, &mut mul_trace, &()), + ]; + + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap() +} + +/// Verify a batched proof with a fresh verifier (AIRs reconstructed, as a real +/// verifier would) against `expected_bus_balance`. +fn batched_verify( + proof: &BatchedMultiProof, + expected_bus_balance: FieldElement, +) -> bool { + let proof_options = ProofOptions::default_test_options(); + let cpu_air = new_cpu_air_with_lookup(&proof_options); + let add_air = new_add_air_with_lookup(&proof_options); + let mul_air = new_mul_air_with_lookup(&proof_options); + let airs: Vec<&dyn AIR> = + vec![&cpu_air, &add_air, &mul_air]; + Verifier::batched_multi_verify( + &airs, + proof, + &mut DefaultTranscript::::new(&[]), + &expected_bus_balance, + ) +} + +/// Sanity anchor: the untampered proof verifies (guards against a false-reject +/// regression that would make every negative below pass vacuously). +#[test_log::test] +fn batched_valid_padding_proof_verifies() { + let proof = valid_padding_proof(); + assert!( + batched_verify(&proof, FieldElement::::zero()), + "a valid all-padding batched proof must verify" + ); +} + +/// Tampering an opened MAIN-trace evaluation breaks the shared main MMCS auth path. +#[test_log::test] +fn batched_rejects_tampered_main_trace_opening() { + let mut proof = valid_padding_proof(); + proof.deep_poly_openings[0].main.per_matrix[0].evaluations[0] += FE::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering an opened COMPOSITION evaluation breaks the shared composition MMCS auth path. +#[test_log::test] +fn batched_rejects_tampered_composition_opening() { + let mut proof = valid_padding_proof(); + proof.deep_poly_openings[0].composition.per_matrix[0].evaluations[0] += + FieldElement::::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering an opened AUX evaluation breaks the shared aux MMCS auth path. +#[test_log::test] +fn batched_rejects_tampered_aux_opening() { + let mut proof = valid_padding_proof(); + proof.deep_poly_openings[0] + .aux + .as_mut() + .expect("padding tables carry an aux (LogUp) trace") + .per_matrix[0] + .evaluations[0] += FieldElement::::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Width-binding: shrinking a per-matrix opening's column count (without changing the +/// flat leaf bytes it would concatenate into) must be caught by the MMCS width check. +#[test_log::test] +fn batched_rejects_main_opening_width_mismatch() { + let mut proof = valid_padding_proof(); + // Drop one column from the opened main row → evaluations.len() != committed width. + proof.deep_poly_openings[0].main.per_matrix[0] + .evaluations + .pop(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering the final FRI value breaks the fold-and-inject terminal check. +#[test_log::test] +fn batched_rejects_tampered_fri_last_value() { + let mut proof = valid_padding_proof(); + proof.fri_last_value += FieldElement::::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering a committed FRI layer root diverges the derived fold challenges and +/// invalidates that layer's opening. +#[test_log::test] +fn batched_rejects_tampered_fri_layer_root() { + let mut proof = valid_padding_proof(); + assert!( + !proof.fri_layers_merkle_roots.is_empty(), + "expect >= 1 FRI layer" + ); + proof.fri_layers_merkle_roots[0][0] ^= 1; + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering a per-query symmetric layer evaluation breaks that FRI layer's opening. +#[test_log::test] +fn batched_rejects_tampered_layer_evaluation_sym() { + let mut proof = valid_padding_proof(); + assert!( + !proof.query_list[0].layers_evaluations_sym.is_empty(), + "expect >= 1 FRI layer eval" + ); + proof.query_list[0].layers_evaluations_sym[0] += FieldElement::::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Tampering a per-table composition-parts OOD value breaks the step-2 composition claim. +#[test_log::test] +fn batched_rejects_tampered_composition_ood() { + let mut proof = valid_padding_proof(); + proof.per_table[1].composition_poly_parts_ood_evaluation[0] += FieldElement::::one(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Dropping a query opening leaves fewer than `fri_number_of_queries` → rejected. +#[test_log::test] +fn batched_rejects_dropped_query_opening() { + let mut proof = valid_padding_proof(); + proof.deep_poly_openings.pop(); + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// Removing the grinding nonce (with grinding_factor > 0) fails the proof-of-work check. +#[test_log::test] +fn batched_rejects_missing_nonce() { + let mut proof = valid_padding_proof(); + proof.nonce = None; + assert!(!batched_verify(&proof, FieldElement::::zero())); +} + +/// A valid Σ = 0 proof must be rejected when a NONZERO bus balance is expected. +#[test_log::test] +fn batched_rejects_wrong_expected_bus_balance() { + let proof = valid_padding_proof(); + assert!(!batched_verify(&proof, FieldElement::::one())); +} diff --git a/crypto/stark/src/tests/bus_tests/completeness_tests.rs b/crypto/stark/src/tests/bus_tests/completeness_tests.rs index 6f4a1655b..334ee0501 100644 --- a/crypto/stark/src/tests/bus_tests/completeness_tests.rs +++ b/crypto/stark/src/tests/bus_tests/completeness_tests.rs @@ -17,7 +17,7 @@ use crate::lookup::{ NullBoundaryConstraintBuilder, Packing, }; use crate::proof::options::ProofOptions; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::multi_prove_batched_ram; use crate::trace::TraceTable; use crate::traits::AIR; use crate::verifier::{IsStarkVerifier, Verifier}; @@ -123,12 +123,12 @@ fn test_multi_table_proof() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -186,12 +186,12 @@ fn test_all_padding() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -249,12 +249,12 @@ fn test_single_operation() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -312,12 +312,12 @@ fn test_duplicate_operations() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -375,17 +375,17 @@ fn test_serialization_roundtrip() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); // Serialize and deserialize let serialized = serde_cbor::to_vec(&multi_proof).expect("serialization failed"); - let deserialized: crate::proof::stark::MultiProof = + let deserialized: crate::proof::stark::BatchedMultiProof = serde_cbor::from_slice(&serialized).expect("deserialization failed"); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &deserialized, &mut DefaultTranscript::::new(&[]), @@ -520,12 +520,12 @@ fn test_bus_value_features() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender_air, &receiver_air]; - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/tests/bus_tests/mod.rs b/crypto/stark/src/tests/bus_tests/mod.rs index a57ca9aaf..1e4685111 100644 --- a/crypto/stark/src/tests/bus_tests/mod.rs +++ b/crypto/stark/src/tests/bus_tests/mod.rs @@ -1,4 +1,5 @@ //! Tests for LogUp bus interactions. +pub mod batched_soundness_tests; pub mod bus_value_tests; pub mod completeness_tests; pub mod multiplicity_tests; diff --git a/crypto/stark/src/tests/bus_tests/multiplicity_tests.rs b/crypto/stark/src/tests/bus_tests/multiplicity_tests.rs index 8bf7492e4..f63633f8e 100644 --- a/crypto/stark/src/tests/bus_tests/multiplicity_tests.rs +++ b/crypto/stark/src/tests/bus_tests/multiplicity_tests.rs @@ -15,7 +15,7 @@ use crate::lookup::{ NullBoundaryConstraintBuilder, Packing, }; use crate::proof::options::ProofOptions; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::multi_prove_batched_ram; use crate::trace::TraceTable; use crate::traits::AIR; use crate::verifier::{IsStarkVerifier, Verifier}; @@ -111,13 +111,13 @@ fn test_multiplicity_one() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender, &receiver]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -219,13 +219,13 @@ fn test_multiplicity_sum() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender, &receiver]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -325,13 +325,13 @@ fn test_multiplicity_negated() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender, &receiver]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/tests/bus_tests/soundness_tests.rs b/crypto/stark/src/tests/bus_tests/soundness_tests.rs index 652f5e87d..8dd778b4e 100644 --- a/crypto/stark/src/tests/bus_tests/soundness_tests.rs +++ b/crypto/stark/src/tests/bus_tests/soundness_tests.rs @@ -15,7 +15,7 @@ use crate::examples::multi_table_lookup::{ }; use crate::proof::options::ProofOptions; use crate::prover::{IsStarkProver, Prover}; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::{multi_prove_batched_ram, multi_prove_ram}; use crate::trace::TraceTable; use crate::traits::AIR; use crate::verifier::{IsStarkVerifier, Verifier}; @@ -120,13 +120,13 @@ fn test_rejects_inflated_composition_part_count() { (&mul_air, &mut mul_trace, &()), ]; let mut multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&cpu_air, &add_air, &mul_air]; // The untampered proof verifies. - assert!(Verifier::multi_verify( + assert!(Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -134,12 +134,12 @@ fn test_rejects_inflated_composition_part_count() { )); // Tamper: inflate the first table's composition-poly part count. - multi_proof.proofs[0] + multi_proof.per_table[0] .composition_poly_parts_ood_evaluation .push(FieldElement::::zero()); assert!( - !Verifier::multi_verify( + !Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -1838,14 +1838,14 @@ fn test_compound_equals_primitive_expansion() { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender, &receiver]; // This should PASS - compound and primitive expansion are equivalent assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/tests/prove_verify_roundtrip_tests.rs b/crypto/stark/src/tests/prove_verify_roundtrip_tests.rs index a387df476..939e6cd50 100644 --- a/crypto/stark/src/tests/prove_verify_roundtrip_tests.rs +++ b/crypto/stark/src/tests/prove_verify_roundtrip_tests.rs @@ -15,8 +15,8 @@ use crate::lookup::{ NullBoundaryConstraintBuilder, Packing, }; use crate::proof::options::ProofOptions; -use crate::proof::stark::MultiProof; -use crate::test_utils::multi_prove_ram; +use crate::proof::stark::BatchedMultiProof; +use crate::test_utils::multi_prove_batched_ram; use crate::traits::AIR; use crate::verifier::{IsStarkVerifier, Verifier}; @@ -135,7 +135,7 @@ fn test_verify_serialized_multi_table_proofs() { (&mul_air, &mut mul_trace, &()), ]; - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap() + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap() }; // ========================================================================= @@ -147,7 +147,7 @@ fn test_verify_serialized_multi_table_proofs() { // At this point, the prover's data is dropped (out of scope above) // The verifier only has the serialized data - let received_proofs: MultiProof = + let received_proofs: BatchedMultiProof = serde_cbor::from_slice(&serialized).expect("Failed to deserialize proofs"); // ========================================================================= @@ -168,7 +168,7 @@ fn test_verify_serialized_multi_table_proofs() { vec![&cpu_air, &add_air, &mul_air]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &received_proofs, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/tests/prover_tests.rs b/crypto/stark/src/tests/prover_tests.rs index a536a206a..fc729131c 100644 --- a/crypto/stark/src/tests/prover_tests.rs +++ b/crypto/stark/src/tests/prover_tests.rs @@ -8,7 +8,7 @@ use crate::{ }, proof::options::ProofOptions, prover::{IsStarkProver, LdeTwiddles, Prover, evaluate_polynomial_on_lde_domain}, - test_utils::multi_prove_ram, + test_utils::{multi_prove_batched_ram, multi_prove_ram}, tests::domain_cache_stats, tests::trace_test_helpers::get_trace_evaluations, trace::{LDETraceTable, get_trace_evaluations_from_lde}, @@ -335,7 +335,7 @@ fn test_multi_prove_mixed_coset_offsets() { (&air_2, &mut trace_2, &pub_inputs), ]; - let multi_proof = multi_prove_ram( + let multi_proof = multi_prove_batched_ram( air_trace_pairs, &mut DefaultTranscript::::new(&[]), ) @@ -350,7 +350,7 @@ fn test_multi_prove_mixed_coset_offsets() { > = vec![&air_1, &air_2]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), @@ -360,6 +360,51 @@ fn test_multi_prove_mixed_coset_offsets() { ); } +/// Scope B Task 2 smoke test: a >=2-table `multi_prove` must still run to +/// completion and hand back a `MultiProof` (one `StarkProof` per table) +/// without panicking, now that Round 1 Phase A absorbs a single batched +/// `MixedMmcs` root instead of N per-table main roots. Verification is +/// deliberately NOT asserted here — the per-table verifier doesn't understand +/// the batched root yet (Scope B Task 7); this test only proves the batched +/// commit wiring executes end-to-end. +#[test_log::test] +fn test_multi_prove_batched_main_mmcs_smoke() { + let mut trace_1 = simple_fibonacci::fibonacci_trace([Felt::from(1), Felt::from(1)], 8); + let mut trace_2 = simple_fibonacci::fibonacci_trace([Felt::from(1), Felt::from(1)], 16); + + let pub_inputs = FibonacciPublicInputs { + a0: Felt::one(), + a1: Felt::one(), + }; + + let proof_options = ProofOptions::default_test_options(); + let air_1 = FibonacciAIR::::new(&proof_options); + let air_2 = FibonacciAIR::::new(&proof_options); + + let air_trace_pairs: Vec<( + &dyn AIR< + Field = GoldilocksField, + FieldExtension = GoldilocksField, + PublicInputs = FibonacciPublicInputs, + >, + &mut _, + &_, + )> = vec![ + (&air_1, &mut trace_1, &pub_inputs), + (&air_2, &mut trace_2, &pub_inputs), + ]; + + let mut transcript = DefaultTranscript::::new(&[]); + let prove_result = multi_prove_ram(air_trace_pairs, &mut transcript); + let multi_proof = prove_result.expect("proving should succeed"); + + assert_eq!( + multi_proof.proofs.len(), + 2, + "multi_prove should return one StarkProof per table" + ); +} + /// Test that the domain cache deduplicates when multiple AIRs share all three key fields /// `(trace_length, blowup, coset_offset)`. Asserts exactly one `Domain`/`LdeTwiddles` /// construction for N identical AIRs and that the resulting proof still verifies. @@ -402,7 +447,7 @@ fn test_multi_prove_dedups_shared_domain_params() { (&air_3, &mut trace_3, &pub_inputs), ]; - let multi_proof = multi_prove_ram( + let multi_proof = multi_prove_batched_ram( air_trace_pairs, &mut DefaultTranscript::::new(&[]), ) @@ -427,7 +472,7 @@ fn test_multi_prove_dedups_shared_domain_params() { > = vec![&air_1, &air_2, &air_3]; assert!( - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index fdee3d2ff..13546e5fe 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -1,7 +1,7 @@ use super::{ config::BatchedMerkleTreeBackend, domain::VerifierDomain, - fri::fri_decommit::FriDecommitment, + fri::{batched::derive_batched_fri_challenges, fri_decommit::FriDecommitment, mmcs::MixedMmcs}, grinding, proof::stark::StarkProof, traits::{AIR, TransitionEvaluationContext}, @@ -10,7 +10,9 @@ use crate::{ config::Commitment, domain::new_verifier_domain, lookup::{LOGUP_CHALLENGE_ALPHA, LOGUP_NUM_CHALLENGES, compute_alpha_powers}, - proof::stark::{DeepPolynomialOpening, MultiProof, PolynomialOpenings}, + proof::stark::{ + BatchedMultiProof, BatchedTableData, DeepPolynomialOpening, MultiProof, PolynomialOpenings, + }, }; use crypto::{fiat_shamir::is_transcript::IsStarkTranscript, merkle_tree::proof::Proof}; #[cfg(not(feature = "test_fiat_shamir"))] @@ -73,6 +75,24 @@ where pub grinding_seed: [u8; 32], } +/// Verifier state carried across the batched (unified-shard) round-4 seam: +/// everything `batched_verify_round_4` needs that rounds 1-3 derived from the +/// transcript (per-table domains/heights, the shared OOD point `z`, the round-2 +/// constraint coefficients, and the shared LogUp challenges). Produced by +/// `batched_verify_rounds_1_to_3`; lets the continuation epoch verifier weave the +/// separate L2G lane in at the seam. +pub struct VmMidState { + pub(crate) domains: Vec>, + pub(crate) heights: Vec, + pub(crate) h_max: usize, + pub(crate) tallest: usize, + pub(crate) needs_lookup_challenges: bool, + pub(crate) lookup_challenges: Vec>, + pub(crate) boundary_coeffs_all: Vec>>, + pub(crate) transition_coeffs_all: Vec>>, + pub(crate) z: FieldElement, +} + pub type DeepPolynomialEvaluations = (Vec>, Vec>); /// The functionality of a STARK verifier providing methods to run the STARK Verify protocol @@ -1249,4 +1269,644 @@ pub trait IsStarkVerifier< true } + + /// Build a lightweight per-table `StarkProof` carrying only the fields + /// `step_2_verify_claimed_composition_polynomial` and + /// `reconstruct_deep_composition_poly_evaluation` actually read + /// (trace_length, OOD evaluations, precomputed root, bus/public inputs). All + /// commitment/opening/FRI fields are placeholders those two helpers never + /// inspect — this lets the batched verifier reuse them unchanged. + fn batched_synthetic_table_proof( + table: &BatchedTableData, + ) -> StarkProof + where + PI: Clone, + { + StarkProof { + trace_length: table.trace_length, + lde_trace_main_merkle_root: [0u8; 32], + lde_trace_aux_merkle_root: None, + lde_trace_precomputed_merkle_root: table.precomputed_root, + trace_ood_evaluations: table.trace_ood_evaluations.clone(), + composition_poly_root: [0u8; 32], + composition_poly_parts_ood_evaluation: table + .composition_poly_parts_ood_evaluation + .clone(), + fri_layers_merkle_roots: Vec::new(), + fri_final_poly_coeffs: Vec::new(), + query_list: Vec::new(), + deep_poly_openings: Vec::new(), + nonce: None, + bus_public_inputs: table.bus_public_inputs.clone(), + public_inputs: table.public_inputs.clone(), + } + } + + /// Verify a `BatchedMultiProof` (unified-shard): ONE linear transcript, ONE + /// shared OOD point z, and ONE FRI over the height-combined per-table DEEP + /// codewords, with all tables opened from three shared mixed-height MMCS + /// trees per query. Mirrors `Prover::multi_prove_batched`. + #[allow(clippy::too_many_lines)] + fn batched_multi_verify( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + ) -> bool + where + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + PI: Clone, + { + let mid = match Self::batched_verify_rounds_1_to_3(airs, proof, transcript) { + Some(m) => m, + None => return false, + }; + Self::batched_verify_round_4(mid, airs, proof, transcript, expected_bus_balance) + } + + /// Rounds 1-3 of the batched (unified-shard) verifier: replays the Fiat-Shamir + /// transcript from Phase A (preprocessed roots + the single main MMCS root) + /// through the OOD absorption, returning the derived `VmMidState` that round 4 + /// consumes. Split out of `batched_multi_verify` (behavior-preserving) so the + /// continuation epoch verifier can weave the separate L2G lane in at the seam. + /// Returns `None` on any structural rejection. + fn batched_verify_rounds_1_to_3( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut (impl IsStarkTranscript + Clone), + ) -> Option> + where + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + PI: Clone, + { + let num_tables = airs.len(); + if num_tables == 0 || num_tables != proof.per_table.len() { + return None; + } + + // Per-table lightweight domains + FRI heights (= lde_log_height). + let domains: Vec> = airs + .iter() + .zip(&proof.per_table) + .map(|(air, t)| new_verifier_domain(*air, t.trace_length)) + .collect(); + let heights: Vec = domains + .iter() + .map(|d| d.lde_length.trailing_zeros() as usize) + .collect(); + let h_max = *heights.iter().max().expect("num_tables > 0"); + // Any tallest table works: all tables at h_max share identical domain + // params (global blowup + coset_offset), so z, the FRI point and the + // query domain are the same whichever we pick. Mirrors the prover's + // `max_by_key` choice (which value it lands on is immaterial here). + let tallest = heights + .iter() + .position(|h| *h == h_max) + .expect("h_max is present"); + + let needs_lookup_challenges = airs.iter().any(|air| air.has_aux_trace()); + + // ===== Round 1 replay ===== + // Phase A: per preprocessed table, append its hardcoded precomputed root + // (checked against the AIR), then the SINGLE batched main-trace MMCS root. + for (air, t) in airs.iter().zip(&proof.per_table) { + // Soundness: composition part count is fixed by the AIR degree bound, + // not chosen by the prover. + if t.trace_length == 0 + || t.composition_poly_parts_ood_evaluation.len() + != air.composition_poly_degree_bound(t.trace_length) / t.trace_length + { + return None; + } + if air.is_preprocessed() { + let expected = air.precomputed_commitment(); + match t.precomputed_root { + Some(actual) if actual == expected => {} + _ => return None, + } + transcript.append_bytes(&expected); + } else if t.precomputed_root.is_some() { + return None; + } + } + transcript.append_bytes(&proof.main_root); + + // Bus-input presence must match the AIR layout (a dishonest prover could + // omit bus_public_inputs to bypass the balance check). + for (air, t) in airs.iter().zip(&proof.per_table) { + if air.has_trace_interaction() != t.bus_public_inputs.is_some() { + return None; + } + } + + // Phase B: shared LogUp challenges. + let lookup_challenges: Vec> = if needs_lookup_challenges { + (0..LOGUP_NUM_CHALLENGES) + .map(|_| transcript.sample_field_element()) + .collect() + } else { + Vec::new() + }; + + // Phase C: single batched aux MMCS root (present iff any table has aux). + if needs_lookup_challenges != proof.aux_root.is_some() { + return None; + } + if let Some(root) = proof.aux_root { + transcript.append_bytes(&root); + } + + // Bus contributions bind before the round-2 challenges. + for t in &proof.per_table { + if let Some(bpi) = &t.bus_public_inputs { + transcript.append_field_element(&bpi.table_contribution); + } + } + + // ===== Round 2: per-table beta -> boundary/transition coeffs, then the + // single batched composition MMCS root. ===== + let mut boundary_coeffs_all: Vec>> = + Vec::with_capacity(num_tables); + let mut transition_coeffs_all: Vec>> = + Vec::with_capacity(num_tables); + for (air, t) in airs.iter().zip(&proof.per_table) { + let beta = transcript.sample_field_element(); + let num_boundary = air + .boundary_constraints( + &t.public_inputs, + &lookup_challenges, + t.bus_public_inputs.as_ref(), + t.trace_length, + ) + .constraints + .len(); + let num_transition = air.context().num_transition_constraints; + let mut coeffs = compute_alpha_powers(&beta, num_boundary + num_transition); + let transition_coeffs: Vec<_> = coeffs.drain(..num_transition).collect(); + transition_coeffs_all.push(transition_coeffs); + boundary_coeffs_all.push(coeffs); + } + transcript.append_bytes(&proof.composition_root); + + // ===== Round 3: shared z (tallest domain), per-table OOD absorbed. ===== + let z = transcript.sample_z_ood_with_domain_params( + domains[tallest].trace_length, + domains[tallest].lde_length, + &domains[tallest].coset_offset, + ); + for t in &proof.per_table { + for col in t.trace_ood_evaluations.columns().iter() { + for elem in col.iter() { + transcript.append_field_element(elem); + } + } + for elem in t.composition_poly_parts_ood_evaluation.iter() { + transcript.append_field_element(elem); + } + } + + Some(VmMidState { + domains, + heights, + h_max, + tallest, + needs_lookup_challenges, + lookup_challenges, + boundary_coeffs_all, + transition_coeffs_all, + z, + }) + } + + /// Round 4 of the batched (unified-shard) verifier: the FRI + query phase over + /// the height-combined per-table DEEP codewords, plus the bus-balance check. + /// Split out of `batched_multi_verify` (behavior-preserving) so the + /// continuation epoch verifier can run it AFTER the L2G lane at the seam. + fn batched_verify_round_4( + mid: VmMidState, + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + ) -> bool + where + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + PI: Clone, + { + let VmMidState { + domains, + heights, + h_max, + tallest, + needs_lookup_challenges, + lookup_challenges, + boundary_coeffs_all, + transition_coeffs_all, + z, + } = mid; + let num_tables = airs.len(); + + // ===== Round 4: shared gamma, per-table DEEP coeffs, batched FRI challenges. ===== + let gamma = transcript.sample_field_element(); + let mut trace_term_coeffs_all: Vec>>> = + Vec::with_capacity(num_tables); + let mut gammas_all: Vec>> = Vec::with_capacity(num_tables); + for (air, t) in airs.iter().zip(&proof.per_table) { + let num_terms_comp = t.composition_poly_parts_ood_evaluation.len(); + let num_terms_trace = air.context().transition_offsets.len() + * air.step_size() + * air.context().trace_columns; + let mut coeffs: Vec<_> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * &gamma)) + .take(num_terms_comp + num_terms_trace) + .collect(); + let trace_term_coeffs: Vec<_> = coeffs + .drain(..num_terms_trace) + .collect::>() + .chunks(air.context().transition_offsets.len() * air.step_size()) + .map(|c| c.to_vec()) + .collect(); + trace_term_coeffs_all.push(trace_term_coeffs); + gammas_all.push(coeffs); + } + + let grinding_factor = airs[0].context().proof_options.grinding_factor; + let num_queries = airs[0].options().fri_number_of_queries; + let fri_domain_size = 1usize << h_max; + let fri_challenges = derive_batched_fri_challenges( + transcript, + &heights, + &proof.fri_layers_merkle_roots, + &proof.fri_last_value, + grinding_factor, + proof.nonce, + num_queries, + fri_domain_size, + ); + let alpha = fri_challenges.alpha; + let betas_fri = fri_challenges.betas; + let iotas = fri_challenges.iotas; + + // Grinding. + if grinding_factor > 0 { + let ok = proof.nonce.is_some_and(|n| { + grinding::is_valid_nonce(&fri_challenges.grinding_seed, n, grinding_factor) + }); + if !ok { + return false; + } + } + + if proof.query_list.len() < num_queries || proof.deep_poly_openings.len() < num_queries { + return false; + } + + // Per-table synthetic proofs + Challenges (reused by step 2 and the query loop). + let synth_proofs: Vec> = proof + .per_table + .iter() + .map(Self::batched_synthetic_table_proof) + .collect(); + let table_challenges: Vec> = (0..num_tables) + .map(|i| Challenges { + z: z.clone(), + boundary_coeffs: boundary_coeffs_all[i].clone(), + transition_coeffs: transition_coeffs_all[i].clone(), + trace_term_coeffs: trace_term_coeffs_all[i].clone(), + gammas: gammas_all[i].clone(), + zetas: Vec::new(), + iotas: Vec::new(), + rap_challenges: lookup_challenges.clone(), + grinding_seed: [0u8; 32], + }) + .collect(); + + // ===== Step 2 (claimed composition polynomial) per table. ===== + for i in 0..num_tables { + if !Self::step_2_verify_claimed_composition_polynomial( + airs[i], + &synth_proofs[i], + &domains[i], + &table_challenges[i], + ) { + return false; + } + } + + // MMCS binding data (all public / from the AIRs). + // Committed main-split width per table = full main columns minus the + // precomputed prefix. `context().trace_columns` counts every committed + // trace column (main + aux), so subtracting the aux and precomputed + // counts yields the main-split width. All three are AIR-intrinsic (not + // proof-supplied), so this binds the MMCS leaf boundaries independently + // of the prover. NB: `trace_layout().0` is NOT usable here — for + // step-packed AIRs (e.g. BitFlags) it is a logical layout figure, not + // the physical column count. + let main_widths: Vec = airs + .iter() + .map(|a| { + a.context().trace_columns + - a.num_auxiliary_rap_columns() + - a.num_precomputed_columns() + }) + .collect(); + let comp_widths: Vec = proof + .per_table + .iter() + .map(|t| t.composition_poly_parts_ood_evaluation.len()) + .collect(); + let aux_indices: Vec = (0..num_tables) + .filter(|&i| airs[i].has_aux_trace()) + .collect(); + let aux_heights: Vec = aux_indices.iter().map(|&i| heights[i]).collect(); + let aux_widths: Vec = aux_indices + .iter() + .map(|&i| airs[i].num_auxiliary_rap_columns()) + .collect(); + let precomputed_indices: Vec = (0..num_tables) + .filter(|&i| airs[i].is_preprocessed()) + .collect(); + + // alpha^i powers for the cross-table combination. + let mut alpha_pows: Vec> = Vec::with_capacity(num_tables); + { + let mut cur = FieldElement::::one(); + for _ in 0..num_tables { + alpha_pows.push(cur.clone()); + cur = &cur * α + } + } + let num_layers = proof.fri_layers_merkle_roots.len(); + + // ===== Per query: MMCS openings, DEEP reconstruction, fold-and-inject FRI. ===== + for (q, &iota) in iotas.iter().enumerate() { + let qo = &proof.deep_poly_openings[q]; + + // 1) Authenticate the shared per-phase openings. + if !MixedMmcs::::verify_batch( + &proof.main_root, + iota, + &qo.main, + &heights, + &main_widths, + ) { + return false; + } + match (&proof.aux_root, &qo.aux) { + (Some(root), Some(aux_op)) => { + if !MixedMmcs::::verify_batch( + root, + iota, + aux_op, + &aux_heights, + &aux_widths, + ) { + return false; + } + } + (None, None) => {} + _ => return false, + } + if !MixedMmcs::::verify_batch( + &proof.composition_root, + iota, + &qo.composition, + &heights, + &comp_widths, + ) { + return false; + } + + // Precomputed openings (one per preprocessed table, in that order). + if qo.precomputed.len() != precomputed_indices.len() { + return false; + } + for (pc, &ti) in precomputed_indices.iter().enumerate() { + let root = match proof.per_table[ti].precomputed_root { + Some(r) => r, + None => return false, + }; + let local = iota >> (h_max - heights[ti]); + if !Self::verify_opening_pair::(&qo.precomputed[pc], &root, local) { + return false; + } + } + + // 2) Reconstruct each table's DEEP value at its opened row pair. + let mut deep_primary = vec![FieldElement::::zero(); num_tables]; + let mut deep_sym = vec![FieldElement::::zero(); num_tables]; + for i in 0..num_tables { + let leaf = iota >> (h_max - heights[i]); + let main_op = &qo.main.per_matrix[i]; + let comp_op = &qo.composition.per_matrix[i]; + let precomp_op = precomputed_indices + .iter() + .position(|&x| x == i) + .map(|pc| &qo.precomputed[pc]); + let aux_op = aux_indices + .iter() + .position(|&x| x == i) + .and_then(|ai| qo.aux.as_ref().map(|a| &a.per_matrix[ai])); + + let mut base_p: Vec> = Vec::new(); + let mut base_s: Vec> = Vec::new(); + if let Some(p) = precomp_op { + base_p.extend_from_slice(&p.evaluations); + base_s.extend_from_slice(&p.evaluations_sym); + } + base_p.extend_from_slice(&main_op.evaluations); + base_s.extend_from_slice(&main_op.evaluations_sym); + let aux_p: &[FieldElement] = + aux_op.map(|a| a.evaluations.as_slice()).unwrap_or(&[]); + let aux_s: &[FieldElement] = + aux_op.map(|a| a.evaluations_sym.as_slice()).unwrap_or(&[]); + + let prim_root = &domains[i].trace_primitive_root; + let ep_p = domains[i] + .lde_coset_element(reverse_index(leaf * 2, domains[i].lde_length as u64)); + let ep_s = domains[i] + .lde_coset_element(reverse_index(leaf * 2 + 1, domains[i].lde_length as u64)); + + deep_primary[i] = match Self::reconstruct_deep_composition_poly_evaluation( + &synth_proofs[i], + &ep_p, + prim_root, + &table_challenges[i], + &base_p, + aux_p, + &comp_op.evaluations, + ) { + Some(v) => v, + None => return false, + }; + deep_sym[i] = match Self::reconstruct_deep_composition_poly_evaluation( + &synth_proofs[i], + &ep_s, + prim_root, + &table_challenges[i], + &base_s, + aux_s, + &comp_op.evaluations_sym, + ) { + Some(v) => v, + None => return false, + }; + } + + // combined[h] at a codeword position selected by `bit` (0 -> primary + // row, 1 -> symmetric row): Sum over tables at height h of alpha^i * deep_i. + let combined_at = |h: usize, bit: usize| -> FieldElement { + let mut acc = FieldElement::::zero(); + for i in 0..num_tables { + if heights[i] == h { + let d = if bit == 0 { + &deep_primary[i] + } else { + &deep_sym[i] + }; + acc += &alpha_pows[i] * d; + } + } + acc + }; + + // 3) Fold-and-inject FRI (inverse of `batched_commit_phase`). + let c_hmax = combined_at(h_max, 0); + let c_hmax_sym = combined_at(h_max, 1); + + let ep0 = domains[tallest] + .lde_coset_element(reverse_index(iota * 2, domains[tallest].lde_length as u64)); + let ep0_inv = match ep0.inv() { + Ok(v) => v, + Err(_) => return false, + }; + + // Initial fold of the (uncommitted) tallest layer with betas_fri[0]. + let mut v = + (&c_hmax + &c_hmax_sym) + ep0_inv.clone() * &betas_fri[0] * (&c_hmax - &c_hmax_sym); + let mut index = iota; + let mut point_inv = ep0_inv.square(); + + let fri_deco = &proof.query_list[q]; + if fri_deco.layers_auth_paths.len() != num_layers + || fri_deco.layers_evaluations_sym.len() != num_layers + { + return false; + } + + let mut fold_ok = true; + for iter in 0..num_layers { + let h = h_max - 1 - iter; + // Inject the tables entering at this height (adds zero if none). + let inj = combined_at(h, index & 1); + v += betas_fri[iter].square() * inj; + + let eval_sym = &fri_deco.layers_evaluations_sym[iter]; + fold_ok &= Self::verify_fri_layer_openings( + &proof.fri_layers_merkle_roots[iter], + &fri_deco.layers_auth_paths[iter], + &v, + eval_sym, + index, + ); + + v = (&v + eval_sym) + point_inv.clone() * &betas_fri[iter + 1] * (&v - eval_sym); + index >>= 1; + point_inv = point_inv.square(); + } + if !fold_ok || v != proof.fri_last_value { + return false; + } + } + + // ===== Bus balance. ===== + if needs_lookup_challenges { + let mut total = FieldElement::::zero(); + for (air, t) in airs.iter().zip(&proof.per_table) { + if air.has_trace_interaction() + && let Some(bpi) = &t.bus_public_inputs + { + total = total + &bpi.table_contribution; + } + } + if total != *expected_bus_balance { + return false; + } + } + + true + } + + /// Continuation epoch verifier: verify the epoch's VM tables with the batched + /// (unified-shard) FRI while verifying the single L2G sub-table as a SEPARATE + /// commitment lane. Mirrors `IsStarkProver::multi_prove_batched_epoch`'s + /// transcript order exactly: + /// + /// 1. Absorb the L2G main root FIRST. + /// 2. `batched_verify_rounds_1_to_3` over the VM tables (to the round-4 seam). + /// 3. At the seam, FORK the transcript (single lane -> no idx bytes; then absorb + /// the L2G aux root, then the L2G bus `table_contribution`) and verify the + /// L2G lane via `verify_rounds_2_to_4` with the shared LogUp challenges. + /// 4. `batched_verify_round_4` for the VM tables on the main transcript. + fn batched_verify_epoch( + vm_refs: &[&dyn AIR], + l2g_ref: &dyn AIR, + vm_proof: &BatchedMultiProof, + l2g_proof: &StarkProof, + transcript: &mut (impl IsStarkTranscript + Clone), + expected_bus_balance: &FieldElement, + ) -> bool + where + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + PI: Clone, + { + // (1) Mirror the prover: absorb the L2G main root FIRST (canonical order). + transcript.append_bytes(&l2g_proof.lde_trace_main_merkle_root); + + // (2) VM rounds 1-3 to the round-4 seam. + let mid = match Self::batched_verify_rounds_1_to_3(vm_refs, vm_proof, transcript) { + Some(m) => m, + None => return false, + }; + + // (3) L2G lane on a fork of the post-seam transcript. Single lane -> no idx + // bytes; absorb the aux root then the bus table_contribution (matches the + // prover's fork in `multi_prove_batched_epoch`). + let mut l2g_fork = transcript.clone(); + if let Some(aux_root) = l2g_proof.lde_trace_aux_merkle_root { + l2g_fork.append_bytes(&aux_root); + } + if let Some(bpi) = l2g_proof.bus_public_inputs.as_ref() { + l2g_fork.append_field_element(&bpi.table_contribution); + } + let l2g_ok = Self::verify_rounds_2_to_4( + l2g_ref, + l2g_proof, + &mut l2g_fork, + mid.lookup_challenges.clone(), + ); + + // (4) VM batched Round 4 continues on the main (un-cloned) transcript. + // + // Bus balance: L2G shares the in-trace Memory / range-check buses with the + // VM tables. The monolithic check summed table_contribution over VM + L2G + // against the COMMIT offset; batched_verify_round_4 sums only the VM lane, + // so fold L2G's contribution into the target: + // Sum_VM table_contribution == expected - L2G_contribution + // i.e. Sum_VM + L2G == expected. L2G's table_contribution is bound to its + // committed trace by the L2G proof verified above, so this stays sound. + let mut vm_expected = expected_bus_balance.clone(); + if l2g_ref.has_trace_interaction() + && let Some(bpi) = l2g_proof.bus_public_inputs.as_ref() + { + vm_expected = &vm_expected - &bpi.table_contribution; + } + let vm_ok = Self::batched_verify_round_4(mid, vm_refs, vm_proof, transcript, &vm_expected); + + l2g_ok && vm_ok + } } diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 54a3ee583..f8a25e371 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -54,7 +54,7 @@ use stark::config::Commitment; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet, EmptyConstraints}; use stark::lookup::{AirWithBuses, AuxiliaryTraceBuildData, NullBoundaryConstraintBuilder}; use stark::proof::options::ProofOptions; -use stark::proof::stark::MultiProof; +use stark::proof::stark::{BatchedMultiProof, MultiProof, StarkProof}; use stark::prover::{IsStarkProver, Prover}; use stark::trace::TraceTable; use stark::traits::AIR; @@ -69,7 +69,7 @@ use crate::tables::types::{GoldilocksExtension, GoldilocksField}; use crate::tables::{MaxRowsConfig, global_memory}; use crate::{ Error, FIXED_TABLE_COUNT, RuntimePageRange, TableCounts, VmAirs, - compute_expected_commit_bus_balance, verify_l2g_commitment_binding, + compute_expected_commit_bus_balance_batched, verify_l2g_commitment_binding, }; type F = GoldilocksField; @@ -343,8 +343,11 @@ struct EpochStart<'a> { /// tables rather than trusting any prover-supplied page config. #[derive(serde::Serialize, serde::Deserialize)] struct EpochProof { - /// The epoch's STARK proof (its tables + the epoch-local L2G sub-table last). - proof: MultiProof, + /// Batched proof over the epoch's VM tables (shared MMCS + unified FRI). + vm_proof: BatchedMultiProof, + /// Standalone proof over the single L2G sub-table (its own tree + own FRI). + /// Its `lde_trace_main_merkle_root` is the binding root (== `l2g_root`). + l2g_proof: StarkProof, /// Bytes this epoch committed — the COMMIT-bus receiver reference. public_output: Vec, /// Statement values the epoch transcript is seeded with (re-derived on verify). @@ -500,26 +503,25 @@ fn prove_epoch( // roots). It is appended to the proof below, not through `air_trace_pairs`. let mut l2g_trace = local_to_global::generate_local_to_global_trace(boundary); - let mut pairs = airs.air_trace_pairs(&mut traces); - pairs.push((&l2g_air, &mut l2g_trace, &())); - let proof = Prover::multi_prove( - pairs, + // VM tables ride the batched (unified-shard) FRI; the single L2G sub-table is a + // SEPARATE commitment lane (its own tree + own FRI) so the L2G<->global root + // binding still holds. Do NOT push L2G into the VM pairs. + let vm_pairs = airs.air_trace_pairs(&mut traces); + let l2g_pair = (&l2g_air as AirRef, &mut l2g_trace, &()); + let (vm_proof, l2g_proof) = Prover::multi_prove_batched_epoch( + vm_pairs, + l2g_pair, &mut seed(), #[cfg(feature = "disk-spill")] stark::storage_mode::StorageMode::Ram, ) .map_err(|e| Error::Prover(format!("{e:?}")))?; - let l2g_root = proof - .proofs - .last() - .ok_or_else(|| { - Error::ContinuationInvariant("epoch proof is missing the L2G sub-table".to_string()) - })? - .lde_trace_main_merkle_root; + let l2g_root = l2g_proof.lde_trace_main_merkle_root; Ok(EpochProof { - proof, + vm_proof, + l2g_proof, public_output, table_counts, runtime_page_ranges, @@ -558,8 +560,10 @@ fn verify_epoch( } else { FIXED_TABLE_COUNT - 1 }; - let expected_proof_count = epoch.table_counts.total() + fixed_tables + 1; - if expected_proof_count != epoch.proof.proofs.len() { + // VM tables ride the batched lane (counted in vm_proof.per_table); the single + // L2G sub-table is a SEPARATE lane (l2g_proof), no longer a +1 here. + let expected_proof_count = epoch.table_counts.total() + fixed_tables; + if expected_proof_count != epoch.vm_proof.per_table.len() { return false; } @@ -573,8 +577,9 @@ fn verify_epoch( is_final, ); let l2g_air = l2g_memory_air(opts, label); - let mut refs = airs.air_refs(); - refs.push(&l2g_air); + // VM tables only — L2G is verified as its own lane, NOT pushed into vm_refs. + let vm_refs = airs.air_refs(); + let l2g_ref: AirRef = &l2g_air; let seed = || { epoch_transcript( @@ -594,29 +599,39 @@ fn verify_epoch( .copied() .unwrap_or(0) as u64; - let expected = match compute_expected_commit_bus_balance( - &refs, - &epoch.proof, + let expected = match compute_expected_commit_bus_balance_batched( + &vm_refs, + &epoch.vm_proof, &epoch.public_output, commit_start_index, + // The epoch transcript absorbs the L2G main root FIRST (mirrors + // `multi_prove_batched_epoch`) so the sampled LogUp `(z, alpha)` — and the + // expected COMMIT-bus offset — match the prover's. Empty output masks a + // mismatch (offset is zero regardless of the challenges). + Some(&epoch.l2g_root), &mut seed(), ) { Some(expected) => expected, None => return false, }; - if !Verifier::multi_verify(&refs, &epoch.proof, &mut seed(), &expected) { + // Two-lane epoch verify: batched VM FRI + standalone L2G lane, woven through + // ONE transcript (L2G main root first, VM rounds 1-3, fork for L2G at the seam, + // then VM round 4). The expected COMMIT balance is shared across both lanes. + if !Verifier::batched_verify_epoch( + &vm_refs, + l2g_ref, + &epoch.vm_proof, + &epoch.l2g_proof, + &mut seed(), + &expected, + ) { return false; } - // The claimed L2G root must be the one this proof actually committed (it is what - // verify_l2g_commitment_binding later ties to the global proof). - epoch - .proof - .proofs - .last() - .map(|p| p.lde_trace_main_merkle_root) - == Some(epoch.l2g_root) + // The claimed L2G root must be the one the L2G lane actually committed (it is + // what verify_l2g_commitment_binding later ties to the global proof). + epoch.l2g_proof.lde_trace_main_merkle_root == epoch.l2g_root } /// Build the cross-epoch global memory proof: every epoch's L2G sub-table on the diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 9c315faf1..eab944c59 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -62,7 +62,7 @@ use crate::test_utils::{ // Re-exported so downstream verifier guests (e.g. the in-VM recursion guest) can // name the proof-options type carried in their private input alongside `VmProof`. pub use stark::proof::options::{GoldilocksCubicProofOptions, ProofOptions}; -use stark::proof::stark::MultiProof; +use stark::proof::stark::{BatchedMultiProof, MultiProof}; /// A run-length encoded range of contiguous zero-initialized 4KB pages. /// @@ -157,8 +157,8 @@ impl TableCounts { /// needed by the verifier to reconstruct the AIR configuration. #[derive(Debug, serde::Serialize, serde::Deserialize)] pub struct VmProof { - /// The multi-table STARK proof. - pub proof: MultiProof, + /// The multi-table STARK proof (unified-shard / batched MMCS). + pub proof: BatchedMultiProof, /// Run-length encoded runtime page ranges. /// These are zero-initialized pages accessed during execution but not /// covered by ELF segments (stack, heap, etc.). @@ -663,6 +663,10 @@ impl VmAirs { /// LogUp challenges (z, alpha). Creates a fresh transcript, appends all main /// trace commitments in the same order as the prover, then samples two /// challenge elements. +/// +/// Only the batched analogue is used in the production path now; this +/// non-batched replay is retained for the monolithic-style test helpers. +#[cfg(test)] pub(crate) fn replay_transcript_phase_a( airs: &[&dyn AIR], multi_proof: &MultiProof, @@ -733,6 +737,10 @@ pub(crate) fn compute_commit_bus_offset( /// Replays Phase A of the transcript to recover (z, alpha), then computes /// the offset from the given public output bytes. Call this after `multi_prove` /// and before `multi_verify`. +/// +/// Superseded by [`compute_expected_commit_bus_balance_batched`] in the +/// production path; retained for the monolithic-style test helpers. +#[cfg(test)] pub(crate) fn compute_expected_commit_bus_balance( airs: &[&dyn AIR], proof: &MultiProof, @@ -744,6 +752,47 @@ pub(crate) fn compute_expected_commit_bus_balance( compute_commit_bus_offset(public_output_bytes, start_index, &z, &alpha) } +/// Batched (unified-shard) analogue of [`replay_transcript_phase_a`]: appends +/// each preprocessed table's hardcoded precomputed root and the SINGLE batched +/// main MMCS root (Phase A of the linear transcript), then samples the shared +/// LogUp `(z, alpha)`. Mirrors `Prover::prove_rounds_1_to_3` Phase A + B. +pub(crate) fn replay_transcript_phase_a_batched( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + l2g_main_root: Option<&Commitment>, + transcript: &mut DefaultTranscript, +) -> (FieldElement, FieldElement) { + // Continuation epochs absorb the standalone L2G main root FIRST (mirrors + // `Prover::multi_prove_batched_epoch`, which appends it before the VM roots); + // monolithic proofs have no separate L2G lane and pass `None`. + if let Some(root) = l2g_main_root { + transcript.append_bytes(root); + } + for air in airs.iter() { + if air.is_preprocessed() { + transcript.append_bytes(&air.precomputed_commitment()); + } + } + transcript.append_bytes(&proof.main_root); + let z: FieldElement = transcript.sample_field_element(); + let alpha: FieldElement = transcript.sample_field_element(); + (z, alpha) +} + +/// Batched analogue of [`compute_expected_commit_bus_balance`] for a +/// [`BatchedMultiProof`]. +pub(crate) fn compute_expected_commit_bus_balance_batched( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + public_output_bytes: &[u8], + start_index: u64, + l2g_main_root: Option<&Commitment>, + transcript: &mut DefaultTranscript, +) -> Option> { + let (z, alpha) = replay_transcript_phase_a_batched(airs, proof, l2g_main_root, transcript); + compute_commit_bus_offset(public_output_bytes, start_index, &z, &alpha) +} + /// Bind the final cross-epoch GlobalMemory proof to the per-epoch proofs. /// /// The final proof commits one local-to-global sub-table per epoch as its first @@ -941,10 +990,10 @@ pub fn prove_with_options_and_inputs( proof_options.fri_final_poly_log_degree, ); - // Phase 4: Prove (multi_prove) + // Phase 4: Prove (unified-shard batched MMCS + single FRI) #[cfg(feature = "instruments")] let __sp = stark::instruments::span("proving"); - let proof = Prover::multi_prove( + let proof = Prover::multi_prove_batched( airs.air_trace_pairs(&mut traces), &mut transcript, #[cfg(feature = "disk-spill")] @@ -1064,13 +1113,13 @@ pub fn verify_with_options( // FIXED_TABLE_COUNT always-present tables, plus page tables. let expected_proof_count = vm_proof.table_counts.total() + FIXED_TABLE_COUNT + page_configs.len(); - if expected_proof_count != vm_proof.proof.proofs.len() { + if expected_proof_count != vm_proof.proof.per_table.len() { return Err(Error::InvalidTableCounts(format!( "table_counts total ({}) + {FIXED_TABLE_COUNT} fixed + {} pages = {}, but proof contains {} sub-proofs", vm_proof.table_counts.total(), page_configs.len(), expected_proof_count, - vm_proof.proof.proofs.len(), + vm_proof.proof.per_table.len(), ))); } @@ -1111,12 +1160,14 @@ pub fn verify_with_options( // independently of the multi_verify transcript, but both must start from // the same statement-bound state. let mut transcript_for_replay = transcript.clone(); - let expected_bus_balance = match compute_expected_commit_bus_balance( + let expected_bus_balance = match compute_expected_commit_bus_balance_batched( &air_refs, &vm_proof.proof, &vm_proof.public_output, // Monolithic proof: commits are indexed from 0. 0, + // No standalone L2G lane in a monolithic proof. + None, &mut transcript_for_replay, ) { Some(balance) => balance, @@ -1126,7 +1177,7 @@ pub fn verify_with_options( stark::profile_markers::step_marker::<{ stark::profile_markers::STEP_AIRS_AND_BUS_BALANCE_DONE }>( ); - Ok(Verifier::multi_verify( + Ok(Verifier::batched_multi_verify( &air_refs, &vm_proof.proof, &mut transcript, diff --git a/prover/src/test_utils.rs b/prover/src/test_utils.rs index 8ed5fdca2..8e24315be 100644 --- a/prover/src/test_utils.rs +++ b/prover/src/test_utils.rs @@ -26,7 +26,7 @@ use stark::lookup::{ AirWithBuses, AuxiliaryTraceBuildData, BusInteraction, BusValue, NullBoundaryConstraintBuilder, }; use stark::proof::options::ProofOptions; -use stark::proof::stark::MultiProof; +use stark::proof::stark::{BatchedMultiProof, MultiProof}; use stark::prover::{IsStarkProver, Prover, ProvingError}; #[cfg(feature = "disk-spill")] use stark::storage_mode::StorageMode; @@ -145,6 +145,21 @@ where ) } +pub fn multi_prove_batched_ram( + air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), +) -> Result, ProvingError> +where + PI: Send + Sync + Clone, +{ + Prover::::multi_prove_batched( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + StorageMode::Ram, + ) +} + // ============================================================================= // Soundness regression helpers (negative AIR tests) // ============================================================================= diff --git a/prover/src/tests/bitwise_tests.rs b/prover/src/tests/bitwise_tests.rs index c824764d3..38a1c0f1a 100644 --- a/prover/src/tests/bitwise_tests.rs +++ b/prover/src/tests/bitwise_tests.rs @@ -5,7 +5,7 @@ use crate::tables::bitwise::{ generate_bitwise_trace, is_preprocessed, preprocessed_commitment, row_index, }; use crate::tables::types::{BusId, FE}; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::{multi_prove_batched_ram, multi_prove_ram}; use math::field::element::FieldElement; use stark::constraints::builder::EmptyConstraints; use stark::lookup::Multiplicity; @@ -626,12 +626,13 @@ mod soundness_tests { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])) + .unwrap(); let airs: Vec<&dyn AIR> = vec![&sender_air, &receiver_air]; - let result = Verifier::multi_verify( + let result = Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/prover/src/tests/lt_bus_tests.rs b/prover/src/tests/lt_bus_tests.rs index e95a81285..8bed0083c 100644 --- a/prover/src/tests/lt_bus_tests.rs +++ b/prover/src/tests/lt_bus_tests.rs @@ -23,7 +23,7 @@ use stark::verifier::{IsStarkVerifier, Verifier}; use crate::tables::lt::{LtOperation, cols, generate_lt_trace}; use crate::tables::types::{BusId, FE, GoldilocksExtension, GoldilocksField}; -use crate::test_utils::multi_prove_ram; +use crate::test_utils::{multi_prove_batched_ram, multi_prove_ram}; type F = GoldilocksField; type E = GoldilocksExtension; @@ -289,12 +289,12 @@ fn prove_and_verify(ops: &[LtOperation]) -> bool { ]; let multi_proof = - multi_prove_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); + multi_prove_batched_ram(air_trace_pairs, &mut DefaultTranscript::::new(&[])).unwrap(); let airs: Vec<&dyn AIR> = vec![&sender_air, &receiver_air]; - Verifier::multi_verify( + Verifier::batched_multi_verify( &airs, &multi_proof, &mut DefaultTranscript::::new(&[]), diff --git a/prover/src/tests/prove_elfs_tests.rs b/prover/src/tests/prove_elfs_tests.rs index 527c092c5..15b704801 100644 --- a/prover/src/tests/prove_elfs_tests.rs +++ b/prover/src/tests/prove_elfs_tests.rs @@ -31,6 +31,7 @@ use executor::vm::execution::Executor; // Import shared utilities use crate::VmAirs; +use crate::test_utils::multi_prove_batched_ram; use crate::test_utils::multi_prove_ram; use crate::test_utils::run_asm_elf; @@ -120,7 +121,7 @@ fn prove_vm_minimal(elf_bytes: &[u8], private_inputs: &[u8], max_rows: &MaxRowsC None, ); let runtime_page_ranges = traces.runtime_page_ranges(); - let proof = multi_prove_ram( + let proof = multi_prove_batched_ram( airs.air_trace_pairs(&mut traces), &mut DefaultTranscript::::new(&[]), ) @@ -164,15 +165,16 @@ fn verify_vm_minimal(vm_proof: &VmProof, elf_bytes: &[u8]) -> bool { ); let air_refs = airs.air_refs(); let mut replay_transcript = DefaultTranscript::::new(&[]); - let expected_bus_balance = crate::compute_expected_commit_bus_balance( + let expected_bus_balance = crate::compute_expected_commit_bus_balance_batched( &air_refs, &vm_proof.proof, &vm_proof.public_output, 0, + None, &mut replay_transcript, ) .expect("fingerprint collision in test"); - Verifier::multi_verify( + Verifier::batched_multi_verify( &air_refs, &vm_proof.proof, &mut DefaultTranscript::::new(&[]),