From 045b5624836ec528d34441293eca3a18a61578f7 Mon Sep 17 00:00:00 2001 From: Nicole Date: Thu, 23 Apr 2026 17:20:38 -0300 Subject: [PATCH 1/3] deduplicate Domain and LdeTwiddles in multi_prove --- crypto/stark/src/prover.rs | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 41ccb8366..ecfdb8c2c 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -645,8 +645,8 @@ pub trait IsStarkProver< fn run_debug_checks( air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>], commitments: &[Round1Commitments], - domains: &[Domain], - twiddle_caches: &[LdeTwiddles], + domains: &[Arc>], + twiddle_caches: &[Arc>], ) where FieldElement: AsBytes, FieldElement: AsBytes, @@ -1523,13 +1523,30 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let phase_start = Instant::now(); + // 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); + 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 domain = new_domain(*air, trace_length); - let twiddles = LdeTwiddles::new(&domain); + let blowup = air.options().blowup_factor as usize; + let coset_offset = air.options().coset_offset; + let key = (trace_length, blowup, coset_offset); + + let (domain, twiddles) = domain_cache + .entry(key) + .or_insert_with(|| { + let d = new_domain(*air, trace_length); + let t = LdeTwiddles::new(&d); + (Arc::new(d), Arc::new(t)) + }) + .clone(); domains.push(domain); twiddle_caches.push(twiddles); From d2b66fed87fdaf99468b7dd72d47e49cc0516ba9 Mon Sep 17 00:00:00 2001 From: Nicole Date: Fri, 24 Apr 2026 11:35:33 -0300 Subject: [PATCH 2/3] free domain cache after pre-pass and add test --- crypto/stark/src/prover.rs | 4 ++ crypto/stark/src/tests/prover_tests.rs | 77 +++++++++++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index ecfdb8c2c..a0fface3a 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1551,6 +1551,10 @@ pub trait IsStarkProver< 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); diff --git a/crypto/stark/src/tests/prover_tests.rs b/crypto/stark/src/tests/prover_tests.rs index 55d58da7e..a34692605 100644 --- a/crypto/stark/src/tests/prover_tests.rs +++ b/crypto/stark/src/tests/prover_tests.rs @@ -1,10 +1,16 @@ +use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use crate::{ domain::Domain, - examples::{quadratic_air::QuadraticAIR, simple_fibonacci}, + examples::{ + quadratic_air::QuadraticAIR, + simple_fibonacci::{self, FibonacciAIR, FibonacciPublicInputs}, + }, proof::options::ProofOptions, prover::{IsStarkProver, Prover, evaluate_polynomial_on_lde_domain}, trace::{LDETraceTable, get_trace_evaluations, get_trace_evaluations_from_lde}, traits::AIR, + verifier::{IsStarkVerifier, Verifier}, }; use math::{ field::{element::FieldElement, goldilocks::GoldilocksField, traits::IsFFTField}, @@ -233,3 +239,72 @@ fn test_decompose_and_extend_d2_matches_original() { assert_eq!(new_result[1][i], original[1][i], "H₁ mismatch at index {i}"); } } + +/// Test that the domain cache 3-tuple key `(trace_length, blowup, coset_offset)` correctly +/// distinguishes AIRs that share the same `(trace_length, blowup)` but differ in +/// `coset_offset`. Both AIRs must get their own `Domain` and the resulting proofs must +/// verify successfully. +#[test_log::test] +fn test_multi_prove_mixed_coset_offsets() { + let proof_options_3 = ProofOptions { + blowup_factor: 2, + fri_number_of_queries: 3, + coset_offset: 3, + grinding_factor: 1, + }; + let proof_options_7 = ProofOptions { + blowup_factor: 2, + fri_number_of_queries: 3, + coset_offset: 7, + grinding_factor: 1, + }; + + // Both AIRs have the same trace length and blowup, but different coset offsets. + 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)], 8); + + let pub_inputs = FibonacciPublicInputs { + a0: Felt::one(), + a1: Felt::one(), + }; + + let air_1 = FibonacciAIR::::new(&proof_options_3); + let air_2 = FibonacciAIR::::new(&proof_options_7); + + 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 multi_proof = Prover::multi_prove( + air_trace_pairs, + &mut DefaultTranscript::::new(&[]), + ) + .expect("proving should succeed"); + + let airs: Vec< + &dyn AIR< + Field = GoldilocksField, + FieldExtension = GoldilocksField, + PublicInputs = FibonacciPublicInputs, + >, + > = vec![&air_1, &air_2]; + + assert!( + Verifier::multi_verify( + &airs, + &multi_proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "verification should succeed when AIRs share (trace_length, blowup) but differ in coset_offset" + ); +} From 31a2bd4b40bf22ce94ea2f4343c9c54709b49731 Mon Sep 17 00:00:00 2001 From: Nicole Date: Fri, 24 Apr 2026 12:02:10 -0300 Subject: [PATCH 3/3] Cover the Domain/LdeTwiddles dedup hit path in multi_prove with a regression test --- crypto/stark/src/prover.rs | 34 +++++++++++ crypto/stark/src/tests/prover_tests.rs | 78 +++++++++++++++++++++++++- 2 files changed, 111 insertions(+), 1 deletion(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index a0fface3a..43086d4fa 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -48,6 +48,34 @@ type AirTracePair<'a, Field, FieldExtension, PI> = ( &'a PI, ); +#[cfg(test)] +pub(crate) mod domain_cache_stats { + use std::cell::Cell; + + thread_local! { + static COUNTS: Cell<(usize, usize)> = const { Cell::new((0, 0)) }; + } + + pub(crate) fn reset() { + COUNTS.with(|c| c.set((0, 0))); + } + + pub(crate) fn get() -> (usize, usize) { + COUNTS.with(Cell::get) + } + + pub(crate) fn record(was_hit: bool) { + COUNTS.with(|c| { + let (hits, misses) = c.get(); + c.set(if was_hit { + (hits + 1, misses) + } else { + (hits, misses + 1) + }); + }); + } +} + /// A default STARK prover implementing `IsStarkProver`. pub struct Prover< Field: IsSubFieldOf + IsFFTField + Send + Sync, @@ -1539,6 +1567,9 @@ pub trait IsStarkProver< 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(|| { @@ -1548,6 +1579,9 @@ pub trait IsStarkProver< }) .clone(); + #[cfg(test)] + domain_cache_stats::record(was_hit); + domains.push(domain); twiddle_caches.push(twiddles); } diff --git a/crypto/stark/src/tests/prover_tests.rs b/crypto/stark/src/tests/prover_tests.rs index a34692605..90cace8cf 100644 --- a/crypto/stark/src/tests/prover_tests.rs +++ b/crypto/stark/src/tests/prover_tests.rs @@ -7,7 +7,7 @@ use crate::{ simple_fibonacci::{self, FibonacciAIR, FibonacciPublicInputs}, }, proof::options::ProofOptions, - prover::{IsStarkProver, Prover, evaluate_polynomial_on_lde_domain}, + prover::{IsStarkProver, Prover, domain_cache_stats, evaluate_polynomial_on_lde_domain}, trace::{LDETraceTable, get_trace_evaluations, get_trace_evaluations_from_lde}, traits::AIR, verifier::{IsStarkVerifier, Verifier}, @@ -308,3 +308,79 @@ fn test_multi_prove_mixed_coset_offsets() { "verification should succeed when AIRs share (trace_length, blowup) but differ in coset_offset" ); } + +/// 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. +#[test_log::test] +fn test_multi_prove_dedups_shared_domain_params() { + domain_cache_stats::reset(); + + let proof_options = ProofOptions { + blowup_factor: 2, + fri_number_of_queries: 3, + coset_offset: 3, + grinding_factor: 1, + }; + + 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)], 8); + let mut trace_3 = simple_fibonacci::fibonacci_trace([Felt::from(1), Felt::from(1)], 8); + + let pub_inputs = FibonacciPublicInputs { + a0: Felt::one(), + a1: Felt::one(), + }; + + let air_1 = FibonacciAIR::::new(&proof_options); + let air_2 = FibonacciAIR::::new(&proof_options); + let air_3 = 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), + (&air_3, &mut trace_3, &pub_inputs), + ]; + + let multi_proof = Prover::multi_prove( + air_trace_pairs, + &mut DefaultTranscript::::new(&[]), + ) + .expect("proving should succeed"); + + let (hits, misses) = domain_cache_stats::get(); + assert_eq!( + misses, 1, + "only one Domain/LdeTwiddles must be constructed for 3 AIRs sharing domain params" + ); + assert_eq!( + hits, 2, + "remaining 2 AIRs must hit the cache instead of reconstructing" + ); + + let airs: Vec< + &dyn AIR< + Field = GoldilocksField, + FieldExtension = GoldilocksField, + PublicInputs = FibonacciPublicInputs, + >, + > = vec![&air_1, &air_2, &air_3]; + + assert!( + Verifier::multi_verify( + &airs, + &multi_proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "verification should succeed when AIRs share all domain parameters" + ); +}