refactor(hl)!: use a trait for common ciphertext lists methods

BREAKING CHANGE: - The `CiphertextList` trait needs to be in scope to use the common methods of the `CompressedCiphertextList` and `CompactCiphertextListExpander` - The `.get` of the `CompactCiphertextListExpander` now returns a `Result<Option>` instead of an `Option<Result>`
zama-ai · Sep 11, 2024 · 199fb58 · 199fb58
1 parent 8ac592a
commit 199fb58
Show file tree

Hide file tree

Showing 12 changed files with 96 additions and 118 deletions.
diff --git a/tfhe/docs/guides/zk-pok.md b/tfhe/docs/guides/zk-pok.md
@@ -12,7 +12,7 @@ Using this feature is straightforward: during encryption, the client generates t
 
 ```rust
 use rand::prelude::*;
-use tfhe::prelude::FheDecrypt;
+use tfhe::prelude::*;
 use tfhe::set_server_key;
 use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
 
@@ -44,8 +44,8 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         // Verify the ciphertexts
         let mut expander = proven_compact_list.verify_and_expand(public_zk_params, &public_key)?;
-        let a: tfhe::FheUint64 = expander.get(0).unwrap()?;
-        let b: tfhe::FheUint64 = expander.get(1).unwrap()?;
+        let a: tfhe::FheUint64 = expander.get(0)?.unwrap();
+        let b: tfhe::FheUint64 = expander.get(1)?.unwrap();
 
         a + b
     };
@@ -78,7 +78,7 @@ This works essentially in the same way as before. Additionally, you need to indi
 
 ```rust
 use rand::prelude::*;
-use tfhe::prelude::FheDecrypt;
+use tfhe::prelude::*;
 use tfhe::set_server_key;
 use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
 
@@ -116,8 +116,8 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         // Verify the ciphertexts
         let mut expander = proven_compact_list.verify_and_expand(public_zk_params, &public_key)?;
-        let a: tfhe::FheUint64 = expander.get(0).unwrap()?;
-        let b: tfhe::FheUint64 = expander.get(1).unwrap()?;
+        let a: tfhe::FheUint64 = expander.get(0)?.unwrap();
+        let b: tfhe::FheUint64 = expander.get(1)?.unwrap();
 
         a + b
     };

diff --git a/tfhe/src/c_api/high_level_api/compact_list.rs b/tfhe/src/c_api/high_level_api/compact_list.rs
@@ -15,6 +15,7 @@ use crate::c_api::high_level_api::utils::{
 #[cfg(feature = "zk-pok")]
 use crate::c_api::high_level_api::zk::{CompactPkePublicParams, ZkComputeLoad};
 use crate::c_api::utils::{catch_panic, get_mut_checked, get_ref_checked};
+use crate::prelude::CiphertextList;
 use std::ffi::c_int;
 
 pub struct CompactCiphertextListBuilder(crate::high_level_api::CompactCiphertextListBuilder);

diff --git a/tfhe/src/c_api/high_level_api/compressed_ciphertext_list.rs b/tfhe/src/c_api/high_level_api/compressed_ciphertext_list.rs
@@ -8,6 +8,7 @@ use crate::c_api::high_level_api::utils::{
     impl_destroy_on_type, impl_serialize_deserialize_on_type,
 };
 use crate::c_api::utils::{catch_panic, get_mut_checked, get_ref_checked};
+use crate::prelude::CiphertextList;
 use std::ffi::c_int;
 
 pub struct CompressedCiphertextListBuilder(crate::high_level_api::CompressedCiphertextListBuilder);

diff --git a/tfhe/src/high_level_api/backward_compatibility/booleans.rs b/tfhe/src/high_level_api/backward_compatibility/booleans.rs
@@ -7,6 +7,7 @@ use crate::high_level_api::booleans::{
     InnerBoolean, InnerBooleanVersionOwned, InnerCompressedFheBool,
 };
 use crate::integer::ciphertext::{CompactCiphertextList, DataKind};
+use crate::prelude::CiphertextList;
 use crate::{
     CompactCiphertextList as HlCompactCiphertextList, CompressedFheBool, Error, FheBool, Tag,
 };
@@ -111,7 +112,7 @@ impl CompactFheBool {
         let block = list
             .inner
             .get::<crate::integer::BooleanBlock>(0)
-            .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+            .map(|b| b.ok_or_else(|| Error::new("Failed to expand compact list".to_string())))??;
 
         let mut ciphertext = FheBool::new(block, Tag::default());
         ciphertext.ciphertext.move_to_device_of_server_key_if_set();
@@ -148,7 +149,9 @@ impl CompactFheBoolList {
                 let block = list
                     .inner
                     .get::<crate::integer::BooleanBlock>(idx)
-                    .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+                    .map(|list| {
+                        list.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))
+                    })??;
 
                 let mut ciphertext = FheBool::new(block, Tag::default());
                 ciphertext.ciphertext.move_to_device_of_server_key_if_set();

diff --git a/tfhe/src/high_level_api/backward_compatibility/integers.rs b/tfhe/src/high_level_api/backward_compatibility/integers.rs
@@ -16,6 +16,7 @@ use crate::integer::ciphertext::{
     CompressedRadixCiphertext as IntegerCompressedRadixCiphertext,
     CompressedSignedRadixCiphertext as IntegerCompressedSignedRadixCiphertext, DataKind,
 };
+use crate::prelude::CiphertextList;
 use crate::shortint::ciphertext::CompressedModulusSwitchedCiphertext;
 use crate::shortint::{Ciphertext, ServerKey};
 use crate::{CompactCiphertextList as HlCompactCiphertextList, Error, Tag};
@@ -277,7 +278,9 @@ where
         let ct = list
             .inner
             .get::<crate::integer::SignedRadixCiphertext>(0)
-            .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+            .map(|list| {
+                list.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))
+            })??;
         Ok(FheInt::new(ct, Tag::default()))
     }
 }
