chore: update CODEOWNERS file

Lru required version update following caro audit

Signed-off-by: Baptiste Roux <baptiste.roux@zama.ai>

.cargo/audit.toml

@@ -2,6 +2,8 @@
 ignore = [
     # Ignoring unmaintained 'paste' advisory as it is a widely used, low-risk build dependency.
     "RUSTSEC-2024-0436",
+    # Ignoring unmaintained 'bincode' crate. Getting rid of it would be too complex on the short term.
+    "RUSTSEC-2025-0141",
 ]
 
 [output]

.github/actions/gpu_setup/action.yml

@@ -23,6 +23,8 @@ runs:
         echo "${CMAKE_SCRIPT_SHA} cmake-${CMAKE_VERSION}-linux-x86_64.sh" > checksum
         sha256sum -c checksum
         sudo bash cmake-"${CMAKE_VERSION}"-linux-x86_64.sh --skip-license --prefix=/usr/ --exclude-subdir
+        sudo apt-get clean
+        sudo rm -rf /var/lib/apt/lists/*
         sudo apt update
         sudo apt remove -y unattended-upgrades
         sudo apt install -y cmake-format libclang-dev

.gitignore

@@ -10,6 +10,7 @@ target/
 **/*.rmeta
 **/Cargo.lock
 **/*.bin
+**/.DS_Store
 
 # Some of our bench outputs
 /tfhe/benchmarks_parameters

CODEOWNERS

@@ -11,7 +11,7 @@
 /tfhe/src/core_crypto/gpu               @agnesLeroy
 /tfhe/src/core_crypto/hpu               @zama-ai/hardware
 
-/tfhe/src/shortint/                     @mayeul-zama
+/tfhe/src/shortint/                     @mayeul-zama @nsarlin-zama
 
 /tfhe/src/integer/                      @tmontaigu
 /tfhe/src/integer/gpu                   @agnesLeroy
@@ -19,8 +19,12 @@
 
 /tfhe/src/high_level_api/               @tmontaigu
 
+/tfhe-zk-pok/                           @nsarlin-zama
+
 /tfhe-benchmark/                        @soonum
 
+/utils/                                 @nsarlin-zama
+
 /Makefile                               @IceTDrinker @soonum
 
 /mockups/tfhe-hpu-mockup                @zama-ai/hardware

Cargo.toml

@@ -36,6 +36,7 @@ rayon = "1.11"
 serde = { version = "1.0", default-features = false }
 wasm-bindgen = "0.2.101"
 getrandom = "0.2.8"
+# The project maintainers consider that this is the last version of the 1.3 branch, any newer version should not be trusted
 bincode = "=1.3.3"
 
 [profile.bench]

backends/tfhe-hpu-backend/Cargo.toml

@@ -40,7 +40,7 @@ rand = "0.8.5"
 regex = "1.10.4"
 bitflags = { version = "2.5.0", features = ["serde"] }
 itertools = "0.11.0"
-lru = "0.12.3"
+lru = "0.16.3"
 bitfield-struct = "0.10.0"
 crossbeam = { version = "0.8.4", features = ["crossbeam-queue"] }
 rayon = { workspace = true }

backends/tfhe-hpu-backend/config_store/v80_archives/psi64.hpu

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:35cc06547a23b862ab9829351d74d944e60ea9dad3ecf593d15f0ce8445d145e
-size 81710610
+oid sha256:934c8131c12010dc837f6a2af5111b83f8f5d42f10485e9b3b971edb24c467f8
+size 82201876

backends/tfhe-hpu-backend/src/fw/program.rs

@@ -160,9 +160,9 @@ impl ProgramInner {
             .filter(|(_, var)| var.is_none())
             .map(|(rid, _)| *rid)
             .collect::<Vec<_>>();
-        demote_order
-            .into_iter()
-            .for_each(|rid| self.regs.demote(&rid));
+        demote_order.into_iter().for_each(|rid| {
+            self.regs.demote(&rid);
+        });
     }
 
     /// Release register entry
@@ -179,7 +179,7 @@ impl ProgramInner {
 
     /// Notify register access to update LRU state
     pub(crate) fn reg_access(&mut self, rid: asm::RegId) {
-        self.regs.promote(&rid)
+        self.regs.promote(&rid);
     }
 
     /// Retrieved least-recent-used heap entry
@@ -220,9 +220,9 @@ impl ProgramInner {
                     .filter(|(_mid, var)| var.is_none())
                     .map(|(mid, _)| *mid)
                     .collect::<Vec<_>>();
-                demote_order
-                    .into_iter()
-                    .for_each(|mid| self.heap.demote(&mid));
+                demote_order.into_iter().for_each(|mid| {
+                    self.heap.demote(&mid);
+                });
             }
             _ => { /*Only release Heap slot*/ }
         }
