feat(shortint): add PRF primitive with re-randomization

- required for some maths properties we need
- limited API surface in shortint to have the minimum amount of functions
to maintain
- factorized the PRF inputs generation to share between rerand and non
rerand version
- seed creation returns the data hashed to represent the output random bits
distribution so that it can be hashed for rerand as well
- added test checking as many properties as possible verifying the rerand
variant is strictly equivalent to the non rerand variant except for the bit
content of the LWE
- integer wrapping and HL to follow

Author: IceTDrinker

tfhe/src/high_level_api/integers/oprf.rs

@@ -591,9 +591,12 @@ mod test {
     };
     use crate::prelude::FheDecrypt;
     use crate::shortint::oprf::test::test_uniformity;
-    use crate::shortint::parameters::test_params::TEST_PARAM_MESSAGE_2_CARRY_2_PBS_KS_GAUSSIAN_2M128;
+    use crate::shortint::parameters::test_params::{
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+        TEST_PARAM_MESSAGE_2_CARRY_2_PBS_KS_GAUSSIAN_2M128,
+    };
     use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128;
-    use crate::{generate_keys, set_server_key, ConfigBuilder, FheUint8, Seed};
+    use crate::{generate_keys, set_server_key, ConfigBuilder, FheInt8, FheUint8, Seed};
     use num_bigint::BigUint;
     use rand::{thread_rng, Rng};
     use rayon::iter::{IntoParallelIterator, ParallelIterator};
@@ -890,6 +893,32 @@ mod test {
         assert!(result_bounded < (1 << 3));
     }
 
+    #[test]
+    fn test_oprf_bounded_zero() {
+        let config = ConfigBuilder::with_custom_parameters(
+            TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+        )
+        .use_dedicated_oprf_key(true)
+        .build();
+
+        let (client_key, server_key) = generate_keys(config);
+        set_server_key(server_key);
+
+        // Do not use static seed in production
+        let ct_unsigned_bounded = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), 0);
+        assert!(ct_unsigned_bounded.is_trivial());
+        let result_unsigned_bounded: u8 = ct_unsigned_bounded.decrypt(&client_key);
+        // PRF with 0 bits is equivalent to modulo 1 meaning only 0 is generated
+        assert_eq!(result_unsigned_bounded, 0);
+
+        // Do not use static seed in production
+        let ct_signed_bounded = FheInt8::generate_oblivious_pseudo_random_bounded(Seed(0), 0);
+        assert!(ct_signed_bounded.is_trivial());
+        let result_signed_bounded: i8 = ct_signed_bounded.decrypt(&client_key);
+        // PRF with 0 bits is equivalent to modulo 1 meaning only 0 is generated
+        assert_eq!(result_signed_bounded, 0);
+    }
+
     #[cfg(feature = "gpu")]
     mod gpu {
         use super::*;
@@ -904,6 +933,29 @@ mod test {
         use rayon::prelude::{IntoParallelRefIterator, ParallelSlice};
         use rayon::ThreadPoolBuilder;
 
+        #[test]
+        fn test_oprf_bounded_zero() {
+            for setup_fn in crate::high_level_api::integers::unsigned::tests::gpu::GPU_SETUP_FN {
+                let client_key = setup_fn();
+
+                // Do not use static seed in production
+                let ct_unsigned_bounded =
+                    FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), 0);
+                assert!(ct_unsigned_bounded.is_trivial());
+                let result_unsigned_bounded: u8 = ct_unsigned_bounded.decrypt(&client_key);
+                // PRF with 0 bits is equivalent to modulo 1 meaning only 0 is generated
+                assert_eq!(result_unsigned_bounded, 0);
+
+                // Do not use static seed in production
+                let ct_signed_bounded =
+                    FheInt8::generate_oblivious_pseudo_random_bounded(Seed(0), 0);
+                assert!(ct_signed_bounded.is_trivial());
+                let result_signed_bounded: i8 = ct_signed_bounded.decrypt(&client_key);
+                // PRF with 0 bits is equivalent to modulo 1 meaning only 0 is generated
+                assert_eq!(result_signed_bounded, 0);
+            }
+        }
+
         #[test]
         fn test_oprf_gpu() {
             for setup_fn in crate::high_level_api::integers::unsigned::tests::gpu::GPU_SETUP_FN {

tfhe/src/high_level_api/integers/signed/base.rs

@@ -912,6 +912,27 @@ where
         self.ciphertext.on_cpu().decrypt_trivial()
     }
 
+    /// Returns true if the ciphertext is a trivial encryption
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tfhe::prelude::*;
+    /// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheInt16};
+    ///
+    /// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
+    /// set_server_key(server_key);
+    ///
+    /// let non_trivial = FheInt16::encrypt(1i16, &client_key);
+    /// assert!(!non_trivial.is_trivial());
+    ///
+    /// let trivial = FheInt16::encrypt_trivial(2i16);
+    /// assert!(trivial.is_trivial());
+    /// ```
+    pub fn is_trivial(&self) -> bool {
+        self.ciphertext.on_cpu().is_trivial()
+    }
+
     /// Reverse the bit of the signed integer
     ///
     /// # Example

tfhe/src/integer/gpu/server_key/radix/oprf.rs

@@ -176,6 +176,7 @@ where
         target_sks: &CudaServerKey,
         streams: &CudaStreams,
     ) -> CudaUnsignedRadixCiphertext {
+        assert!(target_sks.message_modulus.0.is_power_of_two());
         let message_bits_count = target_sks.message_modulus.0.ilog2() as u64;
         let range_bits_count = message_bits_count * num_blocks;
         assert!(range_bits_count > 0);
@@ -291,6 +292,7 @@ where
         target_sks: &CudaServerKey,
         streams: &CudaStreams,
     ) -> CudaSignedRadixCiphertext {
+        assert!(target_sks.message_modulus.0.is_power_of_two());
         let message_bits_count = target_sks.message_modulus.0.ilog2() as u64;
         let range_bits_count = message_bits_count * num_blocks;
         assert!(range_bits_count > 0);
@@ -388,9 +390,10 @@ where
         result
     }
 
