chore(hl): add a feature for extended types

zama-ai · Mar 5, 2025 · 098693d · 098693d
1 parent 5bcadd4
commit 098693d
Show file tree

Hide file tree

Showing 14 changed files with 2,881 additions and 1,203 deletions.
diff --git a/Makefile b/Makefile
@@ -350,6 +350,9 @@ clippy_integer: install_rs_check_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy \
 		--features=integer,experimental \
 		-p $(TFHE_SPEC) -- --no-deps -D warnings
+	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy \
+		--features=integer,experimental,extended-types \
+		-p $(TFHE_SPEC) -- --no-deps -D warnings
 
 .PHONY: clippy # Run clippy lints enabling the boolean, shortint, integer
 clippy: install_rs_check_toolchain
@@ -388,10 +391,10 @@ clippy_c_api: install_rs_check_toolchain
 .PHONY: clippy_js_wasm_api # Run clippy lints enabling the boolean, shortint, integer and the js wasm API
 clippy_js_wasm_api: install_rs_check_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy \
-		--features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,high-level-client-js-wasm-api,zk-pok \
+		--features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,high-level-client-js-wasm-api,zk-pok,extended-types \
 		-p $(TFHE_SPEC) -- --no-deps -D warnings
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy \
-		--features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,high-level-client-js-wasm-api \
+		--features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,high-level-client-js-wasm-api,extended-types \
 		-p $(TFHE_SPEC) -- --no-deps -D warnings
 
 .PHONY: clippy_tasks # Run clippy lints on helper tasks crate.
@@ -407,10 +410,10 @@ clippy_trivium: install_rs_check_toolchain
 .PHONY: clippy_all_targets # Run clippy lints on all targets (benches, examples, etc.)
 clippy_all_targets: install_rs_check_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy --all-targets \
-		--features=boolean,shortint,integer,internal-keycache,zk-pok,strings \
+		--features=boolean,shortint,integer,internal-keycache,zk-pok,strings,extended-types \
 		-p $(TFHE_SPEC) -- --no-deps -D warnings
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy --all-targets \
-		--features=boolean,shortint,integer,internal-keycache,zk-pok,strings,experimental \
+		--features=boolean,shortint,integer,internal-keycache,zk-pok,strings,extended-types,experimental \
 		-p $(TFHE_SPEC) -- --no-deps -D warnings
 
 .PHONY: clippy_tfhe_csprng # Run clippy lints on tfhe-csprng
@@ -528,15 +531,15 @@ build_web_js_api: install_rs_build_toolchain install_wasm_pack
 	cd tfhe && \
 	RUSTFLAGS="$(WASM_RUSTFLAGS)" rustup run "$(RS_BUILD_TOOLCHAIN)" \
 		wasm-pack build --release --target=web \
-		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,zk-pok
+		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,zk-pok,extended-types
 
 .PHONY: build_web_js_api_parallel # Build the js API targeting the web browser with parallelism support
 build_web_js_api_parallel: install_rs_check_toolchain install_wasm_pack
 	cd tfhe && \
 	rustup component add rust-src --toolchain $(RS_CHECK_TOOLCHAIN) && \
 	RUSTFLAGS="$(WASM_RUSTFLAGS) -C target-feature=+atomics,+bulk-memory" rustup run $(RS_CHECK_TOOLCHAIN) \
 		wasm-pack build --release --target=web \
-		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,parallel-wasm-api,zk-pok \
+		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,parallel-wasm-api,zk-pok,extended-types \
 		-Z build-std=panic_abort,std && \
 	find pkg/snippets -type f -iname workerHelpers.js -exec sed -i "s|const pkg = await import('..\/..\/..');|const pkg = await import('..\/..\/..\/tfhe.js');|" {} \;
 	jq '.files += ["snippets"]' tfhe/pkg/package.json > tmp_pkg.json && mv -f tmp_pkg.json tfhe/pkg/package.json
@@ -546,7 +549,7 @@ build_node_js_api: install_rs_build_toolchain install_wasm_pack
 	cd tfhe && \
 	RUSTFLAGS="$(WASM_RUSTFLAGS)" rustup run "$(RS_BUILD_TOOLCHAIN)" \
 		wasm-pack build --release --target=nodejs \
