diff --git a/crypto/crypto/src/fiat_shamir/default_transcript.rs b/crypto/crypto/src/fiat_shamir/default_transcript.rs index 819b0f761..2b00ea74d 100644 --- a/crypto/crypto/src/fiat_shamir/default_transcript.rs +++ b/crypto/crypto/src/fiat_shamir/default_transcript.rs @@ -14,6 +14,20 @@ use rand_chacha::{ChaCha20Rng, rand_core::SeedableRng}; pub struct DefaultTranscript { hasher: Keccak256, + /// Byte reservoir feeding [`IsTranscript::sample_bits`]: bytes squeezed + /// from the sponge that have not been consumed yet. A single 32-byte + /// squeeze serves many small index samples — FRI query indices need only + /// `log2(domain)` bits each — instead of one Keccak permutation per index. + /// Cleared on every absorb so sampled bits stay bound to the full prior + /// transcript (mirrors Plonky3's `HashChallenger`, which clears its output + /// buffer on `observe`). + sample_buffer: Vec, + /// Number of bytes at the front of `sample_buffer` already consumed. + sample_cursor: usize, + /// Count of Keccak squeezes performed (`sample()` calls). A cheap profiling + /// counter for the verifier's Keccak-permutation budget — the FRI query + /// phase dominates it, which is exactly what `sample_bits` amortizes. + keccak_squeezes: usize, phantom: PhantomData, } @@ -21,6 +35,9 @@ impl Clone for DefaultTranscript { fn clone(&self) -> Self { Self { hasher: self.hasher.clone(), + sample_buffer: self.sample_buffer.clone(), + sample_cursor: self.sample_cursor, + keccak_squeezes: self.keccak_squeezes, phantom: PhantomData, } } @@ -34,6 +51,9 @@ where pub fn new(data: &[u8]) -> Self { let mut res = Self { hasher: Keccak256::new(), + sample_buffer: Vec::new(), + sample_cursor: 0, + keccak_squeezes: 0, phantom: PhantomData, }; res.append_bytes(data); @@ -44,8 +64,35 @@ where let mut result_hash: [u8; 32] = self.hasher.finalize_reset().into(); result_hash.reverse(); self.hasher.update(result_hash); + self.keccak_squeezes += 1; result_hash } + + /// Number of Keccak squeezes (`sample()` calls) performed so far. The FRI + /// query phase dominates the verifier's Keccak cost; `sample_bits` + /// amortizes one squeeze across many query indices, so this counter drops + /// by roughly `256 / bits` on that phase. + pub fn keccak_squeezes(&self) -> usize { + self.keccak_squeezes + } + + /// Pull `n` fresh reservoir bytes (`n <= 8`) as a big-endian integer, + /// refilling from the sponge in 32-byte squeezes when drained. Big-endian + /// assembly matches the existing `sample_u64` byte convention. + fn next_reservoir_bytes(&mut self, n: usize) -> u64 { + let mut acc: u64 = 0; + for _ in 0..n { + if self.sample_cursor == self.sample_buffer.len() { + let block = self.sample(); + self.sample_buffer.clear(); + self.sample_buffer.extend_from_slice(&block); + self.sample_cursor = 0; + } + acc = (acc << 8) | self.sample_buffer[self.sample_cursor] as u64; + self.sample_cursor += 1; + } + acc + } } impl Default for DefaultTranscript @@ -65,6 +112,11 @@ where { fn append_bytes(&mut self, new_bytes: &[u8]) { self.hasher.update(new_bytes); + // Absorbing new data invalidates any leftover squeezed bits: a later + // `sample_bits` must reflect everything absorbed so far. `append_field_element` + // routes through here, so field-element absorbs clear the reservoir too. + self.sample_buffer.clear(); + self.sample_cursor = 0; } fn append_field_element(&mut self, element: &FieldElement) { @@ -90,6 +142,18 @@ where } } } + + fn sample_bits(&mut self, bits: usize) -> u64 { + assert!( + (1..64).contains(&bits), + "sample_bits: bits must be in 1..=63" + ); + // Power-of-two range: masking is exactly uniform, so no rejection is + // needed. Draw the fewest whole bytes covering `bits` from the shared + // reservoir and keep the low `bits`. + let raw = self.next_reservoir_bytes(bits.div_ceil(8)); + raw & ((1u64 << bits) - 1) + } } impl IsStarkTranscript for DefaultTranscript diff --git a/crypto/crypto/src/fiat_shamir/is_transcript.rs b/crypto/crypto/src/fiat_shamir/is_transcript.rs index eb011e4d4..7d8799133 100644 --- a/crypto/crypto/src/fiat_shamir/is_transcript.rs +++ b/crypto/crypto/src/fiat_shamir/is_transcript.rs @@ -15,6 +15,26 @@ pub trait IsTranscript { fn sample_field_element(&mut self) -> FieldElement; /// Returns a random index in [0, `upper_bound`). fn sample_u64(&mut self, upper_bound: u64) -> u64; + + /// Returns `bits` uniform random bits as a value in `[0, 2^bits)`. + /// + /// The range is a power of two, so masking is unbiased by construction -- + /// no rejection sampling is needed (unlike [`Self::sample_u64`] for a + /// general bound). This is the right primitive for FRI query indices, + /// whose range is exactly the (power-of-two) folded LDE domain size. + /// + /// Concrete transcripts backed by a wide hash squeeze should amortize one + /// permutation across many `sample_bits` calls; the default below does not + /// (it just reuses the power-of-two `sample_u64` path) and exists only so + /// that alternative [`IsTranscript`] implementations keep compiling. A + /// prover and verifier that use the same transcript type always agree. + fn sample_bits(&mut self, bits: usize) -> u64 { + assert!( + (1..64).contains(&bits), + "sample_bits: bits must be in 1..=63" + ); + self.sample_u64(1u64 << bits) + } } pub trait IsStarkTranscript>: IsTranscript { diff --git a/crypto/crypto/src/tests/default_transcript_tests.rs b/crypto/crypto/src/tests/default_transcript_tests.rs index 065ab8751..d0e242fcd 100644 --- a/crypto/crypto/src/tests/default_transcript_tests.rs +++ b/crypto/crypto/src/tests/default_transcript_tests.rs @@ -170,3 +170,100 @@ fn fork_isolation() { assert_eq!(fork_a.sample(), fork_a_fresh.sample()); } + +#[test] +fn sample_bits_values_are_in_range() { + let mut t = DefaultTranscript::::new(&[0x11, 0x22, 0x33]); + let bits = 24; + for _ in 0..200 { + let v = t.sample_bits(bits); + assert!(v < (1u64 << bits), "sample_bits({bits}) out of range: {v}"); + } +} + +#[test] +fn sample_bits_is_deterministic_across_clones() { + let base = DefaultTranscript::::new(&[0xAB; 8]); + let mut a = base.clone(); + let mut b = base.clone(); + let seq_a: Vec = (0..50).map(|_| a.sample_bits(21)).collect(); + let seq_b: Vec = (0..50).map(|_| b.sample_bits(21)).collect(); + assert_eq!( + seq_a, seq_b, + "identical transcript state must yield identical query bits (prover/verifier agreement)" + ); +} + +#[test] +fn sample_bits_covers_full_small_range() { + let mut t = DefaultTranscript::::new(&[0x42]); + let bits = 3; // range [0, 8) + let mut seen = [false; 8]; + for _ in 0..500 { + seen[t.sample_bits(bits) as usize] = true; + } + assert!( + seen.iter().all(|&s| s), + "sample_bits(3) should reach all 8 values (unbiased masking)" + ); +} + +#[test] +fn absorbing_clears_sample_bits_reservoir() { + // No absorb between two samples: the second reuses buffered bytes, no new squeeze. + let mut t = DefaultTranscript::::new(&[0x09]); + let _ = t.sample_bits(24); + let mid = t.keccak_squeezes(); + let _ = t.sample_bits(24); + assert_eq!( + t.keccak_squeezes(), + mid, + "a second sample_bits should reuse the reservoir (no extra Keccak squeeze)" + ); + + // Absorb between: the reservoir is invalidated and the next sample re-squeezes. + let mut t = DefaultTranscript::::new(&[0x09]); + let _ = t.sample_bits(24); + t.append_bytes(&[0xEE]); + let mid = t.keccak_squeezes(); + let _ = t.sample_bits(24); + assert_eq!( + t.keccak_squeezes(), + mid + 1, + "an absorb must clear the reservoir, forcing a fresh squeeze so bits stay bound to it" + ); +} + +#[test] +fn sample_bits_amortizes_keccak_squeezes() { + let num_queries = 80usize; + let bits = 24usize; // 3 bytes per query + + // New path: sample_bits from the shared reservoir. + let mut t_bits = DefaultTranscript::::new(&[0x01, 0x02]); + let before = t_bits.keccak_squeezes(); + for _ in 0..num_queries { + let _ = t_bits.sample_bits(bits); + } + let bits_squeezes = t_bits.keccak_squeezes() - before; + + // Old path: one sample_u64 (= one Keccak squeeze) per query. + let mut t_u64 = DefaultTranscript::::new(&[0x01, 0x02]); + let before = t_u64.keccak_squeezes(); + for _ in 0..num_queries { + let _ = t_u64.sample_u64(1u64 << bits); + } + let u64_squeezes = t_u64.keccak_squeezes() - before; + + // 80 queries * 3 bytes = 240 bytes -> ceil(240 / 32) = 8 squeezes. + let expected = (num_queries * bits.div_ceil(8)).div_ceil(32); + assert_eq!(bits_squeezes, expected, "reservoir squeeze count"); + assert_eq!( + u64_squeezes, num_queries, + "sample_u64 does one squeeze per query" + ); + assert!( + bits_squeezes * 8 <= u64_squeezes, + "sample_bits should cut query-phase Keccak permutations ~10x: {bits_squeezes} vs {u64_squeezes}" + ); +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8c44c42a9..d6814b9c5 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1583,8 +1583,22 @@ pub trait IsStarkProver< transcript: &mut impl IsStarkTranscript, ) -> Vec { let domain_size = domain.lde_roots_of_unity_coset.len() as u64; + // FRI folds the LDE domain by 2, so query indices live in the folded + // domain of size `domain_size / 2` — always a power of two. That makes + // `sample_bits` unbiased by construction (mask, no rejection) and lets + // one Keccak squeeze serve ~256/bits indices instead of one per query. + let folded = domain_size >> 1; + debug_assert!( + folded.is_power_of_two(), + "FRI query domain size must be a power of two" + ); + let bits = folded.trailing_zeros() as usize; + if bits == 0 { + // Degenerate single-index domain: 0 is the only valid index. + return vec![0; number_of_queries]; + } (0..number_of_queries) - .map(|_| (transcript.sample_u64(domain_size >> 1)) as usize) + .map(|_| transcript.sample_bits(bits) as usize) .collect::>() } diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index 64ae24363..67cce0dbc 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -136,8 +136,21 @@ pub trait IsStarkVerifier< transcript: &mut impl IsStarkTranscript, ) -> Vec { let domain_size = domain.lde_length as u64; + // Must mirror the prover exactly: FRI query indices live in the folded + // (power-of-two) domain of size `domain_size / 2`, sampled via + // `sample_bits` (masking, unbiased) so one Keccak squeeze serves many + // indices instead of one permutation per query. + let folded = domain_size >> 1; + debug_assert!( + folded.is_power_of_two(), + "FRI query domain size must be a power of two" + ); + let bits = folded.trailing_zeros() as usize; + if bits == 0 { + return vec![0; number_of_queries]; + } (0..number_of_queries) - .map(|_| (transcript.sample_u64(domain_size >> 1)) as usize) + .map(|_| transcript.sample_bits(bits) as usize) .collect::>() }