From 1665d779427d6aa2b8107c7aef99cc068fa487c7 Mon Sep 17 00:00:00 2001
From: Piotr Paradzinski <dancingwithopenheart@gmail.com>
Date: Mon, 18 May 2026 15:11:23 +0200
Subject: [PATCH] enable error variances up to 32

---
 crates/fhe-math/src/rq/mod.rs | 14 +++++++-------
 crates/fhe-util/src/lib.rs    | 32 +++++++++++++++-----------------
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/crates/fhe-math/src/rq/mod.rs b/crates/fhe-math/src/rq/mod.rs
index 9c043ecf..2417f92a 100644
--- a/crates/fhe-math/src/rq/mod.rs
+++ b/crates/fhe-math/src/rq/mod.rs
@@ -259,15 +259,15 @@ impl<R: RepresentationTag> Poly<R> {
     /// Generate a small polynomial and convert into the specified
     /// representation.
     ///
-    /// Returns an error if the variance does not belong to [1, ..., 16].
+    /// Returns an error if the variance does not belong to [1, ..., 32].
     pub fn small<T: RngCore + CryptoRng>(
         ctx: &Arc<Context>,
         variance: usize,
         rng: &mut T,
     ) -> Result<Self> {
-        if !(1..=16).contains(&variance) {
+        if !(1..=32).contains(&variance) {
             return Err(Error::Default(
-                "The variance should be an integer between 1 and 16".to_string(),
+                "The variance should be an integer between 1 and 32".to_string(),
             ));
         }
 
@@ -857,16 +857,16 @@ mod tests {
             assert!(e.is_err());
             assert_eq!(
                 e.unwrap_err().to_string(),
-                "The variance should be an integer between 1 and 16"
+                "The variance should be an integer between 1 and 32"
             );
-            let e = Poly::<PowerBasis>::small(&ctx, 17, &mut rng);
+            let e = Poly::<PowerBasis>::small(&ctx, 33, &mut rng);
             assert!(e.is_err());
             assert_eq!(
                 e.unwrap_err().to_string(),
-                "The variance should be an integer between 1 and 16"
+                "The variance should be an integer between 1 and 32"
             );
 
-            for i in 1..=16 {
+            for i in 1..=32 {
                 let p = Poly::<PowerBasis>::small(&ctx, i, &mut rng)?;
                 let coefficients = p.coefficients().to_slice().unwrap();
                 let v = q.center_vec(coefficients);
diff --git a/crates/fhe-util/src/lib.rs b/crates/fhe-util/src/lib.rs
index 94d853be..5ecaa942 100644
--- a/crates/fhe-util/src/lib.rs
+++ b/crates/fhe-util/src/lib.rs
@@ -18,37 +18,30 @@ pub fn is_prime(p: u64) -> bool {
 }
 
 /// Sample a vector of independent centered binomial distributions of a given
-/// variance. Returns an error if the variance is strictly larger than 16.
+/// variance. Returns an error if the variance is not between 1 and 32.
 pub fn sample_vec_cbd<R: RngCore + CryptoRng>(
     vector_size: usize,
     variance: usize,
     rng: &mut R,
 ) -> Result<Vec<i64>, &'static str> {
-    if !(1..=16).contains(&variance) {
-        return Err("The variance should be between 1 and 16");
+    if !(1..=32).contains(&variance) {
+        return Err("The variance should be between 1 and 32");
     }
 
     let mut out = Vec::with_capacity(vector_size);
 
     let number_bits = 4 * variance;
-    let mask_add = ((u64::MAX >> (64 - number_bits)) >> (2 * variance)) as u128;
+    let mask_add = (u128::MAX >> (128 - number_bits)) >> (2 * variance);
     let mask_sub = mask_add << (2 * variance);
 
-    let mut current_pool = 0u128;
-    let mut current_pool_nbits = 0;
-
     for _ in 0..vector_size {
-        if current_pool_nbits < number_bits {
-            current_pool |= (rng.next_u64() as u128) << current_pool_nbits;
-            current_pool_nbits += 64;
-        }
-        debug_assert!(current_pool_nbits >= number_bits);
+        let current_pool =
+            (rng.next_u64() as u128) | ((rng.next_u64() as u128) << 64);
+
         out.push(
             ((current_pool & mask_add).count_ones() as i64)
                 - ((current_pool & mask_sub).count_ones() as i64),
         );
-        current_pool >>= number_bits;
-        current_pool_nbits -= number_bits;
     }
 
     Ok(out)
@@ -233,9 +226,9 @@ mod tests {
     fn sample_cbd() {
         let mut rng = rand::rng();
         assert!(sample_vec_cbd(10, 0, &mut rng).is_err());
-        assert!(sample_vec_cbd(10, 17, &mut rng).is_err());
+        assert!(sample_vec_cbd(10, 33, &mut rng).is_err());
 
-        for var in 1..=16 {
+        for var in 1..=32 {
             for size in 0..=100 {
                 let v = sample_vec_cbd(size, var, &mut rng).unwrap();
                 assert_eq!(v.len(), size);
@@ -248,7 +241,12 @@ mod tests {
             // Verifies that the variance is correct. We could probably refine the bound
             // but for now, we will just check that the rounded value is equal to the
             // variance.
-            assert!(variance(&v).round() == (var as f64));
+            let empirical_variance = variance(&v);
+            let tolerance = (var as f64).sqrt().max(1.0);
+            assert!(
+                (empirical_variance - var as f64).abs() <= tolerance,
+                "empirical variance {empirical_variance} differs from expected variance {var}"
+            );
         }
     }