fix(fiat-shamir): rejection-sample in sample_in_range to remove truncation bias

crates/backend/fiat-shamir/src/challenger.rs

@@ -64,12 +64,60 @@ impl<F: PrimeField64, P: Permutation<[F; WIDTH]>> Challenger<F, P> {
     /// Warning: not perfectly uniform
     pub fn sample_in_range(&mut self, bits: usize, n_samples: usize) -> Vec<usize> {
         assert!(bits < F::bits());
-        let sampled_fe = self.sample_many(n_samples.div_ceil(RATE)).into_iter().flatten();
-        let mut res = Vec::new();
-        for fe in sampled_fe.take(n_samples) {
-            let rand_usize = fe.as_canonical_u64() as usize;
-            res.push(rand_usize & ((1 << bits) - 1));
+
+        let range: u64 = 1u64 << bits;
+        // Valori >= threshold vengono scartati: tenerli darebbe alle residue
+        // basse (0..F::ORDER_U64 % range) massa di probabilita' extra.
+        let threshold: u64 = F::ORDER_U64 - (F::ORDER_U64 % range);
+
+        let mut res = Vec::with_capacity(n_samples);
+        while res.len() < n_samples {
+            let remaining = n_samples - res.len();
+            let batch = self.sample_many(remaining.div_ceil(RATE));
+            for fe in batch.into_iter().flatten() {
+                let candidate = fe.as_canonical_u64();
+                if candidate >= threshold {
+                    continue;
+                }
+                res.push((candidate & (range - 1)) as usize);
+                if res.len() == n_samples {
+                    break;
+                }
+            }
         }
         res
     }
 }
+
+#[cfg(test)]
+mod bias_regression_tests {
+    use super::*;
+    use koala_bear::{KoalaBear, default_koalabear_poseidon1_16};
+
+    #[test]
+    fn sample_in_range_is_not_biased_toward_low_residues() {
+        let bits = 4;
+        let n = 1usize << bits;
+        let samples = 200_000;
+
+        let perm = default_koalabear_poseidon1_16();
+        let mut challenger: Challenger<KoalaBear, _> = Challenger::new(perm, Default::default());
+        challenger.duplex();
+        let out = challenger.sample_in_range(bits, samples);
+
+        let mut counts = vec![0usize; n];
+        for v in out {
+            assert!(v < n);
+            counts[v] += 1;
+        }
+
+        let expected = samples as f64 / n as f64;
+        for (i, &c) in counts.iter().enumerate() {
+            let deviation = (c as f64 - expected).abs() / expected;
+            assert!(
+                deviation < 0.05,
+                "residue {i} deviates {deviation:.3} from uniform (count={c}, expected={expected})"
+            );
+        }
+    }
+}