tfhe/shortint/noise_squashing/
compressed_server_key.rs

1use super::{NoiseSquashingKey, NoiseSquashingKeyConformanceParams, NoiseSquashingPrivateKey};
2use crate::conformance::ParameterSetConformant;
3use crate::core_crypto::algorithms::lwe_bootstrap_key_conversion::par_convert_standard_lwe_bootstrap_key_to_fourier_128;
4use crate::core_crypto::algorithms::lwe_bootstrap_key_generation::par_allocate_and_generate_new_seeded_lwe_bootstrap_key;
5use crate::core_crypto::commons::math::random::Seeder;
6use crate::core_crypto::entities::{Fourier128LweBootstrapKeyOwned, SeededLweBootstrapKeyOwned};
7use crate::shortint::backward_compatibility::noise_squashing::CompressedNoiseSquashingKeyVersions;
8use crate::shortint::client_key::ClientKey;
9use crate::shortint::engine::ShortintEngine;
10use crate::shortint::parameters::{CarryModulus, CoreCiphertextModulus, MessageModulus};
11use crate::shortint::server_key::{
12    CompressedModulusSwitchNoiseReductionKey, ModulusSwitchNoiseReductionKeyConformanceParams,
13};
14use serde::{Deserialize, Serialize};
15use tfhe_versionable::Versionize;
16
17#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Versionize)]
18#[versionize(CompressedNoiseSquashingKeyVersions)]
19pub struct CompressedNoiseSquashingKey {
20    bootstrapping_key: SeededLweBootstrapKeyOwned<u128>,
21    modulus_switch_noise_reduction_key: Option<CompressedModulusSwitchNoiseReductionKey<u64>>,
22    message_modulus: MessageModulus,
23    carry_modulus: CarryModulus,
24    output_ciphertext_modulus: CoreCiphertextModulus<u128>,
25}
26
27impl ClientKey {
28    pub fn new_compressed_noise_squashing_key(
29        &self,
30        noise_squashing_private_key: &NoiseSquashingPrivateKey,
31    ) -> CompressedNoiseSquashingKey {
32        let pbs_parameters = self
33            .parameters
34            .pbs_parameters()
35            .expect("NoiseSquashingKey generation requires PBSParameters");
36
37        assert_eq!(
38            pbs_parameters.message_modulus(),
39            noise_squashing_private_key
40                .noise_squashing_parameters()
41                .message_modulus,
42            "Mismatched MessageModulus between ClientKey {:?} and NoiseSquashingPrivateKey {:?}.",
43            pbs_parameters.message_modulus(),
44            noise_squashing_private_key
45                .noise_squashing_parameters()
46                .message_modulus
47        );
48        assert_eq!(
49            pbs_parameters.carry_modulus(),
50            noise_squashing_private_key
51                .noise_squashing_parameters()
52                .carry_modulus,
53            "Mismatched CarryModulus between ClientKey {:?} and NoiseSquashingPrivateKey {:?}.",
54            pbs_parameters.carry_modulus(),
55            noise_squashing_private_key
56                .noise_squashing_parameters()
57                .carry_modulus
58        );
59
60        let noise_squashing_parameters = noise_squashing_private_key.noise_squashing_parameters();
61
62        let (bootstrapping_key, modulus_switch_noise_reduction_key) =
63            ShortintEngine::with_thread_local_mut(|engine| {
64                let seeded_bsk = par_allocate_and_generate_new_seeded_lwe_bootstrap_key(
65                    &self.lwe_secret_key,
66                    noise_squashing_private_key.post_noise_squashing_secret_key(),
67                    noise_squashing_parameters.decomp_base_log,
68                    noise_squashing_parameters.decomp_level_count,
69                    noise_squashing_parameters.glwe_noise_distribution,
70                    noise_squashing_parameters.ciphertext_modulus,
71                    &mut engine.seeder,
72                );
73
74                let modulus_switch_noise_reduction_key = noise_squashing_parameters
75                    .modulus_switch_noise_reduction_params
76                    .map(|p| {
77                        let seed = engine.seeder.seed();
78                        CompressedModulusSwitchNoiseReductionKey::new(
79                            p,
80                            &self.lwe_secret_key,
81                            engine,
82                            pbs_parameters.ciphertext_modulus(),
83                            pbs_parameters.lwe_noise_distribution(),
84                            seed.into(),
85                        )
86                    });
87
88                (seeded_bsk, modulus_switch_noise_reduction_key)
89            });
90
91        CompressedNoiseSquashingKey {
92            bootstrapping_key,
93            modulus_switch_noise_reduction_key,
94            output_ciphertext_modulus: noise_squashing_parameters.ciphertext_modulus,
95            message_modulus: noise_squashing_parameters.message_modulus,
96            carry_modulus: noise_squashing_parameters.carry_modulus,
97        }
98    }
99}
100
101impl CompressedNoiseSquashingKey {
102    pub fn new(
103        client_key: &ClientKey,
104        noise_squashing_private_key: &NoiseSquashingPrivateKey,
105    ) -> Self {
106        client_key.new_compressed_noise_squashing_key(noise_squashing_private_key)
107    }
108
109    pub fn decompress(&self) -> NoiseSquashingKey {
110        let (bootstrapping_key, modulus_switch_noise_reduction_key) = {
111            let std_bsk = self
112                .bootstrapping_key
113                .as_view()
114                .par_decompress_into_lwe_bootstrap_key();
115
116            let mut fbsk = Fourier128LweBootstrapKeyOwned::new(
117                std_bsk.input_lwe_dimension(),
118                std_bsk.glwe_size(),
119                std_bsk.polynomial_size(),
120                std_bsk.decomposition_base_log(),
121                std_bsk.decomposition_level_count(),
122            );
123
124            par_convert_standard_lwe_bootstrap_key_to_fourier_128(&std_bsk, &mut fbsk);
125
126            (
127                fbsk,
128                self.modulus_switch_noise_reduction_key
129                    .as_ref()
130                    .map(|key| key.decompress()),
131            )
132        };
133
134        NoiseSquashingKey {
135            bootstrapping_key,
136            modulus_switch_noise_reduction_key,
137            message_modulus: self.message_modulus,
138            carry_modulus: self.carry_modulus,
139            output_ciphertext_modulus: self.output_ciphertext_modulus,
140        }
141    }
142
143    pub fn message_modulus(&self) -> MessageModulus {
144        self.message_modulus
145    }
146
147    pub fn carry_modulus(&self) -> CarryModulus {
148        self.carry_modulus
149    }
150
151    pub fn output_ciphertext_modulus(&self) -> CoreCiphertextModulus<u128> {
152        self.output_ciphertext_modulus
153    }
154}
155
156impl ParameterSetConformant for CompressedNoiseSquashingKey {
157    type ParameterSet = NoiseSquashingKeyConformanceParams;
158
159    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
160        let Self {
161            bootstrapping_key,
162            modulus_switch_noise_reduction_key,
163            message_modulus,
164            carry_modulus,
165            output_ciphertext_modulus,
166        } = self;
167
168        let Self::ParameterSet {
169            bootstrapping_key_params: expected_bootstrapping_key_params,
170            modulus_switch_noise_reduction_params: expected_modulus_switch_noise_reduction_params,
171            message_modulus: expected_message_modulus,
172            carry_modulus: expected_carry_modulus,
173        } = parameter_set;
174
175        let modulus_switch_key_ok = match (
176            modulus_switch_noise_reduction_key,
177            expected_modulus_switch_noise_reduction_params,
178        ) {
179            (None, None) => true,
180            (None, Some(_)) => false,
181            (Some(_), None) => false,
182            (Some(key), Some(params)) => {
183                let mod_switch_conformance_params =
184                    ModulusSwitchNoiseReductionKeyConformanceParams {
185                        modulus_switch_noise_reduction_params: *params,
186                        lwe_dimension: bootstrapping_key.input_lwe_dimension(),
187                    };
188
189                key.is_conformant(&mod_switch_conformance_params)
190            }
191        };
192
193        modulus_switch_key_ok
194            && bootstrapping_key.is_conformant(expected_bootstrapping_key_params)
195            && *output_ciphertext_modulus == expected_bootstrapping_key_params.ciphertext_modulus
196            && *message_modulus == *expected_message_modulus
197            && *carry_modulus == *expected_carry_modulus
198    }
199}