@@ -316,7 +319,9 @@ where
                 let ct = list
                     .inner
                     .get::<crate::integer::SignedRadixCiphertext>(idx)
-                    .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+                    .map(|list| {
+                        list.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))
+                    })??;
                 Ok(FheInt::new(ct, Tag::default()))
             })
             .collect::<Result<Vec<_>, _>>()
@@ -353,7 +358,9 @@ where
         let ct = list
             .inner
             .get::<crate::integer::RadixCiphertext>(0)
-            .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+            .map(|ct| {
+                ct.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))
+            })??;
         Ok(FheUint::new(ct, Tag::default()))
     }
 }
@@ -391,7 +398,9 @@ where
                 let ct = list
                     .inner
                     .get::<crate::integer::RadixCiphertext>(idx)
-                    .ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
+                    .map(|ct| {
+                        ct.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))
+                    })??;
                 Ok(FheUint::new(ct, Tag::default()))
             })
             .collect::<Result<Vec<_>, _>>()

diff --git a/tfhe/src/high_level_api/compact_list.rs b/tfhe/src/high_level_api/compact_list.rs
@@ -16,6 +16,7 @@ use crate::integer::parameters::{
     IntegerCompactCiphertextListUnpackingMode,
 };
 use crate::named::Named;
+use crate::prelude::CiphertextList;
 use crate::shortint::MessageModulus;
 #[cfg(feature = "zk-pok")]
 pub use zk::ProvenCompactCiphertextList;
@@ -295,27 +296,27 @@ pub struct CompactCiphertextListExpander {
     tag: Tag,
 }
 
