tfhe/high_level_api/booleans/
squashed_noise.rs

1use super::base::FheBool;
2use crate::backward_compatibility::booleans::{
3    InnerSquashedNoiseBooleanVersionedOwned, SquashedNoiseFheBoolVersions,
4};
5use crate::high_level_api::details::MaybeCloned;
6use crate::high_level_api::errors::UninitializedNoiseSquashing;
7use crate::high_level_api::global_state;
8use crate::high_level_api::global_state::with_internal_keys;
9use crate::high_level_api::keys::InternalServerKey;
10use crate::high_level_api::traits::{FheDecrypt, SquashNoise};
11use crate::integer::ciphertext::SquashedNoiseBooleanBlock;
12use crate::named::Named;
13use crate::{ClientKey, Device, Tag};
14use serde::{Deserializer, Serializer};
15use tfhe_versionable::{Unversionize, UnversionizeError, Versionize, VersionizeOwned};
16
17/// Enum that manages the current inner representation of a boolean.
18pub(in crate::high_level_api) enum InnerSquashedNoiseBoolean {
19    Cpu(SquashedNoiseBooleanBlock),
20}
21
22impl Clone for InnerSquashedNoiseBoolean {
23    fn clone(&self) -> Self {
24        match self {
25            Self::Cpu(inner) => Self::Cpu(inner.clone()),
26        }
27    }
28}
29impl serde::Serialize for InnerSquashedNoiseBoolean {
30    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
31    where
32        S: Serializer,
33    {
34        self.on_cpu().serialize(serializer)
35    }
36}
37
38impl<'de> serde::Deserialize<'de> for InnerSquashedNoiseBoolean {
39    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
40    where
41        D: Deserializer<'de>,
42    {
43        let mut deserialized = Self::Cpu(
44            crate::integer::ciphertext::SquashedNoiseBooleanBlock::deserialize(deserializer)?,
45        );
46        deserialized.move_to_device_of_server_key_if_set();
47        Ok(deserialized)
48    }
49}
50
51// Only CPU data are serialized so we only versionize the CPU type.
52#[derive(serde::Serialize, serde::Deserialize)]
53#[cfg_attr(dylint_lib = "tfhe_lints", allow(serialize_without_versionize))]
54pub(crate) struct InnerSquashedNoiseBooleanVersionOwned(
55    <crate::integer::ciphertext::SquashedNoiseBooleanBlock as VersionizeOwned>::VersionedOwned,
56);
57
58impl Versionize for InnerSquashedNoiseBoolean {
59    type Versioned<'vers> = InnerSquashedNoiseBooleanVersionedOwned;
60
61    fn versionize(&self) -> Self::Versioned<'_> {
62        let data = self.on_cpu();
63        let versioned = data.into_owned().versionize_owned();
64        InnerSquashedNoiseBooleanVersionedOwned::V0(InnerSquashedNoiseBooleanVersionOwned(
65            versioned,
66        ))
67    }
68}
69impl VersionizeOwned for InnerSquashedNoiseBoolean {
70    type VersionedOwned = InnerSquashedNoiseBooleanVersionedOwned;
71
72    fn versionize_owned(self) -> Self::VersionedOwned {
73        let cpu_data = self.on_cpu();
74        InnerSquashedNoiseBooleanVersionedOwned::V0(InnerSquashedNoiseBooleanVersionOwned(
75            cpu_data.into_owned().versionize_owned(),
76        ))
77    }
78}
79
80impl Unversionize for InnerSquashedNoiseBoolean {
81    fn unversionize(versioned: Self::VersionedOwned) -> Result<Self, UnversionizeError> {
82        match versioned {
83            InnerSquashedNoiseBooleanVersionedOwned::V0(v0) => {
84                let mut unversioned = Self::Cpu(
85                    crate::integer::ciphertext::SquashedNoiseBooleanBlock::unversionize(v0.0)?,
86                );
87                unversioned.move_to_device_of_server_key_if_set();
88                Ok(unversioned)
89            }
90        }
91    }
92}
93
94impl InnerSquashedNoiseBoolean {
95    /// Returns the inner cpu ciphertext if self is on the CPU, otherwise, returns a copy
96    /// that is on the CPU
97    pub(crate) fn on_cpu(&self) -> MaybeCloned<'_, SquashedNoiseBooleanBlock> {
98        match self {
99            Self::Cpu(ct) => MaybeCloned::Borrowed(ct),
100        }
101    }
102
103    #[allow(clippy::needless_pass_by_ref_mut)]
104    pub(crate) fn move_to_device(&mut self, device: Device) {
105        match (&self, device) {
106            (Self::Cpu(_), Device::Cpu) => {
107                // Nothing to do, we already are on the correct device
108            }
109            #[cfg(any(feature = "gpu", feature = "hpu"))]
110            _ => panic!("Cuda/Hpu devices do not support noise squashing yet"),
111        }
112    }
113
114    #[inline]
115    pub(crate) fn move_to_device_of_server_key_if_set(&mut self) {
116        if let Some(device) = global_state::device_of_internal_keys() {
117            self.move_to_device(device);
118        }
119    }
120}
121
122#[derive(Clone, serde::Deserialize, serde::Serialize, Versionize)]
123#[versionize(SquashedNoiseFheBoolVersions)]
124pub struct SquashedNoiseFheBool {
125    inner: InnerSquashedNoiseBoolean,
126    tag: Tag,
127}
128
129impl Named for SquashedNoiseFheBool {
130    const NAME: &'static str = "high_level_api::SquashedNoiseFheBool";
131}
132
133impl SquashedNoiseFheBool {
134    pub fn underlying_squashed_noise_ciphertext(
135        &self,
136    ) -> MaybeCloned<'_, SquashedNoiseBooleanBlock> {
137        self.inner.on_cpu()
138    }
139}
140
141impl FheDecrypt<bool> for SquashedNoiseFheBool {
142    fn decrypt(&self, key: &ClientKey) -> bool {
143        key.key
144            .noise_squashing_private_key
145            .as_ref()
146            .map(|noise_squashing_private_key| {
147                noise_squashing_private_key.decrypt_bool(&self.inner.on_cpu())
148            })
149            .expect(
150                "No noise squashing private key in your ClientKey, cannot decrypt. \
151                Did you call `enable_noise_squashing` when creating your Config?",
152            )
153            .unwrap()
154    }
155}
156
157impl SquashNoise for FheBool {
158    type Output = SquashedNoiseFheBool;
159
160    fn squash_noise(&self) -> crate::Result<Self::Output> {
161        with_internal_keys(|keys| match keys {
162            InternalServerKey::Cpu(server_key) => {
163                let noise_squashing_key = server_key
164                    .key
165                    .noise_squashing_key
166                    .as_ref()
167                    .ok_or(UninitializedNoiseSquashing)?;
168
169                Ok(SquashedNoiseFheBool {
170                    inner: InnerSquashedNoiseBoolean::Cpu(
171                        noise_squashing_key.squash_boolean_block_noise(
172                            server_key.key.pbs_key(),
173                            &self.ciphertext.on_cpu(),
174                        )?,
175                    ),
176                    tag: server_key.tag.clone(),
177                })
178            }
179            #[cfg(feature = "gpu")]
180            InternalServerKey::Cuda(_) => Err(crate::error!(
181                "Cuda devices do not support noise squashing yet"
182            )),
183            #[cfg(feature = "hpu")]
184            InternalServerKey::Hpu(_device) => {
185                Err(crate::error!("Hpu devices do not support noise squashing"))
186            }
187        })
188    }
189}