From cdc25d6c60073574ab30771c54c9a26a51989297 Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Fri, 21 Apr 2023 12:36:48 +0200
Subject: [PATCH] chore(core): add more sanity checks on RNG

---
 .../commons/generators/encryption.rs          | 176 ++++++++++++++++++
 .../commons/math/random/generator.rs          |  33 +++-
 2 files changed, 201 insertions(+), 8 deletions(-)

diff --git a/tfhe/src/core_crypto/commons/generators/encryption.rs b/tfhe/src/core_crypto/commons/generators/encryption.rs
index b008f60cb1..91ba3d674c 100644
--- a/tfhe/src/core_crypto/commons/generators/encryption.rs
+++ b/tfhe/src/core_crypto/commons/generators/encryption.rs
@@ -629,4 +629,180 @@ mod test {
     fn noise_gen_native_u128() {
         noise_gen_native::<u128>();
     }
+
+    fn noise_gen_custom_mod<Scalar: UnsignedTorus>(ciphertext_modulus: CiphertextModulus<Scalar>) {
+        let mut gen = new_encryption_random_generator();
+
+        let bits = (Scalar::BITS / 2) as i32;
+
+        for _ in 0..1000 {
+            let val: Scalar =
+                gen.random_noise_custom_mod(StandardDev(2.0f64.powi(-bits)), ciphertext_modulus);
+            assert!(val != Scalar::ZERO);
+        }
+    }
+
+    #[test]
+    fn noise_gen_custom_mod_u32() {
+        noise_gen_custom_mod::<u32>(CiphertextModulus::try_new_power_of_2(31).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_custom_mod_u64() {
+        noise_gen_custom_mod::<u64>(CiphertextModulus::try_new_power_of_2(63).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_custom_mod_u128() {
+        noise_gen_custom_mod::<u128>(CiphertextModulus::try_new_power_of_2(127).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_native_custom_mod_u32() {
+        noise_gen_custom_mod::<u32>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn noise_gen_native_custom_mod_u64() {
+        noise_gen_custom_mod::<u64>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn noise_gen_native_custom_mod_u128() {
+        noise_gen_custom_mod::<u128>(CiphertextModulus::new_native());
+    }
+
+    fn noise_gen_slice_native<Scalar: UnsignedTorus>() {
+        let mut gen = new_encryption_random_generator();
+
+        let bits = (Scalar::BITS / 2) as i32;
+
+        let mut slice = vec![Scalar::ZERO; 1000];
+        gen.fill_slice_with_random_noise(&mut slice, StandardDev(2.0f64.powi(-bits)));
+        assert!(slice.iter().all(|&x| x != Scalar::ZERO))
+    }
+
+    #[test]
+    fn noise_gen_slice_native_u32() {
+        noise_gen_slice_native::<u32>();
+    }
+
+    #[test]
+    fn noise_gen_slice_native_u64() {
+        noise_gen_slice_native::<u64>();
+    }
+
+    #[test]
+    fn noise_gen_slice_native_u128() {
+        noise_gen_slice_native::<u128>();
+    }
+
+    fn noise_gen_slice_custom_mod<Scalar: UnsignedTorus>(
+        ciphertext_modulus: CiphertextModulus<Scalar>,
+    ) {
+        let mut gen = new_encryption_random_generator();
+
+        let bits = (Scalar::BITS / 2) as i32;
+
+        let mut slice = vec![Scalar::ZERO; 1000];
+        gen.fill_slice_with_random_noise_custom_mod(
+            &mut slice,
+            StandardDev(2.0f64.powi(-bits)),
+            ciphertext_modulus,
+        );
+        assert!(slice.iter().all(|&x| x != Scalar::ZERO))
+    }
+
+    #[test]
+    fn noise_gen_slice_custom_mod_u32() {
+        noise_gen_slice_custom_mod::<u32>(CiphertextModulus::try_new_power_of_2(31).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_slice_custom_mod_u64() {
+        noise_gen_slice_custom_mod::<u64>(CiphertextModulus::try_new_power_of_2(63).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_slice_custom_mod_u128() {
+        noise_gen_slice_custom_mod::<u128>(CiphertextModulus::try_new_power_of_2(127).unwrap());
+    }
+
+    #[test]
+    fn noise_gen_slice_native_custom_mod_u32() {
+        noise_gen_slice_custom_mod::<u32>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn noise_gen_slice_native_custom_mod_u64() {
+        noise_gen_slice_custom_mod::<u64>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn noise_gen_slice_native_custom_mod_u128() {
+        noise_gen_slice_custom_mod::<u128>(CiphertextModulus::new_native());
+    }
+
+    fn mask_gen_slice_native<Scalar: UnsignedTorus>() {
+        let mut gen = new_encryption_random_generator();
+
+        let mut slice = vec![Scalar::ZERO; 1000];
+        gen.fill_slice_with_random_mask(&mut slice);
+        assert!(slice.iter().all(|&x| x != Scalar::ZERO))
+    }
+
+    #[test]
+    fn mask_gen_native_u32() {
+        mask_gen_slice_native::<u32>();
+    }
+
+    #[test]
+    fn mask_gen_native_u64() {
+        mask_gen_slice_native::<u64>();
+    }
+
+    #[test]
+    fn mask_gen_native_u128() {
+        mask_gen_slice_native::<u128>();
+    }
+
+    fn mask_gen_slice_custom_mod<Scalar: UnsignedTorus>(
+        ciphertext_modulus: CiphertextModulus<Scalar>,
+    ) {
+        let mut gen = new_encryption_random_generator();
+
+        let mut slice = vec![Scalar::ZERO; 1000];
+        gen.fill_slice_with_random_mask_custom_mod(&mut slice, ciphertext_modulus);
+        assert!(slice.iter().all(|&x| x != Scalar::ZERO))
+    }
+
+    #[test]
+    fn mask_gen_slice_custom_mod_u32() {
+        mask_gen_slice_custom_mod::<u32>(CiphertextModulus::try_new_power_of_2(31).unwrap());
+    }
+
+    #[test]
+    fn mask_gen_slice_custom_mod_u64() {
+        mask_gen_slice_custom_mod::<u64>(CiphertextModulus::try_new_power_of_2(63).unwrap());
+    }
+
+    #[test]
+    fn mask_gen_slice_custom_mod_u128() {
+        mask_gen_slice_custom_mod::<u128>(CiphertextModulus::try_new_power_of_2(127).unwrap());
+    }
+
+    #[test]
+    fn mask_gen_slice_native_custom_mod_u32() {
+        mask_gen_slice_custom_mod::<u32>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn mask_gen_slice_native_custom_mod_u64() {
+        mask_gen_slice_custom_mod::<u64>(CiphertextModulus::new_native());
+    }
+
+    #[test]
+    fn mask_gen_slice_native_custom_mod_u128() {
+        mask_gen_slice_custom_mod::<u128>(CiphertextModulus::new_native());
+    }
 }
diff --git a/tfhe/src/core_crypto/commons/math/random/generator.rs b/tfhe/src/core_crypto/commons/math/random/generator.rs
index 061b5c85db..c2528ecd14 100644
--- a/tfhe/src/core_crypto/commons/math/random/generator.rs
+++ b/tfhe/src/core_crypto/commons/math/random/generator.rs
@@ -175,8 +175,9 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use concrete_csprng::seeders::Seed;
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1u32; 100];
+    /// let mut vec = vec![0u32; 1000];
     /// generator.fill_slice_with_random_uniform(&mut vec);
+    /// assert!(vec.iter().all(|&x| x != 0));
     /// ```
     pub fn fill_slice_with_random_uniform<Scalar>(&mut self, output: &mut [Scalar])
     where
@@ -196,11 +197,12 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// use tfhe::core_crypto::commons::parameters::CiphertextModulus;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1u32; 100];
+    /// let mut vec = vec![0u32; 1000];
     /// generator.fill_slice_with_random_uniform_custom_mod(
     ///     &mut vec,
     ///     CiphertextModulus::try_new_power_of_2(31).unwrap(),
     /// );
+    /// assert!(vec.iter().all(|&x| x != 0));
     /// ```
     pub fn fill_slice_with_random_uniform_custom_mod<Scalar>(
         &mut self,
@@ -243,8 +245,9 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use concrete_csprng::seeders::Seed;
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![2u32; 100];
+    /// let mut vec = vec![0u32; 1000];
     /// generator.fill_slice_with_random_uniform_binary(&mut vec);
+    /// assert!(vec.iter().any(|&x| x != 0));
     /// ```
     pub fn fill_slice_with_random_uniform_binary<Scalar>(&mut self, output: &mut [Scalar])
     where
@@ -348,11 +351,15 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// // check that both samples are in 6 sigmas.
     /// assert!(g1.abs() <= 6.);
     /// assert!(g2.abs() <= 6.);
+    /// assert!(g1 != 0.);
+    /// assert!(g2 != 0.);
     /// // for f64
     /// let (g1, g2): (f64, f64) = generator.random_gaussian(0. as f64, 1. as f64);
     /// // check that both samples are in 6 sigmas.
     /// assert!(g1.abs() <= 6.);
     /// assert!(g2.abs() <= 6.);
+    /// assert!(g1 != 0.);
+    /// assert!(g2 != 0.);
     /// ```
     pub fn random_gaussian<Float, Scalar>(&mut self, mean: Float, std: Float) -> (Scalar, Scalar)
     where
@@ -371,8 +378,9 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use concrete_csprng::seeders::Seed;
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1000f32; 100];
+    /// let mut vec = vec![0f32; 1000];
     /// generator.fill_slice_with_random_gaussian(&mut vec, 0., 1.);
+    /// assert!(vec.iter().all(|&x| x != 0.));
     /// ```
     pub fn fill_slice_with_random_gaussian<Float, Scalar>(
         &mut self,
@@ -404,13 +412,14 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// use tfhe::core_crypto::commons::parameters::CiphertextModulus;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1000u64; 100];
+    /// let mut vec = vec![0u64; 1000];
     /// generator.fill_slice_with_random_gaussian_custom_mod(
     ///     &mut vec,
     ///     0.,
     ///     1.,
     ///     CiphertextModulus::try_new_power_of_2(63).unwrap(),
     /// );
+    /// assert!(vec.iter().all(|&x| x != 0));
     /// ```
     pub fn fill_slice_with_random_gaussian_custom_mod<Float, Scalar>(
         &mut self,
@@ -453,8 +462,9 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use concrete_csprng::seeders::Seed;
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1000u32; 100];
+    /// let mut vec = vec![0u32; 1000];
     /// generator.unsigned_torus_slice_wrapping_add_random_gaussian_assign(&mut vec, 0., 1.);
+    /// assert!(vec.iter().all(|&x| x != 0));
     /// ```
     pub fn unsigned_torus_slice_wrapping_add_random_gaussian_assign<Float, Scalar>(
         &mut self,
@@ -484,9 +494,16 @@ impl<G: ByteRandomGenerator> RandomGenerator<G> {
     /// use concrete_csprng::generators::SoftwareRandomGenerator;
     /// use concrete_csprng::seeders::Seed;
     /// use tfhe::core_crypto::commons::math::random::RandomGenerator;
+    /// use tfhe::core_crypto::commons::parameters::CiphertextModulus;
     /// let mut generator = RandomGenerator::<SoftwareRandomGenerator>::new(Seed(0));
-    /// let mut vec = vec![1000u32; 100];
-    /// generator.unsigned_torus_slice_wrapping_add_random_gaussian_assign(&mut vec, 0., 1.);
+    /// let mut vec = vec![0u32; 1000];
+    /// generator.unsigned_torus_slice_wrapping_add_random_gaussian_custom_mod_assign(
+    ///     &mut vec,
+    ///     0.,
+    ///     1.,
+    ///     CiphertextModulus::try_new_power_of_2(31).unwrap(),
+    /// );
+    /// assert!(vec.iter().all(|&x| x != 0));
     /// ```
     pub fn unsigned_torus_slice_wrapping_add_random_gaussian_custom_mod_assign<Float, Scalar>(
         &mut self,