-    // Core private implementation that calls the OPRF backend.
-    // This function contains the main logic for both bounded and unbounded generation.
-    //
+    /// Core private implementation that calls the OPRF backend.
+    /// This function contains the main logic for both bounded and unbounded generation.
+    ///
+    /// Caller must ensure total_random_bits is non 0 otherwise this function will panic.
     fn generate_multiblocks_oblivious_pseudo_random(
         &self,
         result: &mut CudaRadixCiphertext,
@@ -412,14 +415,18 @@ where
         let polynomial_size = bootstrapping_key.polynomial_size();
         let in_lwe_size = input_lwe_dimension.to_lwe_size();
         let message_bits_count = target_sks.message_modulus.0.ilog2();
+        let carry_bits_count = target_sks.carry_modulus.0.ilog2();
+        let bits_per_block = message_bits_count + carry_bits_count + 1;
 
-        let seeded = create_random_from_seed_modulus_switched(
+        let (seeded, _rle_info) = create_random_from_seed_modulus_switched(
             seed,
             in_lwe_size,
             polynomial_size,
             &[total_random_bits],
-            message_bits_count as u64,
+            message_bits_count.into(),
+            bits_per_block.into(),
         );
+
         let h_seeded_lwe_list: Vec<u64> = seeded
             .into_iter()
             .flat_map(|(seeded, _bits)| {
@@ -471,6 +478,13 @@ where
         }
     }
 
+    /// # Panics
+    ///
+    /// Panics if:
+    /// - `target_sks.message_modulus` is not a power of 2
+    /// - `excluded_upper_bound` is a power of 2 use
+    ///   [`Self::par_generate_oblivious_pseudo_random_unsigned_integer_bounded`] instead
+    /// - `excluded_upper_bound.ilog2() + 1` is greater than the output bit count
     pub fn par_generate_oblivious_pseudo_random_unsigned_custom_range(
         &self,
         seed: impl OprfSeed,
@@ -484,7 +498,13 @@ where
             target_sks.message_modulus.0.is_power_of_two(),
             "Message modulus must be a power of two"
         );
-        let message_bits_count = target_sks.message_modulus.0.ilog2() as u64;
+        assert!(
+            target_sks.carry_modulus.0.is_power_of_two(),
+            "Carry modulus must be a power of two"
+        );
+        let message_bits_count: u64 = target_sks.message_modulus.0.ilog2().into();
+        let carry_bits_count: u64 = target_sks.carry_modulus.0.ilog2().into();
+        let bits_per_block = message_bits_count + carry_bits_count + 1;
 
         assert!(
             !excluded_upper_bound.is_power_of_two(),
@@ -528,12 +548,13 @@ where
             .iter_as::<u64>()
             .collect::<Vec<_>>();
 
-        let seeded = create_random_from_seed_modulus_switched(
+        let (seeded, _rle_info) = create_random_from_seed_modulus_switched(
             seed,
             in_lwe_size,
             polynomial_size,
             &[num_input_random_bits],
             message_bits_count,
+            bits_per_block,
         );
 
         let h_seeded_lwe_list: Vec<u64> = seeded

tfhe/src/shortint/ciphertext/re_randomization.rs

@@ -13,6 +13,8 @@ use crate::core_crypto::entities::{LweCiphertext, LweCompactCiphertextList, Plai
 use crate::core_crypto::prelude::lwe_compact_ciphertext_list_add_assign;
 use crate::shortint::ciphertext::NoiseLevel;
 use crate::shortint::key_switching_key::KeySwitchingKeyMaterialView;
+use crate::shortint::oprf::{OprfSeed, RandomBitsRleLeBytes};
+use crate::shortint::public_key::compact::TFHE_PKE_DOMAIN_SEPARATOR;
 use crate::shortint::{Ciphertext, CompactPublicKey, PBSOrder};
 
 use rayon::prelude::*;
@@ -106,6 +108,22 @@ impl From<blake3::Hasher> for ReRandomizationSeedHasher {
 /// This type cannot be cloned or copied, as a seed should only be used once.
 pub struct ReRandomizationSeed(pub(crate) XofSeed);
 
+impl ReRandomizationSeed {
+    pub(crate) fn new_prf_rerand_seed(
+        hash_algo: ReRandomizationHashAlgo,
+        prf_seed: impl OprfSeed,
+        random_bits_rle_bytes: &RandomBitsRleLeBytes,
+    ) -> Self {
+        let mut seed_gen = ReRandomizationSeedGen::new_prf_rerand_seed_gen(
+            hash_algo,
+            prf_seed,
+            random_bits_rle_bytes,
+        );
+
+        seed_gen.next_seed()
+    }
+}
+
 /// The context that will be hashed and used to generate unique [`ReRandomizationSeed`].
 ///
 /// At this level, the context will directly hash any data passed to it.
@@ -226,6 +244,27 @@ pub struct ReRandomizationSeedGen {
 }
 
 impl ReRandomizationSeedGen {
+    pub(crate) fn new_prf_rerand_seed_gen(
+        hash_algo: ReRandomizationHashAlgo,
+        prf_seed: impl OprfSeed,
+        random_bits_rle_bytes: &RandomBitsRleLeBytes,
+    ) -> Self {
+        const PRF_RERAND_DOMAIN_SEPARATOR: [u8; XofSeed::DOMAIN_SEP_LEN] = *b"PRF_RRND";
+
+        let mut context = ReRandomizationContext::new_with_hasher(
+            TFHE_PKE_DOMAIN_SEPARATOR,
+            ReRandomizationSeedHasher::new(hash_algo, PRF_RERAND_DOMAIN_SEPARATOR),
+        );
+
+        let prf_seed = prf_seed.into_bytes();
+        let prf_seed = prf_seed.as_ref();
+
+        context.add_bytes(prf_seed);
+        context.add_bytes(random_bits_rle_bytes.get());
+
+        context.finalize()
+    }
+
     pub fn next_seed(&mut self) -> ReRandomizationSeed {
         let current_seed_index = self.next_seed_index;
         self.next_seed_index += 1;