-		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,zk-pok
+		-- --features=boolean-client-js-wasm-api,shortint-client-js-wasm-api,integer-client-js-wasm-api,zk-pok,extended-types
 
 .PHONY: build_tfhe_csprng # Build tfhe_csprng
 build_tfhe_csprng: install_rs_build_toolchain

diff --git a/tfhe/Cargo.toml b/tfhe/Cargo.toml
@@ -92,6 +92,9 @@ internal-keycache = ["dep:fs2"]
 gpu = ["dep:tfhe-cuda-backend"]
 zk-pok = ["dep:tfhe-zk-pok"]
 
+# Adds more FheUint/FheInt types to the HL
+extended-types = []
+
 pbs-stats = []
 noise-asserts = []
 

diff --git a/tfhe/build.rs b/tfhe/build.rs
@@ -59,6 +59,8 @@ fn gen_c_api() {
         "gpu",
         #[cfg(feature = "zk-pok")]
         "zk-pok",
+        #[cfg(feature = "extended-types")]
+        "extended-types",
     ];
 
     let parse_expand_vec = if parse_expand_features_vec.is_empty() {

diff --git a/tfhe/src/c_api/high_level_api/mod.rs b/tfhe/src/c_api/high_level_api/mod.rs
@@ -61,6 +61,7 @@ pub enum FheTypes {
     Type_FheInt1024 = 32,
     Type_FheInt2048 = 33,
 
+    // Extended types
     Type_FheUint24 = 34,
     Type_FheUint40 = 35,
     Type_FheUint48 = 36,

diff --git a/tfhe/src/high_level_api/compact_list.rs b/tfhe/src/high_level_api/compact_list.rs
@@ -634,6 +634,47 @@ mod tests {
         }
     }
 
+    #[cfg(feature = "extended-types")]
+    #[test]
+    fn test_compact_list_extended_types() {
+        let config = crate::ConfigBuilder::default().build();
+
+        let ck = crate::ClientKey::generate(config);
+        let sk = crate::ServerKey::new(&ck);
+        let pk = crate::CompactPublicKey::new(&ck);
+
+        set_server_key(sk);
+
+        let compact_list = CompactCiphertextList::builder(&pk)
+            .push_with_num_bits(-17i64, 40)
+            .unwrap()
+            .push_with_num_bits(3u8, 24)
+            .unwrap()
+            .build_packed();
+
+        let serialized = bincode::serialize(&compact_list).unwrap();
+        let compact_list: CompactCiphertextList = bincode::deserialize(&serialized).unwrap();
+        let expander = compact_list.expand().unwrap();
+
+        {
+            let a: crate::FheInt40 = expander.get(0).unwrap().unwrap();
+            let b: crate::FheUint24 = expander.get(1).unwrap().unwrap();
+
+            let a: i64 = a.decrypt(&ck);
+            assert_eq!(a, -17);
+            let b: u8 = b.decrypt(&ck);
+            assert_eq!(b, 3);
+        }
+
+        {
+            // Incorrect type
+            assert!(expander.get::<FheUint32>(0).is_err());
+
+            // Correct type but wrong number of bits
+            assert!(expander.get::<FheInt64>(0).is_err());
+        }
+    }
+
     #[test]
     fn test_compact_list_with_casting() {
         let config = crate::ConfigBuilder::with_custom_parameters(

diff --git a/tfhe/src/high_level_api/integers/mod.rs b/tfhe/src/high_level_api/integers/mod.rs
@@ -1,21 +1,32 @@
 expand_pub_use_fhe_type!(
     pub use unsigned{
         FheUint2, FheUint4, FheUint6, FheUint8, FheUint10, FheUint12, FheUint14, FheUint16,
-        FheUint24, FheUint32, FheUint40, FheUint48, FheUint56, FheUint64, FheUint72, FheUint80,
-        FheUint88, FheUint96, FheUint104, FheUint112, FheUint120, FheUint128, FheUint136,
-        FheUint144, FheUint152, FheUint160, FheUint168, FheUint176, FheUint184, FheUint192,
-        FheUint200, FheUint208, FheUint216, FheUint224, FheUint232, FheUint240, FheUint248,
-        FheUint256, FheUint512, FheUint1024, FheUint2048,
+        FheUint32, FheUint64, FheUint128, FheUint160, FheUint256, FheUint512, FheUint1024,
+        FheUint2048,
+    };
+);
+#[cfg(feature = "extended-types")]
+expand_pub_use_fhe_type!(
+    pub use unsigned{
+        FheUint24, FheUint40, FheUint48, FheUint56, FheUint72, FheUint80,FheUint88, FheUint96,
+        FheUint104, FheUint112, FheUint120, FheUint136, FheUint144, FheUint152, FheUint168,
+        FheUint176, FheUint184, FheUint192, FheUint200, FheUint208, FheUint216, FheUint224,
+        FheUint232, FheUint240, FheUint248,
     };
 );
 
 expand_pub_use_fhe_type!(
     pub use signed{
-        FheInt2, FheInt4, FheInt6, FheInt8, FheInt10, FheInt12, FheInt14, FheInt16, FheInt24,
-        FheInt32, FheInt40, FheInt48, FheInt56, FheInt64, FheInt72, FheInt80, FheInt88, FheInt96,
-        FheInt104, FheInt112, FheInt120, FheInt128, FheInt136, FheInt144, FheInt152, FheInt160,
-        FheInt168, FheInt176, FheInt184, FheInt192, FheInt200, FheInt208, FheInt216, FheInt224,
-        FheInt232, FheInt240, FheInt248, FheInt256, FheInt512, FheInt1024, FheInt2048,
+        FheInt2, FheInt4, FheInt6, FheInt8, FheInt10, FheInt12, FheInt14, FheInt16, FheInt32,
+        FheInt64, FheInt128, FheInt160, FheInt256, FheInt512, FheInt1024, FheInt2048,
+    };
+);
+#[cfg(feature = "extended-types")]
+expand_pub_use_fhe_type!(
+    pub use signed{
+        FheInt24, FheInt40, FheInt48, FheInt56, FheInt72, FheInt80, FheInt88, FheInt96, FheInt104,
+        FheInt112, FheInt120, FheInt136, FheInt144, FheInt152, FheInt168, FheInt176, FheInt184,
+        FheInt192, FheInt200, FheInt208, FheInt216, FheInt224, FheInt232, FheInt240, FheInt248,
     };
 );
 

diff --git a/tfhe/src/high_level_api/integers/signed/mod.rs b/tfhe/src/high_level_api/integers/signed/mod.rs
@@ -17,10 +17,15 @@ pub(in crate::high_level_api) use inner::{RadixCiphertext, RadixCiphertextVersio
 
 expand_pub_use_fhe_type!(
     pub use static_{
-        FheInt2, FheInt4, FheInt6, FheInt8, FheInt10, FheInt12, FheInt14, FheInt16, FheInt24,
-        FheInt32, FheInt40, FheInt48, FheInt56, FheInt64, FheInt72, FheInt80, FheInt88, FheInt96,
-        FheInt104, FheInt112, FheInt120, FheInt128, FheInt136, FheInt144, FheInt152, FheInt160,
-        FheInt168, FheInt176, FheInt184, FheInt192, FheInt200, FheInt208, FheInt216, FheInt224,
-        FheInt232, FheInt240, FheInt248, FheInt256, FheInt512, FheInt1024, FheInt2048
+        FheInt2, FheInt4, FheInt6, FheInt8, FheInt10, FheInt12, FheInt14, FheInt16, FheInt32,
+        FheInt64,FheInt128, FheInt160, FheInt256, FheInt512, FheInt1024, FheInt2048
+    };
+);
+#[cfg(feature = "extended-types")]
+expand_pub_use_fhe_type!(
+    pub use static_{
+        FheInt24, FheInt40, FheInt48, FheInt56, FheInt72, FheInt80, FheInt88, FheInt96, FheInt104,
+        FheInt112, FheInt120, FheInt136, FheInt144, FheInt152, FheInt168, FheInt176, FheInt184,
+        FheInt192, FheInt200, FheInt208, FheInt216, FheInt224, FheInt232, FheInt240, FheInt248
     };
 );