feat : updated challange indices logic for clear abstraction

Author: ocdbytes

benches/zook_vs_whir.rs

@@ -33,8 +33,8 @@ use whir::{
     parameters::ProtocolParameters,
     protocols::{
         params::{
-            DecodingRegime, FoldingFactor, KneeWeight, Mode, PowBudget, ProtocolConfig,
-            RateSchedule, SecuritySpec, TuningSpec,
+            DecodingRegime, FoldingFactor, Mode, PowBudget, ProtocolConfig, RateSchedule,
+            SecuritySpec, TuningSpec,
         },
         whir as whir_non_zk,
     },
@@ -370,21 +370,15 @@ fn main() {
             Err(e) => eprintln!("  zook_standard_stepping failed: {e}"),
         }
 
-        eprintln!("[2^{log_size}] zook_zk_adaptive ...");
-        match bench_zook(
-            log_size,
-            Mode::ZeroKnowledge,
-            RateSchedule::Adaptive {
-                knee_weight: KneeWeight::DEFAULT,
-            },
-        ) {
+        eprintln!("[2^{log_size}] zook_zk_stepping ...");
+        match bench_zook(log_size, Mode::ZeroKnowledge, RateSchedule::Stepping) {
             Ok(s) => {
                 let pb = s.proof_bytes;
                 record_cell(
                     &mut csv,
                     log_size,
                     size,
-                    "zook_zk_adaptive",
+                    "zook_zk_stepping",
                     "commit",
                     &s.commit,
                     pb,
@@ -393,7 +387,7 @@ fn main() {
                     &mut csv,
                     log_size,
                     size,
-                    "zook_zk_adaptive",
+                    "zook_zk_stepping",
                     "prove",
                     &s.prove,
                     pb,
@@ -402,7 +396,7 @@ fn main() {
                     &mut csv,
                     log_size,
                     size,
-                    "zook_zk_adaptive",
+                    "zook_zk_stepping",
                     "verify",
                     &s.verify,
                     pb,

src/protocols/challenge_indices.rs

@@ -21,30 +21,10 @@ where
         return if deduplicate { vec![0] } else { vec![0; count] };
     }
 
-    // Size the entropy chunk so `2^(8·size_bytes) ≥ next_pow2(num_leaves)`.
-    // For pow2 `num_leaves`, the entropy space is an exact multiple of
-    // `num_leaves` and rejection never triggers — bit-identical to the
-    // pre-rejection implementation. For non-pow2 `num_leaves`, rejection
-    // sampling eliminates the modular bias of a plain `% num_leaves`.
-    let bits_needed = num_leaves.next_power_of_two().ilog2() as usize;
-    let size_bytes = bits_needed.div_ceil(8);
-
-    // Largest multiple of `num_leaves` below `2^(8·size_bytes)`. u128
-    // accommodates `size_bytes ≤ 16` without shift overflow on 64-bit hosts.
-    let entropy_space: u128 = 1u128 << (8 * size_bytes);
-    let num_leaves_u = num_leaves as u128;
-    let threshold: u128 = (entropy_space / num_leaves_u) * num_leaves_u;
-
+    let domain = IndexDomain::new(num_leaves);
     let mut indices = Vec::with_capacity(count);
-    while indices.len() < count {
-        let mut candidate: u128 = 0;
-        for _ in 0..size_bytes {
-            candidate = (candidate << 8) | u128::from(transcript.verifier_message::<u8>());
-        }
-        if candidate < threshold {
-            indices.push((candidate % num_leaves_u) as usize);
-        }
-        // else: candidate falls in the biased tail — reject and redraw.
+    for _ in 0..count {
+        indices.push(domain.sample(transcript));
     }
 
     if deduplicate {
@@ -54,6 +34,69 @@ where
     indices
 }
 
+/// Rejection-sampling domain for transcript-derived row indices.
+///
+/// For power-of-two domains, `entropy_space` is an exact multiple of
+/// `num_leaves`, so rejection never triggers and sampling is bit-identical to
+/// the legacy `% num_leaves` implementation. For non-power-of-two domains,
+/// candidates in the biased tail are rejected and redrawn.
+#[derive(Clone, Copy, Debug)]
+struct IndexDomain {
+    num_leaves: usize,
+    size_bytes: usize,
+    threshold: u128,
+}
+
+impl IndexDomain {
+    fn new(num_leaves: usize) -> Self {
+        debug_assert!(num_leaves > 1);
+        let bits_needed = ceil_log2(num_leaves);
+        let size_bytes = bits_needed.div_ceil(8);
+        let entropy_bits = 8 * size_bytes;
+        debug_assert!(entropy_bits < u128::BITS as usize);
+
+        let entropy_space = 1u128 << entropy_bits;
+        let num_leaves_u = num_leaves as u128;
+        let threshold = (entropy_space / num_leaves_u) * num_leaves_u;
+
+        Self {
+            num_leaves,
+            size_bytes,
+            threshold,
+        }
+    }
+
+    fn sample<T>(&self, transcript: &mut T) -> usize
+    where
+        T: VerifierMessage,
+        u8: Decoding<[T::U]>,
+    {
+        loop {
+            let candidate = self.draw_candidate(transcript);
+            if candidate < self.threshold {
+                return (candidate % self.num_leaves as u128) as usize;
+            }
+        }
+    }
+
+    fn draw_candidate<T>(&self, transcript: &mut T) -> u128
+    where
+        T: VerifierMessage,
+        u8: Decoding<[T::U]>,
+    {
+        let mut candidate = 0u128;
+        for _ in 0..self.size_bytes {
+            candidate = (candidate << 8) | u128::from(transcript.verifier_message::<u8>());
+        }
+        candidate
+    }
+}
+
+fn ceil_log2(n: usize) -> usize {
+    debug_assert!(n > 1);
+    usize::BITS as usize - (n - 1).leading_zeros() as usize
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;