feat(c_api): add oprf

zama-ai · Jan 31, 2024 · 7b4bb6a · 7b4bb6a
1 parent fd084d5
commit 7b4bb6a
Show file tree

Hide file tree

Showing 3 changed files with 192 additions and 12 deletions.
diff --git a/tfhe/c_api_tests/test_high_level_integers.c b/tfhe/c_api_tests/test_high_level_integers.c
@@ -534,6 +534,54 @@ void test_try_decrypt_trivial(const ClientKey *client_key) {
   fhe_uint16_destroy(non_trivial);
 }
 
+void test_oprf(const ClientKey *client_key) {
+  {
+    FheUint8 *ct = NULL;
+
+    int status = generate_oblivious_pseudo_random_fhe_uint8(&ct, 0, 0);
+    assert(status == 0);
+
+    uint8_t decrypted;
+    status = fhe_uint8_decrypt(ct, client_key, &decrypted);
+    // nothing to assert here, as decrypted can be any uint8_t value
+
+    fhe_uint8_destroy(ct);
+
+    status = generate_oblivious_pseudo_random_bits_fhe_uint8(&ct, 0, 0, 2);
+    assert(status == 0);
+
+    status = fhe_uint8_decrypt(ct, client_key, &decrypted);
+    assert(status == 0);
+
+    assert(decrypted < 4);
+
+    fhe_uint8_destroy(ct);
+  }
+
+  {
+    FheInt8 *ct = NULL;
+
+    int status = generate_oblivious_pseudo_random_full_signed_range_fhe_int8(&ct, 0, 0);
+    assert(status == 0);
+
+    int8_t decrypted;
+    status = fhe_int8_decrypt(ct, client_key, &decrypted);
+    assert(status == 0);
+    // nothing to assert here, as decrypted can be any int8_t value
+
+    fhe_int8_destroy(ct);
+
+    status = generate_oblivious_pseudo_random_unsigned_fhe_int8(&ct, 0, 0, 2);
+    assert(status == 0);
+
+    status = fhe_int8_decrypt(ct, client_key, &decrypted);
+    assert(status == 0);
+
+    assert(decrypted < 4);
+    assert(decrypted >= 0);
+  }
+}
+
 int main(void) {
   int ok = 0;
   {
@@ -575,6 +623,7 @@ int main(void) {
     test_int8_overflowing_sub(client_key);
     test_int8_overflowing_mul(client_key);
     test_try_decrypt_trivial(client_key);
+    test_oprf(client_key);
 
     client_key_destroy(client_key);
     public_key_destroy(public_key);

diff --git a/tfhe/src/c_api/high_level_api/integers.rs b/tfhe/src/c_api/high_level_api/integers.rs
@@ -1,16 +1,15 @@
-use crate::c_api::high_level_api::keys::CompactPublicKey;
-use crate::high_level_api::prelude::*;
-use std::ops::{
-    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
-    Mul, MulAssign, Neg, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub, SubAssign,
-};
-
 use crate::c_api::high_level_api::booleans::FheBool;
 use crate::c_api::high_level_api::i128::I128;
 use crate::c_api::high_level_api::i256::I256;
+use crate::c_api::high_level_api::keys::CompactPublicKey;
 use crate::c_api::high_level_api::u128::U128;
 use crate::c_api::high_level_api::u256::U256;
 use crate::c_api::utils::*;
+use crate::high_level_api::prelude::*;
+use std::ops::{
+    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
+    Mul, MulAssign, Neg, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub, SubAssign,
+};
 use std::os::raw::c_int;
 
 macro_rules! define_all_cast_into_for_integer_type {
@@ -659,3 +658,134 @@ pub unsafe extern "C" fn compact_fhe_int256_list_try_encrypt_with_compact_public
 }
 
 define_all_cast_into_for_integer_type!(FheBool);
+
+macro_rules! impl_oprf_for_uint {
+    (
+        name: $name:ident
+    ) => {
+
+        ::paste::paste! {
+            #[no_mangle]
+            pub unsafe extern "C" fn [<generate_oblivious_pseudo_random_ $name:snake>](
+                out_result: *mut *mut $name,
+                seed_low_bytes: u64,
+                seed_high_bytes: u64,
+            ) -> c_int {
+                use crate::high_level_api::IntegerId;
+                $crate::c_api::utils::catch_panic(|| {
+                    let seed_low_bytes: u128 = seed_low_bytes.into();
+                    let seed_high_bytes: u128 = seed_high_bytes.into();
+                    let seed = crate::Seed((seed_high_bytes << 64) | seed_low_bytes);
+
+                    let result = crate::FheUint::generate_oblivious_pseudo_random(
+                        seed,
+                        <crate::[<$name Id>] as IntegerId>::num_bits() as u64
+                    );
+                    *out_result = Box::into_raw(Box::new($name(result)));
+                })
+            }
+        }
+
+        ::paste::paste! {
+            #[no_mangle]
+            pub unsafe extern "C" fn [<generate_oblivious_pseudo_random_bits_ $name:snake>](
+                out_result: *mut *mut $name,
+                seed_low_bytes: u64,
+                seed_high_bytes: u64,
+                random_bits_count: u64,
+            ) -> c_int {
+
+                $crate::c_api::utils::catch_panic(|| {
+                    let seed_low_bytes: u128 = seed_low_bytes.into();
+                    let seed_high_bytes: u128 = seed_high_bytes.into();
+                    let seed = crate::Seed((seed_high_bytes << 64) | seed_low_bytes);
+
+                    let result = crate::FheUint::generate_oblivious_pseudo_random(seed, random_bits_count);
+                    *out_result = Box::into_raw(Box::new($name(result)));
+                })
+            }
+        }
+    };
+}
+
+macro_rules! impl_oprf_for_int {
+    (
+        name: $name:ident
+    ) => {
+
+        ::paste::paste! {
+            #[no_mangle]
+            pub unsafe extern "C" fn [<generate_oblivious_pseudo_random_unsigned_ $name:snake>](
+                out_result: *mut *mut $name,
+                seed_low_bytes: u64,
+                seed_high_bytes: u64,
+                random_bits_count: u64,
+            ) -> c_int {
+                $crate::c_api::utils::catch_panic(|| {
+                    let seed_low_bytes: u128 = seed_low_bytes.into();
+                    let seed_high_bytes: u128 = seed_high_bytes.into();
+                    let seed = crate::Seed((seed_high_bytes << 64) | seed_low_bytes);
+
+                    let result =
+                        crate::FheInt::generate_oblivious_pseudo_random(
+                            seed,
+                            crate::high_level_api::SignedRandomizationSpec::Unsigned {
+                                random_bits_count
+                            },
+                        );
+                    *out_result = Box::into_raw(Box::new($name(result)));
+                })
+            }
+        }
+
+        ::paste::paste! {
+            #[no_mangle]
+            pub unsafe extern "C" fn [<generate_oblivious_pseudo_random_full_signed_range_ $name:snake>](
+                out_result: *mut *mut $name,
+                seed_low_bytes: u64,
+                seed_high_bytes: u64,
+            ) -> c_int {
+                $crate::c_api::utils::catch_panic(|| {
+                    let seed_low_bytes: u128 = seed_low_bytes.into();
+                    let seed_high_bytes: u128 = seed_high_bytes.into();
+                    let seed = crate::Seed((seed_high_bytes << 64) | seed_low_bytes);
+
+                    let result = crate::FheInt::generate_oblivious_pseudo_random(
+                        seed,
+                        crate::high_level_api::SignedRandomizationSpec::FullSigned,
+                    );
+                    *out_result = Box::into_raw(Box::new($name(result)));
+
+                })
+            }
+        }
+    };
+}
+
+impl_oprf_for_uint!(name: FheUint2);
+impl_oprf_for_uint!(name: FheUint4);
+impl_oprf_for_uint!(name: FheUint6);
+impl_oprf_for_uint!(name: FheUint8);
+impl_oprf_for_uint!(name: FheUint10);
+impl_oprf_for_uint!(name: FheUint12);
+impl_oprf_for_uint!(name: FheUint14);
+impl_oprf_for_uint!(name: FheUint16);
+impl_oprf_for_uint!(name: FheUint32);
+impl_oprf_for_uint!(name: FheUint64);
+impl_oprf_for_uint!(name: FheUint128);
+impl_oprf_for_uint!(name: FheUint160);
+impl_oprf_for_uint!(name: FheUint256);
+
+impl_oprf_for_int!(name: FheInt2);
+impl_oprf_for_int!(name: FheInt4);
+impl_oprf_for_int!(name: FheInt6);
+impl_oprf_for_int!(name: FheInt8);
+impl_oprf_for_int!(name: FheInt10);
+impl_oprf_for_int!(name: FheInt12);
+impl_oprf_for_int!(name: FheInt14);
+impl_oprf_for_int!(name: FheInt16);
+impl_oprf_for_int!(name: FheInt32);
+impl_oprf_for_int!(name: FheInt64);
+impl_oprf_for_int!(name: FheInt128);
+impl_oprf_for_int!(name: FheInt160);
+impl_oprf_for_int!(name: FheInt256);
diff --git a/tfhe/src/high_level_api/mod.rs b/tfhe/src/high_level_api/mod.rs
@@ -24,21 +24,22 @@ macro_rules! expand_pub_use_fhe_type(
 );
 
 pub use crate::core_crypto::commons::math::random::Seed;
+pub use crate::integer::oprf::SignedRandomizationSpec;
 pub use config::{Config, ConfigBuilder};
 pub use errors::{Error, OutOfRangeError};
 pub use global_state::{set_server_key, unset_server_key, with_server_key_as_context};
+
+pub use integers::{
+    CompactFheInt, CompactFheIntList, CompactFheUint, CompactFheUintList, CompressedFheInt,
+    CompressedFheUint, FheInt, FheUint, IntegerId,
+};
 #[cfg(feature = "gpu")]
 pub use keys::CudaServerKey;
 pub use keys::{
     generate_keys, ClientKey, CompactPublicKey, CompressedCompactPublicKey, CompressedPublicKey,
     CompressedServerKey, KeySwitchingKey, PublicKey, ServerKey,
 };
 
-pub use integers::{
-    CompactFheInt, CompactFheIntList, CompactFheUint, CompactFheUintList, CompressedFheInt,
-    CompressedFheUint, FheInt, FheUint,
-};
-
 #[cfg(test)]
 mod tests;