refactor(integer): provide recompose_unsigned, recompose_signed functions

- to decrypt values and make it flexible wrt input type for the
recomposition, e.g. using u128 for noise squashed primitives

Author: IceTDrinker

tfhe/src/high_level_api/array/cpu/integers.rs

@@ -12,8 +12,9 @@ use crate::array::traits::{
 };
 use crate::high_level_api::global_state;
 use crate::high_level_api::integers::{FheIntId, FheUintId};
-use crate::integer::block_decomposition::{DecomposableInto, RecomposableFrom};
-use crate::integer::client_key::RecomposableSignedInteger;
+use crate::integer::block_decomposition::{
+    DecomposableInto, RecomposableFrom, RecomposableSignedInteger,
+};
 use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocable;
 use crate::integer::server_key::{Reciprocable, ScalarMultiplier};
 use crate::integer::{IntegerRadixCiphertext, RadixCiphertext, SignedRadixCiphertext};

tfhe/src/high_level_api/array/dynamic/signed.rs

@@ -11,8 +11,7 @@ use crate::high_level_api::array::{
 };
 use crate::high_level_api::global_state;
 use crate::high_level_api::integers::FheIntId;
-use crate::integer::block_decomposition::DecomposableInto;
-use crate::integer::client_key::RecomposableSignedInteger;
+use crate::integer::block_decomposition::{DecomposableInto, RecomposableSignedInteger};
 use crate::integer::SignedRadixCiphertext;
 use crate::prelude::{FheDecrypt, FheTryEncrypt};
 use crate::{ClientKey, Device, Error};

tfhe/src/high_level_api/array/gpu/integers.rs

@@ -15,8 +15,9 @@ use crate::core_crypto::gpu::CudaStreams;
 use crate::high_level_api::global_state;
 use crate::high_level_api::global_state::with_thread_local_cuda_streams;
 use crate::high_level_api::integers::{FheIntId, FheUintId};
-use crate::integer::block_decomposition::{DecomposableInto, RecomposableFrom};
-use crate::integer::client_key::RecomposableSignedInteger;
+use crate::integer::block_decomposition::{
+    DecomposableInto, RecomposableFrom, RecomposableSignedInteger,
+};
 use crate::integer::gpu::ciphertext::{
     CudaIntegerRadixCiphertext, CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext,
 };

tfhe/src/high_level_api/integers/signed/base.rs

@@ -7,7 +7,7 @@ use crate::high_level_api::global_state;
 use crate::high_level_api::integers::{FheUint, FheUintId, IntegerId};
 use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::traits::Tagged;
-use crate::integer::client_key::RecomposableSignedInteger;
+use crate::integer::block_decomposition::RecomposableSignedInteger;
 use crate::integer::parameters::RadixCiphertextConformanceParams;
 use crate::named::Named;
 use crate::prelude::CastFrom;

tfhe/src/high_level_api/integers/signed/encrypt.rs

@@ -4,8 +4,7 @@ use crate::high_level_api::global_state;
 use crate::high_level_api::global_state::with_thread_local_cuda_streams;
 use crate::high_level_api::integers::FheIntId;
 use crate::high_level_api::keys::InternalServerKey;
-use crate::integer::block_decomposition::DecomposableInto;
-use crate::integer::client_key::RecomposableSignedInteger;
+use crate::integer::block_decomposition::{DecomposableInto, RecomposableSignedInteger};
 #[cfg(feature = "gpu")]
 use crate::integer::gpu::ciphertext::CudaSignedRadixCiphertext;
 use crate::prelude::{FheDecrypt, FheTrivialEncrypt, FheTryEncrypt, FheTryTrivialEncrypt};

tfhe/src/integer/block_decomposition.rs

@@ -1,4 +1,4 @@
-use crate::core_crypto::prelude::{CastFrom, CastInto, Numeric};
+use crate::core_crypto::prelude::{CastFrom, CastInto, Numeric, SignedNumeric};
 use crate::integer::bigint::static_signed::StaticSignedBigInt;
 use crate::integer::bigint::static_unsigned::StaticUnsignedBigInt;
 use core::ops::{AddAssign, BitAnd, ShlAssign, ShrAssign};
@@ -88,6 +88,57 @@ impl<const N: usize> RecomposableFrom<u8> for StaticUnsignedBigInt<N> {}
 impl<const N: usize> DecomposableInto<u64> for StaticUnsignedBigInt<N> {}
 impl<const N: usize> DecomposableInto<u8> for StaticUnsignedBigInt<N> {}
 
