perf(transcript): sample FRI query indices via a masked bit reservoir

crypto/crypto/src/fiat_shamir/default_transcript.rs

@@ -14,13 +14,30 @@ use rand_chacha::{ChaCha20Rng, rand_core::SeedableRng};
 
 pub struct DefaultTranscript<F: HasDefaultTranscript> {
     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<u8>,
+    /// 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<F>,
 }
 
 impl<F: HasDefaultTranscript> Clone for DefaultTranscript<F> {
     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<F> Default for DefaultTranscript<F>
@@ -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<F>) {
@@ -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<F, S> IsStarkTranscript<F, S> for DefaultTranscript<F>

crypto/crypto/src/fiat_shamir/is_transcript.rs

@@ -15,6 +15,26 @@ pub trait IsTranscript<F: IsField> {
     fn sample_field_element(&mut self) -> FieldElement<F>;
     /// 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<F: IsField, S: IsField + IsSubFieldOf<F>>: IsTranscript<F> {

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::<Degree3GoldilocksExtensionField>::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::<Degree3GoldilocksExtensionField>::new(&[0xAB; 8]);
+    let mut a = base.clone();
+    let mut b = base.clone();
+    let seq_a: Vec<u64> = (0..50).map(|_| a.sample_bits(21)).collect();
+    let seq_b: Vec<u64> = (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::<Degree3GoldilocksExtensionField>::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::<Degree3GoldilocksExtensionField>::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::<Degree3GoldilocksExtensionField>::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::<Degree3GoldilocksExtensionField>::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::<Degree3GoldilocksExtensionField>::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}"
+    );
+}

crypto/stark/src/prover.rs

@@ -1563,8 +1563,22 @@ pub trait IsStarkProver<
         transcript: &mut impl IsStarkTranscript<FieldExtension, Field>,
     ) -> Vec<usize> {
         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::<Vec<usize>>()
     }
 

crypto/stark/src/verifier.rs

@@ -131,8 +131,21 @@ pub trait IsStarkVerifier<
         transcript: &mut impl IsStarkTranscript<FieldExtension, Field>,
     ) -> Vec<usize> {
         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::<Vec<usize>>()
     }