add oprf over any range

Author: mayeul-zama

tfhe/Cargo.toml

@@ -27,6 +27,7 @@ rand_distr = "0.4.3"
 criterion = "0.5.1"
 doc-comment = "0.3.3"
 serde_json = "1.0.94"
+num-bigint = "0.4.6"
 # clap has to be pinned as its minimum supported rust version
 # changes often between minor releases, which breaks our CI
 clap = { version = "=4.5.30", features = ["derive"] }

tfhe/src/high_level_api/integers/oprf.rs

@@ -150,6 +150,66 @@ impl<Id: FheUintId> FheUint<Id> {
             }
         })
     }
+
+    /// Generates an encrypted unsigned integer
+    /// taken almost uniformly in `[0, excluded_upper_bound[` using the given seed.
+    /// The encrypted value is oblivious to the server.
+    /// It can be useful to make server random generation deterministic.
+    ///
+    /// The norm-1 distance (defined as ∆(, ) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)| between the actual distribution and the target uniform distribution is below the `max_distance` argument.
+    ///
+    /// A safe value for `max_distance` is `2^-128`. It is the default value if None is provided.
+    ///
+    /// Higher values allow better performance but must be considered carefully in the context of their target application
+    /// as it may have serious unintended consequences.
+    ///
+    /// ```rust
+    /// use tfhe::prelude::FheDecrypt;
+    /// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, Seed};
+    ///
+    /// let config = ConfigBuilder::default().build();
+    /// let (client_key, server_key) = generate_keys(config);
+    ///
+    /// set_server_key(server_key);
+    ///
+    /// let excluded_upper_bound = 3;
+    ///
+    /// let ct_res = FheUint8::generate_oblivious_pseudo_random_custom_range(Seed(0), excluded_upper_bound, None);
+    ///
+    /// let dec_result: u16 = ct_res.decrypt(&client_key);
+    /// assert!(dec_result < excluded_upper_bound as u16);
+    /// ```
+    pub fn generate_oblivious_pseudo_random_custom_range(
+        seed: Seed,
+        excluded_upper_bound: u64,
+        max_distance: Option<f64>,
+    ) -> Self {
+        global_state::with_internal_keys(|key| match key {
+            InternalServerKey::Cpu(key) => {
+                let num_blocks_output = Id::num_blocks(key.message_modulus()) as u64;
+
+                let ct = key
+                    .pbs_key()
+                    .par_generate_oblivious_pseudo_random_unsigned_custom_range2(
+                        seed,
+                        excluded_upper_bound,
+                        num_blocks_output,
+                        max_distance,
+                    );
+
+                Self::new(ct, key.tag.clone(), ReRandomizationMetadata::default())
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(cuda_key) => {
+                panic!("Gpu does not support this operation yet.")
+            }
+            #[cfg(feature = "hpu")]
+            InternalServerKey::Hpu(_device) => {
+                panic!("Hpu does not support this operation yet.")
+            }
+        })
+    }
+
     #[cfg(feature = "gpu")]
     /// Returns the amount of memory required to execute generate_oblivious_pseudo_random_bounded
     ///

tfhe/src/integer/oprf.rs

@@ -1,6 +1,7 @@
 use super::{RadixCiphertext, ServerKey, SignedRadixCiphertext};
 use crate::core_crypto::commons::generators::DeterministicSeeder;
 use crate::core_crypto::prelude::DefaultRandomGenerator;
+use crate::shortint::MessageModulus;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
 pub use tfhe_csprng::seeders::{Seed, Seeder};
@@ -229,6 +230,116 @@ impl ServerKey {
 
         result
     }
+
+    /// Generates an encrypted `num_blocks_output` blocks unsigned integer
+    /// taken almost uniformly in `[0, excluded_upper_bound[` using the given seed.
+    /// The encrypted value is oblivious to the server.
+    /// It can be useful to make server random generation deterministic.
+    ///
+    /// The norm-1 distance (defined as ∆(, ) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)| between the actual distribution and the target uniform distribution is below the `max_distance` argument.
+    ///
+    /// A safe value for `max_distance` is `2^-128`. It is the default value if None is provided.
+    ///
+    /// Higher values allow better performance but must be considered carefully in the context of their target application
+    /// as it may have serious unintended consequences.
+    ///
+    /// ```rust
+    /// use tfhe::integer::gen_keys_radix;
+    /// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128;
+    /// use tfhe::Seed;
+    ///
+    /// let size = 4;
+    ///
+    /// let (cks, sks) = gen_keys_radix(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128, size);
+    ///
+    /// let excluded_upper_bound = 3;
+    /// let num_blocks_output = 8;
+    ///
+    /// let ct_res = sks.par_generate_oblivious_pseudo_random_unsigned_custom_range2(Seed(0), excluded_upper_bound,num_blocks_output, None);
+    ///
+    /// let dec_result: u64 = cks.decrypt(&ct_res);
+    /// assert!(dec_result < excluded_upper_bound);
+    /// ```
+    pub fn par_generate_oblivious_pseudo_random_unsigned_custom_range2(
+        &self,
+        seed: Seed,
+        excluded_upper_bound: u64,
+        num_blocks_output: u64,
+        max_distance: Option<f64>,
+    ) -> RadixCiphertext {
+        let max_distance = max_distance.unwrap_or(2_f64.powi(-128));
+
+        let message_modulus = self.message_modulus();
+
+        let num_input_random_bits = num_input_random_bits_for_max_distance(
+            excluded_upper_bound,
+            max_distance,
+            message_modulus,
+        );
+
+        self.par_generate_oblivious_pseudo_random_unsigned_custom_range(
+            seed,
+            num_input_random_bits,
+            excluded_upper_bound,
+            num_blocks_output,
+        )
+    }
+}
+
+fn num_input_random_bits_for_max_distance(
+    excluded_upper_bound: u64,
+    max_distance: f64,
+    message_modulus: MessageModulus,
+) -> u64 {
+    let log_message_modulus = message_modulus.0.ilog2() as u64;
+
+    let mut random_block_count = 1;
+
+    let random_block_count = loop {
+        let random_bit_count = random_block_count * log_message_modulus;
+
+        let remainder = mod_pow_2(random_bit_count, excluded_upper_bound) as f64;
+
+        let distance = remainder * (excluded_upper_bound as f64 - remainder)
+            / (2_f64.powi(random_bit_count as i32) * excluded_upper_bound as f64);
+
+        if distance < max_distance {
+            break random_block_count;
+        }
+
+        random_block_count += 1;
+    };
+
+    random_block_count * log_message_modulus
+}
+
+// Computes 2^exponent % modulus
+fn mod_pow_2(exponent: u64, modulus: u64) -> u64 {
+    if modulus == 1 {
+        return 0;
+    }
+
+    let mut result: u128 = 1;
+    let mut base: u128 = 2; // We are calculating 2^i
+
+    // We cast exponent to u128 to match the loop, though u64 is fine
+    let mut exp = exponent;
+    let mod_val = modulus as u128;
+
+    while exp > 0 {
+        // If exponent is odd, multiply result with base
+        if exp % 2 == 1 {
+            result = (result * base) % mod_val;
+        }
+
+        // Square the base
+        base = (base * base) % mod_val;
+
+        // Divide exponent by 2
+        exp /= 2;
+    }
+
+    result as u64
 }
 
 impl ServerKey {
@@ -362,3 +473,101 @@ impl ServerKey {
         SignedRadixCiphertext::from(blocks)
     }
 }
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use crate::integer::gen_keys_radix;
+    use crate::shortint::oprf::test::test_uniformity;
+    use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128;
+    use crate::Seed;
+    use num_bigint::BigUint;
+    use rand::{thread_rng, Rng};
+
+    // Helper: The "Oracle" implementation using BigInt
+    // This is slow but mathematically guaranteed to be correct.
+    fn oracle_mod_pow_2(exponent: u64, modulus: u64) -> u64 {
+        if modulus == 0 {
+            panic!("div by 0");
+        }
+        if modulus == 1 {
+            return 0;
+        }
+
+        let base = BigUint::from(2u32);
+        let exp = BigUint::from(exponent);
+        let modu = BigUint::from(modulus);
+
+        let res = base.modpow(&exp, &modu);
+        res.iter_u64_digits().next().unwrap_or(0)
+    }
+
+    #[test]
+    fn test_edge_cases() {
+        // 2^0 % 10 = 1
+        assert_eq!(mod_pow_2(0, 10), 1, "Failed exponent 0");
+
+        // 2^10 % 1 = 0
+        assert_eq!(mod_pow_2(10, 1), 0, "Failed modulus 1");
+
+        // 2^1 % 10 = 2
+        assert_eq!(mod_pow_2(1, 10), 2, "Failed exponent 1");
+
+        // 2^3 % 5 = 8 % 5 = 3
+        assert_eq!(mod_pow_2(3, 5), 3, "Failed small calc");
+    }
+
+    #[test]
+    fn test_boundaries_and_overflow() {
+        assert_eq!(mod_pow_2(2, u64::MAX), 4);
+
+        assert_eq!(mod_pow_2(u64::MAX, 3), 2);
+
+        assert_eq!(mod_pow_2(5, 32), 0);
+    }
+
+    #[test]
+    fn test_fuzzing_against_oracle() {
+        let mut rng = thread_rng();
+        for _ in 0..1_000_000 {
+            let exp: u64 = rng.gen();
+            let mod_val: u64 = rng.gen();
+
+            let mod_val = if mod_val == 0 { 1 } else { mod_val };
+
+            let expected = oracle_mod_pow_2(exp, mod_val);
+            let actual = mod_pow_2(exp, mod_val);
+
+            assert_eq!(
+                actual, expected,
+                "Mismatch! 2^{} % {} => Ours: {}, Oracle: {}",
+                exp, mod_val, actual, expected
+            );
+        }
+    }
+
+    #[test]
+    fn test_uniformity_par_generate_oblivious_pseudo_random_unsigned_custom_range2() {
+        let num_blocks = 8;
+
+        let sample_count: usize = 1_000;
+
+        let p_value_limit: f64 = 0.000_1;
+
+        let (cks, sks) = gen_keys_radix(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128, num_blocks);
+
+        let excluded_upper_bound = 3;
+
+        test_uniformity(sample_count, p_value_limit, excluded_upper_bound, &|seed| {
+            let img = sks.par_generate_oblivious_pseudo_random_unsigned_custom_range2(
+                Seed(seed as u128),
+                excluded_upper_bound,
+                num_blocks as u64,
+                Some(2_f64.powi(-32)),
+            );
+
+            cks.decrypt(&img)
+        });
+    }
+}