+pub trait RecomposableSignedInteger:
+    RecomposableFrom<u64>
+    + std::ops::Neg<Output = Self>
+    + std::ops::Shr<u32, Output = Self>
+    + std::ops::BitOrAssign<Self>
+    + std::ops::BitOr<Self, Output = Self>
+    + std::ops::Mul<Self, Output = Self>
+    + CastFrom<Self>
+    + SignedNumeric
+{
+}
+
+impl RecomposableSignedInteger for i8 {}
+impl RecomposableSignedInteger for i16 {}
+impl RecomposableSignedInteger for i32 {}
+impl RecomposableSignedInteger for i64 {}
+impl RecomposableSignedInteger for i128 {}
+
+impl<const N: usize> RecomposableSignedInteger for StaticSignedBigInt<N> {}
+
+pub trait SignExtendable:
+    std::ops::Shl<u32, Output = Self> + std::ops::Shr<u32, Output = Self> + SignedNumeric
+{
+}
+
+impl<T> SignExtendable for T where T: RecomposableSignedInteger {}
+
+/// This function takes a signed integer of type `T` for which `num_bits_set`
+/// have been set.
+///
+/// It will set the most significant bits to the value of the bit
+/// at pos `num_bits_set - 1`.
+///
+/// This is used to correctly decrypt a signed radix ciphertext into a clear type
+/// that has more bits than the original ciphertext.
+///
+/// This is like doing i8 as i16, i16 as i64, i16 as i8, etc
+pub(in crate::integer) fn sign_extend_partial_number<T>(unpadded_value: T, num_bits_set: u32) -> T
+where
+    T: SignExtendable,
+{
+    if num_bits_set >= T::BITS as u32 {
+        return unpadded_value;
+    }
+
+    // Shift to put the last set bit in the position of the sign bit of T
+    // When right shifting this will do the sign extend automatically
+    let shift = T::BITS as u32 - num_bits_set;
+    (unpadded_value << shift) >> shift
+}
+
 #[derive(Copy, Clone)]
 #[repr(u32)]
 pub enum PaddingBitValue {
@@ -256,25 +307,6 @@ pub struct BlockRecomposer<T> {
     bit_pos: u32,
 }
 
-impl<T> BlockRecomposer<T>
-where
-    T: Recomposable,
-{
-    pub fn value(&self) -> T {
-        let is_signed = (T::ONE << (T::BITS as u32 - 1)) < T::ZERO;
-        if self.bit_pos >= (T::BITS as u32 - u32::from(is_signed)) {
-            self.data
-        } else {
-            let valid_mask = (T::ONE << self.bit_pos) - T::ONE;
-            self.data & valid_mask
-        }
-    }
-
-    pub fn unmasked_value(&self) -> T {
-        self.data
-    }
-}
-
 impl<T> BlockRecomposer<T>
 where
     T: Recomposable,
@@ -292,12 +324,21 @@ where
             bit_pos,
         }
     }
-}
 
-impl<T> BlockRecomposer<T>
-where
-    T: Recomposable,
-{
+    pub fn value(&self) -> T {
+        let is_signed = (T::ONE << (T::BITS as u32 - 1)) < T::ZERO;
+        if self.bit_pos >= (T::BITS as u32 - u32::from(is_signed)) {
+            self.data
+        } else {
+            let valid_mask = (T::ONE << self.bit_pos) - T::ONE;
+            self.data & valid_mask
+        }
+    }
+
+    pub fn unmasked_value(&self) -> T {
+        self.data
+    }
+
     pub fn add_unmasked<V>(&mut self, block: V) -> bool
     where
         T: CastFrom<V>,
@@ -328,6 +369,34 @@ where
 
         true
     }
+
+    pub(crate) fn recompose_unsigned<U>(input: impl Iterator<Item = U>, bits_in_block: u32) -> T
+    where
+        T: RecomposableFrom<U>,
+    {
+        let mut recomposer = Self::new(bits_in_block);
+        for limb in input {
+            if !recomposer.add_unmasked(limb) {
+                break;
+            }
+        }
+
+        recomposer.value()
+    }
+
+    pub(crate) fn recompose_signed<U>(input: impl Iterator<Item = U>, bits_in_block: u32) -> T
+    where
+        T: RecomposableFrom<U> + SignExtendable,
+    {
+        let mut recomposer = Self::new(bits_in_block);
+        for limb in input {
+            if !recomposer.add_unmasked(limb) {
+                break;
+            }
+        }
+
+        sign_extend_partial_number(recomposer.value(), recomposer.bit_pos)
+    }
 }
 
 #[cfg(test)]