-impl CompactCiphertextListExpander {
-    pub fn len(&self) -> usize {
+impl CiphertextList for CompactCiphertextListExpander {
+    fn len(&self) -> usize {
         self.inner.len()
     }
 
-    pub fn is_empty(&self) -> bool {
+    fn is_empty(&self) -> bool {
         self.len() == 0
     }
 
-    pub fn get_kind_of(&self, index: usize) -> Option<crate::FheTypes> {
+    fn get_kind_of(&self, index: usize) -> Option<crate::FheTypes> {
         self.inner.get_kind_of(index).and_then(|data_kind| {
             crate::FheTypes::from_data_kind(data_kind, self.inner.message_modulus())
         })
     }
 
-    pub fn get<T>(&self, index: usize) -> Option<crate::Result<T>>
+    fn get<T>(&self, index: usize) -> crate::Result<Option<T>>
     where
         T: Expandable + Tagged,
     {
         let mut expanded = self.inner.get::<T>(index);
-        if let Some(Ok(inner)) = &mut expanded {
+        if let Ok(Some(inner)) = &mut expanded {
             inner.tag_mut().set_data(self.tag.data());
         }
         expanded
@@ -468,15 +469,15 @@ mod tests {
             let e: u8 = e.decrypt(&ck);
             assert_eq!(e, 3);
 
-            assert!(expander.get::<FheBool>(5).is_none());
+            assert!(expander.get::<FheBool>(5).unwrap().is_none());
         }
 
         {
             // Incorrect type
-            assert!(expander.get::<FheInt64>(0).unwrap().is_err());
+            assert!(expander.get::<FheInt64>(0).is_err());
 
             // Correct type but wrong number of bits
-            assert!(expander.get::<FheUint16>(0).unwrap().is_err());
+            assert!(expander.get::<FheUint16>(0).is_err());
         }
     }
 
@@ -529,15 +530,15 @@ mod tests {
             let d: u8 = d.decrypt(&ck);
             assert_eq!(d, 3);
 
-            assert!(expander.get::<FheBool>(4).is_none());
+            assert!(expander.get::<FheBool>(4).unwrap().is_none());
         }
 
         {
             // Incorrect type
-            assert!(expander.get::<FheInt64>(0).unwrap().is_err());
+            assert!(expander.get::<FheInt64>(0).is_err());
 
             // Correct type but wrong number of bits
-            assert!(expander.get::<FheUint16>(0).unwrap().is_err());
+            assert!(expander.get::<FheUint16>(0).is_err());
         }
     }
 
@@ -596,15 +597,15 @@ mod tests {
             let d: u8 = d.decrypt(&ck);
             assert_eq!(d, 3);
 
-            assert!(expander.get::<FheBool>(4).is_none());
+            assert!(expander.get::<FheBool>(4).unwrap().is_none());
         }
 
         {
             // Incorrect type
-            assert!(expander.get::<FheInt64>(0).unwrap().is_err());
+            assert!(expander.get::<FheInt64>(0).is_err());
 
             // Correct type but wrong number of bits
-            assert!(expander.get::<FheUint16>(0).unwrap().is_err());
+            assert!(expander.get::<FheUint16>(0).is_err());
         }
     }
 }
diff --git a/tfhe/src/high_level_api/compressed_ciphertext_list.rs b/tfhe/src/high_level_api/compressed_ciphertext_list.rs
@@ -6,7 +6,7 @@ use crate::core_crypto::commons::math::random::{Deserialize, Serialize};
 use crate::high_level_api::integers::{FheIntId, FheUintId};
 use crate::integer::ciphertext::{Compressible, DataKind, Expandable};
 use crate::named::Named;
-use crate::prelude::Tagged;
+use crate::prelude::{CiphertextList, Tagged};
 use crate::shortint::Ciphertext;
 use crate::{FheBool, FheInt, FheUint, Tag};
 
@@ -97,16 +97,16 @@ impl Tagged for CompressedCiphertextList {
     }
 }
 
-impl CompressedCiphertextList {
-    pub fn len(&self) -> usize {
+impl CiphertextList for CompressedCiphertextList {
+    fn len(&self) -> usize {
         self.inner.len()
     }
 
-    pub fn is_empty(&self) -> bool {
+    fn is_empty(&self) -> bool {
         self.len() == 0
     }
 
-    pub fn get_kind_of(&self, index: usize) -> Option<crate::FheTypes> {
+    fn get_kind_of(&self, index: usize) -> Option<crate::FheTypes> {
         Some(match self.inner.get_kind_of(index)? {
             DataKind::Unsigned(n) => {
                 let num_bits_per_block = self.inner.packed_list.message_modulus.0.ilog2() as usize;
@@ -152,7 +152,7 @@ impl CompressedCiphertextList {
         })
     }
 
-    pub fn get<T>(&self, index: usize) -> crate::Result<Option<T>>
+    fn get<T>(&self, index: usize) -> crate::Result<Option<T>>
     where
         T: Expandable + Tagged,
     {

diff --git a/tfhe/src/high_level_api/prelude.rs b/tfhe/src/high_level_api/prelude.rs
@@ -6,10 +6,10 @@
 //! use tfhe::prelude::*;
 //! ```
 pub use crate::high_level_api::traits::{
-    BitSlice, DivRem, FheBootstrap, FheDecrypt, FheEncrypt, FheEq, FheKeyswitch, FheMax, FheMin,
-    FheNumberConstant, FheOrd, FheTrivialEncrypt, FheTryEncrypt, FheTryTrivialEncrypt, IfThenElse,
-    OverflowingAdd, OverflowingMul, OverflowingSub, RotateLeft, RotateLeftAssign, RotateRight,
-    RotateRightAssign, Tagged,
+    BitSlice, CiphertextList, DivRem, FheBootstrap, FheDecrypt, FheEncrypt, FheEq, FheKeyswitch,
+    FheMax, FheMin, FheNumberConstant, FheOrd, FheTrivialEncrypt, FheTryEncrypt,
+    FheTryTrivialEncrypt, IfThenElse, OverflowingAdd, OverflowingMul, OverflowingSub, RotateLeft,
+    RotateLeftAssign, RotateRight, RotateRightAssign, Tagged,
 };
 
 pub use crate::conformance::ParameterSetConformant;

diff --git a/tfhe/src/high_level_api/traits.rs b/tfhe/src/high_level_api/traits.rs
@@ -2,6 +2,7 @@ use std::ops::RangeBounds;
 
 use crate::error::InvalidRangeError;
 use crate::high_level_api::ClientKey;
+use crate::integer::ciphertext::Expandable;
 use crate::{FheBool, Tag};
 
 /// Trait used to have a generic way of creating a value of a FHE type
@@ -199,3 +200,12 @@ pub trait Tagged {
 
     fn tag_mut(&mut self) -> &mut Tag;
 }
+
+pub trait CiphertextList {
+    fn len(&self) -> usize;
+    fn is_empty(&self) -> bool;
+    fn get_kind_of(&self, index: usize) -> Option<crate::FheTypes>;
+    fn get<T>(&self, index: usize) -> crate::Result<Option<T>>
+    where
+        T: Expandable + Tagged;
+}
diff --git a/tfhe/src/integer/ciphertext/compact_list.rs b/tfhe/src/integer/ciphertext/compact_list.rs
@@ -283,12 +283,13 @@ impl CompactCiphertextListExpander {
             .map(|block| (block, current_info))
     }
 
-    pub fn get<T>(&self, index: usize) -> Option<crate::Result<T>>
+    pub fn get<T>(&self, index: usize) -> crate::Result<Option<T>>
     where
         T: Expandable,
     {
         self.blocks_of(index)
             .map(|(blocks, kind)| T::from_expanded_blocks(blocks.to_owned(), kind))
+            .transpose()
     }
 
     pub(crate) fn message_modulus(&self) -> MessageModulus {

diff --git a/tfhe/src/js_on_wasm_api/js_high_level_api/integers.rs b/tfhe/src/js_on_wasm_api/js_high_level_api/integers.rs
@@ -1036,12 +1036,13 @@ macro_rules! define_expander_get_method {
                     #[wasm_bindgen]
                     pub fn [<get_uint $num_bits>] (&mut self, index: usize) -> Result<[<FheUint $num_bits>], JsError> {
                         catch_panic_result(|| {
-                           self.0.get::<crate::[<FheUint $num_bits>]>(index)
-                                .map_or_else(
-                                    || Err(JsError::new(&format!("Index {index} is out of bounds"))),
-                                    |a| a.map_err(into_js_error),
-                                )
-                                .map([<FheUint $num_bits>])
+                            self.0.get::<crate::[<FheUint $num_bits>]>(index)
+                                .map_err(into_js_error)
+                                .map(|val|
+                                      val.map_or_else(
+                                          || Err(JsError::new(&format!("Index {index} is out of bounds"))),
+                                          |val| Ok([<FheUint $num_bits>](val))
+                                    ))?
                         })
                     }
                 )*
@@ -1062,11 +1063,12 @@ macro_rules! define_expander_get_method {
                     pub fn [<get_int $num_bits>] (&mut self, index: usize) -> Result<[<FheInt $num_bits>], JsError> {
                         catch_panic_result(|| {
                            self.0.get::<crate::[<FheInt $num_bits>]>(index)
-                                .map_or_else(
-                                    || Err(JsError::new(&format!("Index {index} is out of bounds"))),
-                                    |a| a.map_err(into_js_error),
-                                )
-                                .map([<FheInt $num_bits>])
+                                .map_err(into_js_error)
+                                .map(|val|
+                                      val.map_or_else(
+                                          || Err(JsError::new(&format!("Index {index} is out of bounds"))),
+                                          |val| Ok([<FheInt $num_bits>](val))
+                                    ))?
                         })
                     }
                 )*
@@ -1088,11 +1090,13 @@ impl CompactCiphertextListExpander {
         catch_panic_result(|| {
             self.0
                 .get::<crate::FheBool>(index)
-                .map_or_else(
-                    || Err(JsError::new(&format!("Index {index} is out of bounds"))),
-                    |a| a.map_err(into_js_error),
-                )
-                .map(FheBool)
+                .map_err(into_js_error)
+                .map(|val| {
+                    val.map_or_else(
+                        || Err(JsError::new(&format!("Index {index} is out of bounds"))),
+                        |val| Ok(FheBool(val)),
+                    )
+                })?
         })
     }