@@ -231,7 +231,9 @@ impl ProgramInner {
     /// Notify heap access to update LRU state
     pub(crate) fn heap_access(&mut self, mid: asm::MemId) {
         match mid {
-            asm::MemId::Heap { .. } => self.heap.promote(&mid),
+            asm::MemId::Heap { .. } => {
+                self.heap.promote(&mid);
+            }
             _ => { /* Do Nothing slot do not below to heap*/ }
         }
     }

tfhe-benchmark/.gitignore

@@ -0,0 +1 @@
+benchmarks_parameters/*

tfhe-benchmark/benches/high_level_api/bench.rs

@@ -2,7 +2,9 @@ use benchmark::utilities::{
     hlapi_throughput_num_ops, write_to_json, BenchmarkType, BitSizesSet, EnvConfig, OperatorType,
 };
 use criterion::{black_box, Criterion, Throughput};
+use oprf::oprf_any_range2;
 use rand::prelude::*;
+use rayon::prelude::*;
 use std::marker::PhantomData;
 use std::ops::*;
 use tfhe::core_crypto::prelude::Numeric;
@@ -11,34 +13,42 @@ use tfhe::keycache::NamedParam;
 use tfhe::named::Named;
 use tfhe::prelude::*;
 use tfhe::{
-    ClientKey, CompressedServerKey, FheIntegerType, FheUint10, FheUint12, FheUint128, FheUint14,
-    FheUint16, FheUint2, FheUint32, FheUint4, FheUint6, FheUint64, FheUint8, FheUintId, IntegerId,
-    KVStore,
+    ClientKey, CompressedServerKey, FheIntegerType, FheUint, FheUint10, FheUint12, FheUint128,
+    FheUint14, FheUint16, FheUint2, FheUint32, FheUint4, FheUint6, FheUint64, FheUint8, FheUintId,
+    IntegerId, KVStore,
 };
 
-use rayon::prelude::*;
+mod oprf;
 
-fn bench_fhe_type<FheType>(
+trait BenchWait {
+    fn wait_bench(&self);
+}
+
+impl<Id: FheUintId> BenchWait for FheUint<Id> {
+    fn wait_bench(&self) {
+        self.wait()
+    }
+}
+
+impl<T1: FheWait, T2> BenchWait for (T1, T2) {
+    fn wait_bench(&self) {
+        self.0.wait()
+    }
+}
+
+fn bench_fhe_type_op<FheType, F, R>(
     c: &mut Criterion,
     client_key: &ClientKey,
     type_name: &str,
     bit_size: usize,
+    display_name: &str,
+    func_name: &str,
+    func: F,
 ) where
+    F: Fn(&FheType, &FheType) -> R,
+    R: BenchWait,
     FheType: FheEncrypt<u128, ClientKey>,
     FheType: FheWait,
-    for<'a> &'a FheType: Add<&'a FheType, Output = FheType>
-        + Sub<&'a FheType, Output = FheType>
-        + Mul<&'a FheType, Output = FheType>
-        + BitAnd<&'a FheType, Output = FheType>
-        + BitOr<&'a FheType, Output = FheType>
-        + BitXor<&'a FheType, Output = FheType>
-        + Shl<&'a FheType, Output = FheType>
-        + Shr<&'a FheType, Output = FheType>
-        + RotateLeft<&'a FheType, Output = FheType>
-        + RotateRight<&'a FheType, Output = FheType>
-        + OverflowingAdd<&'a FheType, Output = FheType>
-        + OverflowingSub<&'a FheType, Output = FheType>,
-    for<'a> FheType: FheMin<&'a FheType, Output = FheType> + FheMax<&'a FheType, Output = FheType>,
 {
     let mut bench_group = c.benchmark_group(type_name);
     let mut bench_prefix = "hlapi".to_string();
@@ -71,170 +81,90 @@ fn bench_fhe_type<FheType>(
     let lhs = FheType::encrypt(rng.gen(), client_key);
     let rhs = FheType::encrypt(rng.gen(), client_key);
 
-    let mut bench_id;
+    let bench_id = format!("{bench_prefix}::{func_name}::{param_name}::{type_name}");
 
-    bench_id = format!("{bench_prefix}::add::{param_name}::{type_name}");
     bench_group.bench_function(&bench_id, |b| {
         b.iter(|| {
-            let res = &lhs + &rhs;
-            res.wait();
+            let res = func(&lhs, &rhs);
+            res.wait_bench();
             black_box(res)
         })
     });
-    write_record(bench_id, "add");
-
-    bench_id = format!("{bench_prefix}::overflowing_add::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let (res, flag) = lhs.overflowing_add(&rhs);
-            res.wait();
-            black_box((res, flag))
-        })
-    });
-    write_record(bench_id, "overflowing_add");
-
-    bench_id = format!("{bench_prefix}::overflowing_sub::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let (res, flag) = lhs.overflowing_sub(&rhs);
-            res.wait();
-            black_box((res, flag))
-        })
-    });
-    write_record(bench_id, "overflowing_sub");
-
-    bench_id = format!("{bench_prefix}::sub::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs - &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "sub");
-
-    bench_id = format!("{bench_prefix}::mul::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs * &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "mul");
-
-    bench_id = format!("{bench_prefix}::bitand::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs & &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "bitand");
-
-    bench_id = format!("{bench_prefix}::bitor::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs | &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "bitor");
-
-    bench_id = format!("{bench_prefix}::bitxor::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs ^ &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "bitxor");
-
-    bench_id = format!("{bench_prefix}::left_shift::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs << &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "left_shift");
-
-    bench_id = format!("{bench_prefix}::right_shift::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = &lhs >> &rhs;
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "right_shift");
-
-    bench_id = format!("{bench_prefix}::left_rotate::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = (&lhs).rotate_left(&rhs);
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "left_rotate");
-
-    bench_id = format!("{bench_prefix}::right_rotate::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = (&lhs).rotate_right(&rhs);
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "right_rotate");
-
-    bench_id = format!("{bench_prefix}::min::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = lhs.min(&rhs);
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "min");
-
-    bench_id = format!("{bench_prefix}::max::{param_name}::{type_name}");
-    bench_group.bench_function(&bench_id, |b| {
-        b.iter(|| {
-            let res = lhs.max(&rhs);
-            res.wait();
-            black_box(res)
-        })
-    });
-    write_record(bench_id, "max");
+    write_record(bench_id, display_name);
 }
 
-macro_rules! bench_type {
-    ($fhe_type:ident) => {
+macro_rules! bench_type_op (
+    (type_name: $fhe_type:ident, display_name: $display_name:literal, operation: $op:ident) => {
         ::paste::paste! {
-            fn [<bench_ $fhe_type:snake>](c: &mut Criterion, cks: &ClientKey) {
-                bench_fhe_type::<$fhe_type>(c, cks, stringify!($fhe_type), $fhe_type::num_bits());
+            fn [<bench_ $fhe_type:snake _ $op>](c: &mut Criterion, cks: &ClientKey) {
+                bench_fhe_type_op::<$fhe_type, _, _>(
+                    c,
+                    cks,
+                    stringify!($fhe_type),
+                    $fhe_type::num_bits(),
+                    $display_name,
+                    stringify!($op),
+                    |lhs, rhs| lhs.$op(rhs)
+                );
             }
         }
     };
+);
+
+macro_rules! generate_typed_benches {
+    ($fhe_type:ident) => {
+        bench_type_op!(type_name: $fhe_type, display_name: "add", operation: add);
+        bench_type_op!(type_name: $fhe_type, display_name: "overflowing_add", operation: overflowing_add);
+        bench_type_op!(type_name: $fhe_type, display_name: "sub", operation: sub);
+        bench_type_op!(type_name: $fhe_type, display_name: "overflowing_sub", operation: overflowing_sub);
+        bench_type_op!(type_name: $fhe_type, display_name: "mul", operation: mul);
+        bench_type_op!(type_name: $fhe_type, display_name: "bitand", operation: bitand);
+        bench_type_op!(type_name: $fhe_type, display_name: "bitor", operation: bitor);
+        bench_type_op!(type_name: $fhe_type, display_name: "bitxor", operation: bitxor);
+        bench_type_op!(type_name: $fhe_type, display_name: "left_shift", operation: shl);
+        bench_type_op!(type_name: $fhe_type, display_name: "right_shift", operation: shr);
+        bench_type_op!(type_name: $fhe_type, display_name: "left_rotate", operation: rotate_left);
+        bench_type_op!(type_name: $fhe_type, display_name: "right_rotate", operation: rotate_right);
+        bench_type_op!(type_name: $fhe_type, display_name: "min", operation: min);
+        bench_type_op!(type_name: $fhe_type, display_name: "max", operation: max);
+    };
 }
 
-bench_type!(FheUint2);
-bench_type!(FheUint4);
-bench_type!(FheUint6);
-bench_type!(FheUint8);
-bench_type!(FheUint10);
-bench_type!(FheUint12);
-bench_type!(FheUint14);
-bench_type!(FheUint16);
-bench_type!(FheUint32);
-bench_type!(FheUint64);
-bench_type!(FheUint128);
+// Generate benches for all FheUint types
+generate_typed_benches!(FheUint2);
+generate_typed_benches!(FheUint4);
+generate_typed_benches!(FheUint6);
+generate_typed_benches!(FheUint8);
+generate_typed_benches!(FheUint10);
+generate_typed_benches!(FheUint12);
+generate_typed_benches!(FheUint14);
+generate_typed_benches!(FheUint16);
+generate_typed_benches!(FheUint32);
+generate_typed_benches!(FheUint64);
+generate_typed_benches!(FheUint128);
+
+macro_rules! run_benches {
+    ($c:expr, $cks:expr, $($fhe_type:ident),+ $(,)?) => {
+        $(
+            ::paste::paste! {
+                [<bench_ $fhe_type:snake _add>]($c, $cks);
+                [<bench_ $fhe_type:snake _overflowing_add>]($c, $cks);
+                [<bench_ $fhe_type:snake _sub>]($c, $cks);
+                [<bench_ $fhe_type:snake _overflowing_sub>]($c, $cks);
+                [<bench_ $fhe_type:snake _mul>]($c, $cks);
+                [<bench_ $fhe_type:snake _bitand>]($c, $cks);
+                [<bench_ $fhe_type:snake _bitor>]($c, $cks);
+                [<bench_ $fhe_type:snake _bitxor>]($c, $cks);
+                [<bench_ $fhe_type:snake _shl>]($c, $cks);
+                [<bench_ $fhe_type:snake _shr>]($c, $cks);
+                [<bench_ $fhe_type:snake _rotate_left>]($c, $cks);
+                [<bench_ $fhe_type:snake _rotate_right>]($c, $cks);
+                [<bench_ $fhe_type:snake _min>]($c, $cks);
+                [<bench_ $fhe_type:snake _max>]($c, $cks);
+            }
+        )+
+    };
+}
 
 trait TypeDisplay {
     fn fmt(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -444,7 +374,7 @@ fn main() {
 
     match env_config.bit_sizes_set {
         BitSizesSet::Fast => {
-            bench_fhe_uint64(&mut c, &cks);
+            run_benches!(&mut c, &cks, FheUint64);
 
             // KVStore Benches
             if benched_device == tfhe::Device::Cpu {
@@ -452,17 +382,11 @@ fn main() {
             }
         }
         _ => {
-            bench_fhe_uint2(&mut c, &cks);
-            bench_fhe_uint4(&mut c, &cks);
-            bench_fhe_uint6(&mut c, &cks);
-            bench_fhe_uint8(&mut c, &cks);
-            bench_fhe_uint10(&mut c, &cks);
-            bench_fhe_uint12(&mut c, &cks);
-            bench_fhe_uint14(&mut c, &cks);
-            bench_fhe_uint16(&mut c, &cks);
-            bench_fhe_uint32(&mut c, &cks);
-            bench_fhe_uint64(&mut c, &cks);
-            bench_fhe_uint128(&mut c, &cks);
+            // Call all benchmarks for all types
+            run_benches!(
+                &mut c, &cks, FheUint2, FheUint4, FheUint6, FheUint8, FheUint10, FheUint12,
+                FheUint14, FheUint16, FheUint32, FheUint64, FheUint128
+            );
 
             // KVStore Benches
             if benched_device == tfhe::Device::Cpu {
@@ -481,5 +405,7 @@ fn main() {
         }
     }
 
+    oprf_any_range2();
+
     c.final_summary();
 }

tfhe-benchmark/benches/high_level_api/oprf.rs

@@ -0,0 +1,44 @@
+use benchmark::params_aliases::BENCH_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+use criterion::{black_box, criterion_group, Criterion};
+use std::num::NonZeroU64;
+use tfhe::{set_server_key, ClientKey, ConfigBuilder, FheUint64, RangeForRandom, Seed, ServerKey};
+
+pub fn oprf_any_range(c: &mut Criterion) {
+    let bench_name = "hlapi::oprf_any_range";
+
+    let mut bench_group = c.benchmark_group(bench_name);
+    bench_group
+        .sample_size(15)
+        .measurement_time(std::time::Duration::from_secs(30));
+
+    let param = BENCH_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+
+    let config = ConfigBuilder::with_custom_parameters(param).build();
+    let cks = ClientKey::generate(config);
+    let sks = ServerKey::new(&cks);
+
+    rayon::broadcast(|_| set_server_key(sks.clone()));
+    set_server_key(sks);
+
+    for excluded_upper_bound in [3, 52] {
+        let range = RangeForRandom::new_from_excluded_upper_bound(
+            NonZeroU64::new(excluded_upper_bound).unwrap(),
+        );
+
+        let bench_id_oprf = format!("{bench_name}::bound_{excluded_upper_bound}");
+
+        bench_group.bench_function(&bench_id_oprf, |b| {
+            b.iter(|| {
+                _ = black_box(FheUint64::generate_oblivious_pseudo_random_custom_range(
+                    Seed(0),
+                    &range,
+                    None,
+                ));
+            })
+        });
+    }
+
+    bench_group.finish()
+}
+
+criterion_group!(oprf_any_range2, oprf_any_range);

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/docs/fhe-computation/advanced-features/encrypted-prf.md

@@ -2,14 +2,30 @@
 
 This document explains the mechanism and steps to generate an oblivious encrypted random value using only server keys.
 
-The goal is to give to the server the possibility to generate a random value, which will be obtained in an encrypted format and will remain unknown to the server. The implementation is based on [this article](https://eprint.iacr.org/2024/665).
+The goal is to give to the server the possibility to generate a random value, which will be obtained in an encrypted format and will remain unknown to the server.
 
-This is possible through two methods on `FheUint` and `FheInt`: 
+The main method for this is `FheUint::generate_oblivious_pseudo_random_custom_range` which returns an integer in the given range.
+Currently the range can only be in the form `[0, excluded_upper_bound[` with any `excluded_upper_bound` in `[1, 2^64[`
+It follows a distribution close to the uniform.
+
+This function guarantees the norm-1 distance (defined as ∆(P,Q) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)|)
+between the actual distribution and the target uniform distribution will be below the `max_distance` argument (which must be in ]0, 1[).
+The higher the distance, the more dissimilar the actual distribution is from the target uniform distribution.
+
+The default value for `max_distance` is `2^-128` 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.
+
+If the range is a power of 2, the distribution is uniform (for any `max_distance`) and the cost is smaller.
+
+
+For powers of 2 specifically there are two methods on `FheUint` and `FheInt` (based on [this article](https://eprint.iacr.org/2024/665)): 
 - `generate_oblivious_pseudo_random` which return an integer taken uniformly in the full integer range (`[0; 2^N[` for a `FheUintN` and `[-2^(N-1); 2^(N-1)[` for a `FheIntN`).
 - `generate_oblivious_pseudo_random_bounded` which return an integer taken uniformly in `[0; 2^random_bits_count[`. For a `FheUintN`, we must have  `random_bits_count <= N`. For a `FheIntN`, we must have  `random_bits_count <= N - 1`.
 
-Both methods functions take a seed `Seed` as input, which could be any `u128` value.
-They both rely on the use of the usual server key.
+
+These method functions take a seed `Seed` as input, which could be any `u128` value.
+They rely on the use of the usual server key.
 The output is reproducible, i.e., the function is deterministic from the inputs: assuming the same hardware, seed and server key, this function outputs the same random encrypted value.
 
 
@@ -18,7 +34,8 @@ Here is an example of the usage:
 
 ```rust
 use tfhe::prelude::FheDecrypt;
-use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, FheInt8, Seed};
+use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, FheInt8, RangeForRandom, Seed};
+use std::num::NonZeroU64;
 
 pub fn main() {
     let config = ConfigBuilder::default().build();
@@ -26,23 +43,30 @@ pub fn main() {
 
     set_server_key(server_key);
 
-    let random_bits_count = 3;
-
-    let ct_res = FheUint8::generate_oblivious_pseudo_random(Seed(0));
+    let excluded_upper_bound = NonZeroU64::new(3).unwrap();
+    let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound);
 
+    // in [0, excluded_upper_bound[ = {0, 1, 2}
+    let ct_res = FheUint8::generate_oblivious_pseudo_random_custom_range(Seed(0), &range, None);
     let dec_result: u8 = ct_res.decrypt(&client_key);
 
-    let ct_res = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count);
+    let random_bits_count = 3;
 
+    // in [0, 2^8[
+    let ct_res = FheUint8::generate_oblivious_pseudo_random(Seed(0));
+    let dec_result: u8 = ct_res.decrypt(&client_key);
+
+    // in [0, 2^random_bits_count[ = [0, 8[
+    let ct_res = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count);
     let dec_result: u8 = ct_res.decrypt(&client_key);
     assert!(dec_result < (1 << random_bits_count));
 
+    // in [-2^7, 2^7[
     let ct_res = FheInt8::generate_oblivious_pseudo_random(Seed(0));
-    
     let dec_result: i8 = ct_res.decrypt(&client_key);
     
+    // in [0, 2^random_bits_count[ = [0, 8[
     let ct_res = FheInt8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count);
-
     let dec_result: i8 = ct_res.decrypt(&client_key);
     assert!(dec_result < (1 << random_bits_count));
 }

tfhe/src/core_crypto/commons/math/random/tests.rs

@@ -540,10 +540,12 @@ pub fn sup_diff(cumulative_bins: &[u64], theoretical_cdf: &[f64]) -> f64 {
         .iter()
         .copied()
         .zip_eq(theoretical_cdf.iter().copied())
-        .map(|(x, theoretical_cdf)| {
+        .enumerate()
+        .map(|(i, (x, theoretical_cdf))| {
             let empirical_cdf = x as f64 / number_of_samples as f64;
 
-            if theoretical_cdf == 1.0 {
+            if i == cumulative_bins.len() - 1 {
+                assert_eq!(theoretical_cdf, 1.0);
                 assert_eq!(empirical_cdf, 1.0);
             }
 

tfhe/src/high_level_api/integers/oprf.rs

@@ -4,7 +4,9 @@ use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::re_randomization::ReRandomizationMetadata;
 #[cfg(feature = "gpu")]
 use crate::integer::gpu::ciphertext::{CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext};
+use crate::shortint::MessageModulus;
 use crate::{FheInt, Seed};
+use std::num::NonZeroU64;
 
 impl<Id: FheUintId> FheUint<Id> {
     /// Generates an encrypted unsigned integer
@@ -92,7 +94,7 @@ impl<Id: FheUintId> FheUint<Id> {
             }
         })
     }
-    /// Generates an encrypted `num_block` blocks unsigned integer
+    /// Generates an encrypted unsigned integer
     /// taken uniformly in `[0, 2^random_bits_count[` using the given seed.
     /// The encrypted value is oblivious to the server.
     /// It can be useful to make server random generation deterministic.
@@ -150,6 +152,103 @@ impl<Id: FheUintId> FheUint<Id> {
             }
         })
     }
+
+    /// Generates an encrypted unsigned integer
+    /// taken almost uniformly in the given range using the given seed.
+    /// Currently the range can only be in the form `[0, excluded_upper_bound[`
+    /// with any `excluded_upper_bound` in `[1, 2^64[`.
+    ///
+    /// The encrypted value is oblivious to the server.
+    /// It can be useful to make server random generation deterministic.
+    ///
+    /// This function guarantees the the norm-1 distance
+    /// (defined as ∆(P,Q) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)|)
+    /// between the actual distribution and the target uniform distribution
+    /// will be below the `max_distance` argument (which must be in ]0, 1[).
+    /// The higher the distance, the more dissimilar the actual distribution is
+    /// from the target uniform distribution.
+    ///
+    /// The default value for `max_distance` is `2^-128` 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.
+    ///
+    /// If the range is a power of 2, the distribution is uniform (for any `max_distance`) and
+    /// the cost is smaller.
+    ///
+    /// ```rust
+    /// use std::num::NonZeroU64;
+    /// use tfhe::prelude::FheDecrypt;
+    /// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, RangeForRandom, Seed};
+    ///
+    /// let config = ConfigBuilder::default().build();
+    /// let (client_key, server_key) = generate_keys(config);
+    ///
+    /// set_server_key(server_key);
+    ///
+    /// let excluded_upper_bound = NonZeroU64::new(3).unwrap();
+    ///
+    /// let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound);
+    ///
+    /// let ct_res = FheUint8::generate_oblivious_pseudo_random_custom_range(Seed(0), &range, None);
+    ///
+    /// let dec_result: u16 = ct_res.decrypt(&client_key);
+    /// assert!(dec_result < excluded_upper_bound.get() as u16);
+    /// ```
+    pub fn generate_oblivious_pseudo_random_custom_range(
+        seed: Seed,
+        range: &RangeForRandom,
+        max_distance: Option<f64>,
+    ) -> Self {
+        let excluded_upper_bound = range.excluded_upper_bound;
+
+        if excluded_upper_bound.is_power_of_two() {
+            let random_bits_count = excluded_upper_bound.ilog2() as u64;
+
+            Self::generate_oblivious_pseudo_random_bounded(seed, random_bits_count)
+        } else {
+            let max_distance = max_distance.unwrap_or_else(|| 2_f64.powi(-128));
+
+            assert!(
+                0_f64 < max_distance && max_distance < 1_f64,
+                "max_distance (={max_distance}) should be in ]0, 1["
+            );
+
+            global_state::with_internal_keys(|key| match key {
+                InternalServerKey::Cpu(key) => {
+                    let message_modulus = key.message_modulus();
+
+                    let num_input_random_bits = num_input_random_bits_for_max_distance(
+                        excluded_upper_bound,
+                        max_distance,
+                        message_modulus,
+                    );
+
+                    let num_blocks_output = Id::num_blocks(key.message_modulus()) as u64;
+
+                    let ct = key
+                        .pbs_key()
+                        .par_generate_oblivious_pseudo_random_unsigned_custom_range(
+                            seed,
+                            num_input_random_bits,
+                            excluded_upper_bound,
+                            num_blocks_output,
+                        );
+
+                    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
     ///
@@ -273,7 +372,7 @@ impl<Id: FheIntId> FheInt<Id> {
             }
         })
     }
-    /// Generates an encrypted `num_block` blocks signed integer
+    /// Generates an encrypted signed integer
     /// taken uniformly in `[0, 2^random_bits_count[` using the given seed.
     /// The encrypted value is oblivious to the server.
     /// It can be useful to make server random generation deterministic.
@@ -367,10 +466,350 @@ impl<Id: FheIntId> FheInt<Id> {
     }
 }
 
+pub struct RangeForRandom {
+    excluded_upper_bound: NonZeroU64,
+}
+
+impl RangeForRandom {
+    pub fn new_from_excluded_upper_bound(excluded_upper_bound: NonZeroU64) -> Self {
+        Self {
+            excluded_upper_bound,
+        }
+    }
+}
+
+fn num_input_random_bits_for_max_distance(
+    excluded_upper_bound: NonZeroU64,
+    max_distance: f64,
+    message_modulus: MessageModulus,
+) -> u64 {
+    assert!(message_modulus.0.is_power_of_two());
+    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 distance = distance(excluded_upper_bound.get(), random_bit_count);
+
+        if distance < max_distance {
+            break random_block_count;
+        }
+
+        random_block_count += 1;
+    };
+
+    random_block_count * log_message_modulus
+}
+
+fn distance(excluded_upper_bound: u64, random_bit_count: u64) -> f64 {
+    let remainder = mod_pow_2(random_bit_count, excluded_upper_bound);
+
+    remainder as f64 * (excluded_upper_bound - remainder) as f64
+        / (2_f64.powi(random_bit_count as i32) * excluded_upper_bound as f64)
+}
+
+// Computes 2^exponent % modulus
+fn mod_pow_2(exponent: u64, modulus: u64) -> u64 {
+    assert_ne!(modulus, 0);
+
+    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
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use crate::integer::server_key::radix_parallel::tests_unsigned::test_oprf::{
+        oprf_density_function, p_value_upper_bound_oprf_almost_uniformity_from_values,
+        probability_density_function_from_density,
+    };
+    use crate::prelude::FheDecrypt;
+    use crate::shortint::oprf::test::test_uniformity;
+    use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128;
+    use crate::{generate_keys, set_server_key, ClientKey, ConfigBuilder, FheUint8, Seed};
+    use num_bigint::BigUint;
+    use rand::{thread_rng, Rng};
+    use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
+    // 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 {
+        assert_ne!(modulus, 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_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^{exp} % {mod_val} => Ours: {actual}, Oracle: {expected}",
+            );
+        }
+    }
+
+    #[test]
+    fn test_distance_with_uniform() {
+        for excluded_upper_bound in 1..20 {
+            for num_input_random_bits in 0..20 {
+                let density = oprf_density_function(excluded_upper_bound, num_input_random_bits);
+
+                let theoretical_pdf = probability_density_function_from_density(&density);
+
+                let p_uniform = 1. / excluded_upper_bound as f64;
+
+                let actual_distance: f64 = 1. / 2.
+                    * theoretical_pdf
+                        .iter()
+                        .map(|p| (*p - p_uniform).abs())
+                        .sum::<f64>();
+
+                let theoretical_distance = distance(excluded_upper_bound, num_input_random_bits);
+
+                assert!(
+                    (theoretical_distance - actual_distance).abs()
+                        <= theoretical_distance / 1_000_000.,
+                    "{theoretical_distance} != {actual_distance}"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_uniformity_scalar_mul_shift() {
+        let max_distance = 2_f64.powi(-20);
+
+        let message_modulus = MessageModulus(4);
+
+        let excluded_upper_bound = 3;
+
+        let num_input_random_bits = num_input_random_bits_for_max_distance(
+            NonZeroU64::new(excluded_upper_bound).unwrap(),
+            max_distance,
+            message_modulus,
+        );
+
+        let sample_count: usize = 10_000_000;
+
+        let p_value_limit: f64 = 0.001;
+
+        // The distribution is not exactly uniform
+        // This check ensures than with the given low max_distance,
+        // the distribution is indistinguishable from the uniform with at the given sample count
+        test_uniformity(sample_count, p_value_limit, excluded_upper_bound, |_seed| {
+            oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits)
+        });
+    }
+
+    fn oprf_clear_equivalent(excluded_upper_bound: u64, num_input_random_bits: u64) -> u64 {
+        let random_input_upper_bound = 1 << num_input_random_bits;
+
+        let random_input = thread_rng().gen_range(0..random_input_upper_bound);
+
+        (random_input * excluded_upper_bound) >> num_input_random_bits
+    }
+
+    #[test]
+    fn test_uniformity_generate_oblivious_pseudo_random_custom_range() {
+        let base_sample_count: usize = 10_000;
+
+        let p_value_limit: f64 = 0.001;
+
+        let params = PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128;
+        let config = ConfigBuilder::with_custom_parameters(params).build();
+
+        let (cks, sks) = generate_keys(config);
+        rayon::broadcast(|_| set_server_key(sks.clone()));
+
+        let message_modulus = params.message_modulus;
+
+        // [0.7, 0.1] for `max_distance` chosen to have `num_input_random_bits` be [2, 4]
+        // for any of the listed `excluded_upper_bound`
+        for (expected_num_input_random_bits, max_distance, excluded_upper_bounds) in
+            [(2, 0.7, [3, 5, 6, 7]), (4, 0.1, [3, 5, 6, 7])]
+        {
+            for excluded_upper_bound in excluded_upper_bounds {
+                let sample_count = base_sample_count * excluded_upper_bound as usize;
+
+                let excluded_upper_bound = NonZeroU64::new(excluded_upper_bound).unwrap();
+
+                let num_input_random_bits = num_input_random_bits_for_max_distance(
+                    excluded_upper_bound,
+                    max_distance,
+                    message_modulus,
+                );
+
+                assert_eq!(num_input_random_bits, expected_num_input_random_bits);
+
+                test_uniformity_generate_oblivious_pseudo_random_custom_range2(
+                    sample_count,
+                    p_value_limit,
+                    message_modulus,
+                    &cks,
+                    excluded_upper_bound,
+                    max_distance,
+                );
+            }
+        }
+    }
+
+    fn test_uniformity_generate_oblivious_pseudo_random_custom_range2(
+        sample_count: usize,
+        p_value_limit: f64,
+        message_modulus: MessageModulus,
+        cks: &ClientKey,
+        excluded_upper_bound: NonZeroU64,
+        max_distance: f64,
+    ) {
+        let num_input_random_bits = num_input_random_bits_for_max_distance(
+            excluded_upper_bound,
+            max_distance,
+            message_modulus,
+        );
+
+        let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound);
+
+        let real_values: Vec<u64> = (0..sample_count)
+            .into_par_iter()
+            .map(|_| {
+                let img = FheUint8::generate_oblivious_pseudo_random_custom_range(
+                    Seed(rand::thread_rng().gen::<u128>()),
+                    &range,
+                    Some(max_distance),
+                );
+
+                img.decrypt(cks)
+            })
+            .collect();
+
+        let excluded_upper_bound = excluded_upper_bound.get();
+
+        let uniform_values: Vec<u64> = (0..sample_count)
+            .into_par_iter()
+            .map(|_| thread_rng().gen_range(0..excluded_upper_bound))
+            .collect();
+
+        let clear_oprf_value_lower_num_input_random_bits = (0..sample_count)
+            .into_par_iter()
+            .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits - 1))
+            .collect();
+
+        let clear_oprf_value_same_num_input_random_bits = (0..sample_count)
+            .into_par_iter()
+            .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits))
+            .collect();
+
+        let clear_oprf_value_higher_num_input_random_bits = (0..sample_count)
+            .into_par_iter()
+            .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits + 1))
+            .collect();
+
+        for (values, should_have_low_p_value) in [
+            (&real_values, false),
+            // to test that the same distribution passes
+            (&clear_oprf_value_same_num_input_random_bits, false),
+            // to test that other distribution don't pass
+            // (makes sure the test is statistically powerful)
+            (&uniform_values, true),
+            (&clear_oprf_value_lower_num_input_random_bits, true),
+            (&clear_oprf_value_higher_num_input_random_bits, true),
+        ] {
+            let p_value_upper_bound = p_value_upper_bound_oprf_almost_uniformity_from_values(
+                values,
+                num_input_random_bits,
+                excluded_upper_bound,
+            );
+
+            println!("p_value_upper_bound: {p_value_upper_bound}");
+
+            if should_have_low_p_value {
+                assert!(
+                    p_value_upper_bound < p_value_limit,
+                    "p_value_upper_bound (={p_value_upper_bound}) expected to be smaller than {p_value_limit}"
+                );
+            } else {
+                assert!(
+                    p_value_limit < p_value_upper_bound ,
+                    "p_value_upper_bound (={p_value_upper_bound}) expected to be bigger than {p_value_limit}"
+                );
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 #[cfg(feature = "gpu")]
 #[allow(unused_imports)]
-mod test {
+mod test_gpu {
     use crate::prelude::*;
     use crate::{
         generate_keys, set_server_key, ConfigBuilder, FheInt128, FheUint32, FheUint64, GpuIndex,

tfhe/src/high_level_api/mod.rs

@@ -48,6 +48,7 @@ macro_rules! export_concrete_array_types {
 }
 
 pub use crate::core_crypto::commons::math::random::{Seed, XofSeed};
+pub use crate::high_level_api::integers::oprf::RangeForRandom;
 pub use crate::integer::server_key::MatchValues;
 use crate::{error, Error, Versionize};
 use backward_compatibility::compressed_ciphertext_list::SquashedNoiseCiphertextStateVersions;

tfhe/src/integer/oprf.rs

@@ -2,6 +2,7 @@ use super::{RadixCiphertext, ServerKey, SignedRadixCiphertext};
 use crate::core_crypto::commons::generators::DeterministicSeeder;
 use crate::core_crypto::prelude::DefaultRandomGenerator;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+use std::num::NonZeroU64;
 
 pub use tfhe_csprng::seeders::{Seed, Seeder};
 
@@ -163,6 +164,7 @@ impl ServerKey {
     /// as `num_input_random_bits`
     ///
     /// ```rust
+    /// use std::num::NonZeroU64;
     /// use tfhe::integer::gen_keys_radix;
     /// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128;
     /// use tfhe::Seed;
@@ -173,7 +175,7 @@ impl ServerKey {
     /// let (cks, sks) = gen_keys_radix(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128, size);
     ///
     /// let num_input_random_bits = 5;
-    /// let excluded_upper_bound = 3;
+    /// let excluded_upper_bound = NonZeroU64::new(3).unwrap();
     /// let num_blocks_output = 8;
     ///
     /// let ct_res = sks.par_generate_oblivious_pseudo_random_unsigned_custom_range(
@@ -186,15 +188,17 @@ impl ServerKey {
     /// // Decrypt:
     /// let dec_result: u64 = cks.decrypt(&ct_res);
     ///
-    /// assert!(dec_result < excluded_upper_bound);
+    /// assert!(dec_result < excluded_upper_bound.get());
     /// ```
     pub fn par_generate_oblivious_pseudo_random_unsigned_custom_range(
         &self,
         seed: Seed,
         num_input_random_bits: u64,
-        excluded_upper_bound: u64,
+        excluded_upper_bound: NonZeroU64,
         num_blocks_output: u64,
     ) -> RadixCiphertext {
+        let excluded_upper_bound = excluded_upper_bound.get();
+
         assert!(self.message_modulus().0.is_power_of_two());
         let message_bits_count = self.message_modulus().0.ilog2() as u64;
 

tfhe/src/integer/server_key/radix_parallel/tests_long_run/test_random_op_sequence.rs

@@ -10,6 +10,7 @@ use crate::integer::{BooleanBlock, IntegerKeyKind, RadixCiphertext, RadixClientK
 use crate::shortint::parameters::*;
 use crate::{ClientKey, CompressedServerKey, MatchValues, Seed, Tag};
 use std::cmp::{max, min};
+use std::num::NonZeroU64;
 use std::sync::Arc;
 
 create_parameterized_test!(random_op_sequence {
@@ -498,7 +499,18 @@ where
         &ServerKey::par_generate_oblivious_pseudo_random_unsigned_integer_bounded,
     );
     let oprf_custom_range_executor = OpSequenceCpuFunctionExecutor::new(
-        &ServerKey::par_generate_oblivious_pseudo_random_unsigned_custom_range,
+        &|sk: &ServerKey,
+          seed: Seed,
+          num_input_random_bits: u64,
+          excluded_upper_bound: u64,
+          num_blocks_output: u64| {
+            sk.par_generate_oblivious_pseudo_random_unsigned_custom_range(
+                seed,
+                num_input_random_bits,
+                NonZeroU64::new(excluded_upper_bound).unwrap_or(NonZeroU64::new(1).unwrap()),
+                num_blocks_output,
+            )
+        },
     );
 
     let mut oprf_ops: Vec<(OprfExecutor, String)> = vec![(

tfhe/src/integer/server_key/radix_parallel/tests_unsigned/test_oprf.rs

@@ -9,6 +9,7 @@ use crate::integer::{IntegerKeyKind, RadixCiphertext, RadixClientKey, ServerKey}
 use crate::shortint::parameters::*;
 use statrs::distribution::ContinuousCDF;
 use std::collections::HashMap;
+use std::num::NonZeroU64;
 use std::sync::Arc;
 use tfhe_csprng::seeders::Seed;
 
@@ -36,9 +37,19 @@ fn oprf_any_range_unsigned<P>(param: P)
 where
     P: Into<TestParameters>,
 {
-    let executor = CpuFunctionExecutor::new(
-        &ServerKey::par_generate_oblivious_pseudo_random_unsigned_custom_range,
-    );
+    let executor =
+        CpuFunctionExecutor::new(&|sk: &ServerKey,
+                                   seed: Seed,
+                                   num_input_random_bits: u64,
+                                   excluded_upper_bound: u64,
+                                   num_blocks_output: u64| {
+            sk.par_generate_oblivious_pseudo_random_unsigned_custom_range(
+                seed,
+                num_input_random_bits,
+                NonZeroU64::new(excluded_upper_bound).unwrap(),
+                num_blocks_output,
+            )
+        });
     oprf_any_range_test(param, executor);
 }
 
@@ -46,9 +57,19 @@ fn oprf_almost_uniformity_unsigned<P>(param: P)
 where
     P: Into<TestParameters>,
 {
-    let executor = CpuFunctionExecutor::new(
-        &ServerKey::par_generate_oblivious_pseudo_random_unsigned_custom_range,
-    );
+    let executor =
+        CpuFunctionExecutor::new(&|sk: &ServerKey,
+                                   seed: Seed,
+                                   num_input_random_bits: u64,
+                                   excluded_upper_bound: u64,
+                                   num_blocks_output: u64| {
+            sk.par_generate_oblivious_pseudo_random_unsigned_custom_range(
+                seed,
+                num_input_random_bits,
+                NonZeroU64::new(excluded_upper_bound).unwrap(),
+                num_blocks_output,
+            )
+        });
     oprf_almost_uniformity_test(param, executor);
 }
 
@@ -89,7 +110,7 @@ where
     );
 }
 
-pub fn oprf_uniformity_test<P, E>(param: P, mut executor: E)
+pub(crate) fn oprf_uniformity_test<P, E>(param: P, mut executor: E)
 where
     P: Into<TestParameters>,
     E: for<'a> FunctionExecutor<(Seed, u64, u64), RadixCiphertext>,
@@ -113,7 +134,7 @@ where
     });
 }
 
-pub fn oprf_any_range_test<P, E>(param: P, mut executor: E)
+pub(crate) fn oprf_any_range_test<P, E>(param: P, mut executor: E)
 where
     P: Into<TestParameters>,
     E: for<'a> FunctionExecutor<(Seed, u64, u64, u64), RadixCiphertext>,
@@ -149,7 +170,7 @@ where
     }
 }
 
-pub fn oprf_almost_uniformity_test<P, E>(param: P, mut executor: E)
+pub(crate) fn oprf_almost_uniformity_test<P, E>(param: P, mut executor: E)
 where
     P: Into<TestParameters>,
     E: for<'a> FunctionExecutor<(Seed, u64, u64, u64), RadixCiphertext>,
@@ -165,40 +186,70 @@ where
     let num_input_random_bits: u64 = 4;
     let num_blocks_output = 64;
     let excluded_upper_bound = 10;
-    let random_input_upper_bound = 1 << num_input_random_bits;
-
-    let mut density = vec![0_usize; excluded_upper_bound as usize];
-    for i in 0..random_input_upper_bound {
-        let index = ((i * excluded_upper_bound) as f64 / random_input_upper_bound as f64) as usize;
-        density[index] += 1;
-    }
-
-    let theoretical_pdf: Vec<f64> = density
-        .iter()
-        .map(|count| *count as f64 / random_input_upper_bound as f64)
-        .collect();
 
     let values: Vec<u64> = (0..sample_count)
         .map(|seed| {
             let img = executor.execute((
                 Seed(seed as u128),
                 num_input_random_bits,
-                excluded_upper_bound as u64,
+                excluded_upper_bound,
                 num_blocks_output,
             ));
             cks.decrypt(&img)
         })
         .collect();
 
+    let p_value_upper_bound = p_value_upper_bound_oprf_almost_uniformity_from_values(
+        &values,
+        num_input_random_bits,
+        excluded_upper_bound,
+    );
+
+    assert!(p_value_limit < p_value_upper_bound);
+}
+
+pub(crate) fn p_value_upper_bound_oprf_almost_uniformity_from_values(
+    values: &[u64],
+    num_input_random_bits: u64,
+    excluded_upper_bound: u64,
+) -> f64 {
+    let density = oprf_density_function(excluded_upper_bound, num_input_random_bits);
+
+    let theoretical_pdf = probability_density_function_from_density(&density);
+
     let mut bins = vec![0_u64; excluded_upper_bound as usize];
-    for value in values {
+    for value in values.iter().copied() {
         bins[value as usize] += 1;
     }
 
     let cumulative_bins = cumulate(&bins);
     let theoretical_cdf = cumulate(&theoretical_pdf);
     let sup_diff = sup_diff(&cumulative_bins, &theoretical_cdf);
-    let p_value_upper_bound = dkw_alpha_from_epsilon(sample_count as f64, sup_diff);
 
-    assert!(p_value_limit < p_value_upper_bound);
+    dkw_alpha_from_epsilon(values.len() as f64, sup_diff)
+}
+
+pub(crate) fn oprf_density_function(
+    excluded_upper_bound: u64,
+    num_input_random_bits: u64,
+) -> Vec<usize> {
+    let random_input_upper_bound = 1 << num_input_random_bits;
+
+    let mut density = vec![0_usize; excluded_upper_bound as usize];
+
+    for i in 0..random_input_upper_bound {
+        let output = ((i * excluded_upper_bound) >> num_input_random_bits) as usize;
+
+        density[output] += 1;
+    }
+    density
+}
+
+pub(crate) fn probability_density_function_from_density(density: &[usize]) -> Vec<f64> {
+    let total_count: usize = density.iter().copied().sum();
+
+    density
+        .iter()
+        .map(|count| *count as f64 / total_count as f64)
+        .collect()
 }

tfhe/src/shortint/oprf.rs

@@ -475,8 +475,12 @@ pub(crate) mod test {
         }
     }
 
-    pub fn test_uniformity<F>(sample_count: usize, p_value_limit: f64, distinct_values: u64, f: F)
-    where
+    pub(crate) fn test_uniformity<F>(
+        sample_count: usize,
+        p_value_limit: f64,
+        distinct_values: u64,
+        f: F,
+    ) where
         F: Sync + Fn(usize) -> u64,
     {
         let p_value = uniformity_p_value(f, sample_count, distinct_values);
@@ -487,7 +491,7 @@ pub(crate) mod test {
         );
     }
 
-    fn uniformity_p_value<F>(f: F, sample_count: usize, distinct_values: u64) -> f64
+    pub(crate) fn uniformity_p_value<F>(f: F, sample_count: usize, distinct_values: u64) -> f64
     where
         F: Sync + Fn(usize) -> u64,
     {
@@ -495,8 +499,11 @@ pub(crate) mod test {
 
         let mut values_count = HashMap::new();
 
-        for i in &values {
-            assert!(*i < distinct_values, "i {} dv{}", *i, distinct_values);
+        for i in values.iter().copied() {
+            assert!(
+                i < distinct_values,
+                "i (={i}) is supposed to be smaller than distinct_values (={distinct_values})",
+            );
 
             *values_count.entry(i).or_insert(0) += 1;
         }

tfhe/web_wasm_parallel_tests/worker.js

@@ -727,8 +727,15 @@ async function compactPublicKeyZeroKnowledgeBench() {
             serialized_size = list.safe_serialize(BigInt(10000000)).length;
           }
           const mean = timing / bench_loops;
+
+          let base_bench_str = "compact_fhe_uint_proven_encryption_";
+          let supportsThreads = await threads();
+          if (!supportsThreads) {
+            base_bench_str += "unsafe_coop_";
+          }
+
           const common_bench_str =
-            "compact_fhe_uint_proven_encryption_" +
+            base_bench_str +
             params.zk_scheme +
             "_" +
             bits_to_encrypt +