tfhe/src/integer/ciphertext/base.rs

@@ -4,8 +4,9 @@ use crate::core_crypto::prelude::UnsignedNumeric;
 use crate::integer::backward_compatibility::ciphertext::{
     BaseCrtCiphertextVersions, BaseRadixCiphertextVersions, BaseSignedRadixCiphertextVersions,
 };
-use crate::integer::block_decomposition::{BlockRecomposer, RecomposableFrom};
-use crate::integer::client_key::{sign_extend_partial_number, RecomposableSignedInteger};
+use crate::integer::block_decomposition::{
+    BlockRecomposer, RecomposableFrom, RecomposableSignedInteger,
+};
 use crate::shortint::ciphertext::NotTrivialCiphertextError;
 use crate::shortint::parameters::CiphertextConformanceParams;
 use crate::shortint::Ciphertext;
@@ -102,15 +103,21 @@ impl RadixCiphertext {
     where
         Clear: UnsignedNumeric + RecomposableFrom<u64>,
     {
+        if !self.blocks.iter().all(|b| b.is_trivial()) {
+            return Err(NotTrivialCiphertextError);
+        }
+
         let bits_in_block = self.blocks[0].message_modulus.0.ilog2();
-        let mut recomposer = BlockRecomposer::<Clear>::new(bits_in_block);
 
-        for encrypted_block in &self.blocks {
-            let decrypted_block = encrypted_block.decrypt_trivial_message_and_carry()?;
-            recomposer.add_unmasked(decrypted_block);
-        }
+        let decrypted_block_iter = self
+            .blocks
+            .iter()
+            .map(|block| block.decrypt_trivial_message_and_carry().unwrap());
 
-        Ok(recomposer.value())
+        Ok(BlockRecomposer::recompose_unsigned(
+            decrypted_block_iter,
+            bits_in_block,
+        ))
     }
 }
 
@@ -204,17 +211,21 @@ impl SignedRadixCiphertext {
     where
         Clear: RecomposableSignedInteger,
     {
+        if !self.blocks.iter().all(|b| b.is_trivial()) {
+            return Err(NotTrivialCiphertextError);
+        }
+
         let bits_in_block = self.blocks[0].message_modulus.0.ilog2();
-        let mut recomposer = BlockRecomposer::<Clear>::new(bits_in_block);
 
-        for encrypted_block in &self.blocks {
-            let decrypted_block = encrypted_block.decrypt_trivial_message_and_carry()?;
-            recomposer.add_unmasked(decrypted_block);
-        }
+        let decrypted_block_iter = self
+            .blocks
+            .iter()
+            .map(|block| block.decrypt_trivial_message_and_carry().unwrap());
 
-        let num_bits_in_ctxt = bits_in_block * self.blocks.len() as u32;
-        let unpadded_value = recomposer.value();
-        Ok(sign_extend_partial_number(unpadded_value, num_bits_in_ctxt))
+        Ok(BlockRecomposer::recompose_signed(
+            decrypted_block_iter,
+            bits_in_block,
+        ))
     }
 }
 

tfhe/src/integer/client_key/mod.rs

@@ -9,14 +9,14 @@ pub(crate) mod secret_encryption_key;
 pub(crate) mod utils;
 
 use super::backward_compatibility::client_key::ClientKeyVersions;
-use super::block_decomposition::{DecomposableInto, RecomposableFrom};
 use super::ciphertext::{
     CompressedRadixCiphertext, CompressedSignedRadixCiphertext, RadixCiphertext,
     SignedRadixCiphertext,
 };
-use crate::core_crypto::prelude::{CastFrom, SignedNumeric, UnsignedNumeric};
-use crate::integer::bigint::static_signed::StaticSignedBigInt;
-use crate::integer::block_decomposition::BlockRecomposer;
+use crate::core_crypto::prelude::{SignedNumeric, UnsignedNumeric};
+use crate::integer::block_decomposition::{
+    BlockRecomposer, DecomposableInto, RecomposableFrom, RecomposableSignedInteger,
+};
 use crate::integer::ciphertext::boolean_value::BooleanBlock;
 use crate::integer::ciphertext::{CompressedCrtCiphertext, CrtCiphertext};
 use crate::integer::client_key::utils::i_crt;
@@ -33,53 +33,6 @@ use secret_encryption_key::SecretEncryptionKeyView;
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
 
-pub trait RecomposableSignedInteger:
-    RecomposableFrom<u64>
-    + std::ops::Neg<Output = Self>
-    + std::ops::Shr<u32, Output = Self>
-    + std::ops::BitOrAssign<Self>
-    + std::ops::BitOr<Self, Output = Self>
-    + std::ops::Mul<Self, Output = Self>
-    + CastFrom<Self>
-    + SignedNumeric
-{
-}
-
-impl RecomposableSignedInteger for i8 {}
-impl RecomposableSignedInteger for i16 {}
-impl RecomposableSignedInteger for i32 {}
-impl RecomposableSignedInteger for i64 {}
-impl RecomposableSignedInteger for i128 {}
-
-impl<const N: usize> RecomposableSignedInteger for StaticSignedBigInt<N> {}
-
-/// This function takes a signed integer of type `T` for which `num_bits_set`
-/// have been set.
-///
-/// It will set the most significant bits to the value of the bit
-/// at pos `num_bits_set - 1`.
-///
-/// This is used to correctly decrypt a signed radix ciphertext into a clear type
-/// that has more bits than the original ciphertext.
-///
-/// This is like doing i8 as i16, i16 as i64, i6 as i8, etc
-pub(in crate::integer) fn sign_extend_partial_number<T>(unpadded_value: T, num_bits_set: u32) -> T
-where
-    T: RecomposableSignedInteger,
-{
-    if num_bits_set >= T::BITS as u32 {
-        return unpadded_value;
-    }
-    let sign_bit_pos = num_bits_set - 1;
-    let sign_bit = (unpadded_value >> sign_bit_pos) & T::ONE;
-
-    // Creates a padding mask
-    // where bits above num_bits_set
-    // are 1s if sign bit is `1` else `0`
-    let padding = (T::MAX * sign_bit) << num_bits_set;
-    padding | unpadded_value
-}
-
 /// A structure containing the client key, which must be kept secret.
 ///
 /// This key can be used to encrypt both in Radix and CRT
@@ -381,18 +334,8 @@ impl ClientKey {
         }
 
         let bits_in_block = self.key.parameters.message_modulus().0.ilog2();
-        let mut recomposer = BlockRecomposer::<T>::new(bits_in_block);
-
-        for encrypted_block in blocks {
-            let decrypted_block = decrypt_block(&self.key, encrypted_block);
-            if !recomposer.add_unmasked(decrypted_block) {
-                // End of T::BITS reached no need to try more
-                // recomposition
-                break;
-            }
-        }
-
-        recomposer.value()
+        let decrypted_block_iter = blocks.iter().map(|block| decrypt_block(&self.key, block));
+        BlockRecomposer::recompose_unsigned(decrypted_block_iter, bits_in_block)
     }
 
     pub fn encrypt_signed_radix<T>(&self, message: T, num_blocks: usize) -> SignedRadixCiphertext
@@ -470,15 +413,16 @@ impl ClientKey {
         let message_modulus = self.parameters().message_modulus().0;
         assert!(message_modulus.is_power_of_two());
 
-        // Decrypting a signed value is the same as decrypting an unsigned value
-        // but, in the signed case,
-        // we have to take care of the case when the clear type T has more bits
-        // than what the ciphertext encrypts.
-        let unpadded_value = self.decrypt_radix_impl(&ctxt.blocks, decrypt_block);
+        if ctxt.blocks.is_empty() {
+            return T::ZERO;
+        }
 
-        let num_bits_in_message = message_modulus.ilog2();
-        let num_bits_in_ctxt = num_bits_in_message * ctxt.blocks.len() as u32;
-        sign_extend_partial_number(unpadded_value, num_bits_in_ctxt)
+        let bits_in_block = self.key.parameters.message_modulus().0.ilog2();
+        let decrypted_block_iter = ctxt
+            .blocks
+            .iter()
+            .map(|block| decrypt_block(&self.key, block));
+        BlockRecomposer::recompose_signed(decrypted_block_iter, bits_in_block)
     }
 
     /// Encrypts one block.