use tfhe_versionable::Versionize;
use crate::core_crypto::prelude::{
allocate_and_generate_new_binary_glwe_secret_key,
allocate_and_generate_new_lwe_packing_keyswitch_key,
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key,
par_allocate_and_generate_new_seeded_lwe_bootstrap_key,
par_allocate_and_generate_new_seeded_lwe_multi_bit_bootstrap_key, GlweSecretKey,
GlweSecretKeyOwned,
};
use crate::shortint::backward_compatibility::list_compression::{
CompressionPrivateKeysVersions, NoiseSquashingCompressionPrivateKeyVersions,
};
use crate::shortint::ciphertext::CompressedSquashedNoiseCiphertextList;
use crate::shortint::client_key::ClientKey;
use crate::shortint::engine::ShortintEngine;
use crate::shortint::noise_squashing::NoiseSquashingPrivateKeyView;
use crate::shortint::parameters::{CompressionParameters, NoiseSquashingCompressionParameters};
use crate::shortint::server_key::{
CompressedModulusSwitchConfiguration, ModulusSwitchConfiguration, ShortintBootstrappingKey,
ShortintCompressedBootstrappingKey,
};
use crate::shortint::{EncryptionKeyChoice, ShortintParameterSet};
use std::fmt::Debug;
use super::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
};
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Versionize)]
#[versionize(CompressionPrivateKeysVersions)]
pub struct CompressionPrivateKeys {
pub post_packing_ks_key: GlweSecretKeyOwned<u64>,
pub params: CompressionParameters,
}
impl CompressionPrivateKeys {
pub(crate) fn new_compression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressionKey {
ShortintEngine::with_thread_local_mut(|engine| {
self.new_compression_key_with_engine(glwe_secret_key, pbs_params, engine)
})
}
pub(crate) fn new_compression_key_with_engine(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
engine: &mut ShortintEngine,
) -> CompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
assert!(
compression_params.storage_log_modulus().0
<= pbs_params
.polynomial_size()
.to_blind_rotation_input_modulus_log()
.0,
"Compression parameters say to store more bits than useful"
);
let packing_key_switching_key = allocate_and_generate_new_lwe_packing_keyswitch_key(
&glwe_secret_key.as_lwe_secret_key(),
&self.post_packing_ks_key,
compression_params.packing_ks_base_log(),
compression_params.packing_ks_level(),
compression_params.packing_ks_key_noise_distribution(),
pbs_params.ciphertext_modulus(),
&mut engine.encryption_generator,
);
CompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe(),
storage_log_modulus: compression_params.storage_log_modulus(),
}
}
pub(crate) fn new_compressed_compression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressedCompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key(
&glwe_secret_key.as_lwe_secret_key(),
&self.post_packing_ks_key,
compression_params.packing_ks_base_log(),
compression_params.packing_ks_level(),
compression_params.packing_ks_key_noise_distribution(),
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
CompressedCompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe(),
storage_log_modulus: compression_params.storage_log_modulus(),
}
}
pub(crate) fn new_decompression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> DecompressionKey {
self.new_decompression_key_with_params(glwe_secret_key, pbs_params, self.params)
}
pub(crate) fn new_decompression_key_with_params(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
compression_params: CompressionParameters,
) -> DecompressionKey {
ShortintEngine::with_thread_local_mut(|engine| {
self.new_decompression_key_with_params_and_engine(
glwe_secret_key,
pbs_params,
compression_params,
engine,
)
})
}
pub(crate) fn new_decompression_key_with_params_and_engine(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
compression_params: CompressionParameters,
engine: &mut ShortintEngine,
) -> DecompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
match compression_params {
CompressionParameters::Classic(classic_compression_parameters) => {
let blind_rotate_key = engine.new_classic_bootstrapping_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
pbs_params.glwe_noise_distribution(),
classic_compression_parameters.br_base_log,
classic_compression_parameters.br_level,
pbs_params.ciphertext_modulus(),
);
DecompressionKey {
bsk: ShortintBootstrappingKey::Classic {
bsk: blind_rotate_key,
modulus_switch_noise_reduction_key: ModulusSwitchConfiguration::Standard,
},
lwe_per_glwe: classic_compression_parameters.lwe_per_glwe,
}
}
CompressionParameters::MultiBit(multi_bit_compression_parameters) => {
let multi_bit_blind_rotate_key = engine.new_multibit_bootstrapping_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
pbs_params.glwe_noise_distribution(),
multi_bit_compression_parameters.br_base_log,
multi_bit_compression_parameters.br_level,
multi_bit_compression_parameters.decompression_grouping_factor,
pbs_params.ciphertext_modulus(),
);
let thread_count = ShortintEngine::get_thread_count_for_multi_bit_pbs(
pbs_params.lwe_dimension(),
pbs_params.glwe_dimension(),
pbs_params.polynomial_size(),
pbs_params.pbs_base_log(),
pbs_params.pbs_level(),
multi_bit_compression_parameters.decompression_grouping_factor,
);
DecompressionKey {
bsk: ShortintBootstrappingKey::MultiBit {
fourier_bsk: multi_bit_blind_rotate_key,
thread_count,
deterministic_execution: true,
},
lwe_per_glwe: multi_bit_compression_parameters.lwe_per_glwe,
}
}
}
}
pub(crate) fn new_compressed_decompression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressedDecompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
match compression_params {
CompressionParameters::Classic(classic_compression_parameters) => {
let blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
par_allocate_and_generate_new_seeded_lwe_bootstrap_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
classic_compression_parameters.br_base_log,
classic_compression_parameters.br_level,
pbs_params.glwe_noise_distribution(),
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
CompressedDecompressionKey {
bsk: ShortintCompressedBootstrappingKey::Classic {
bsk: blind_rotate_key,
modulus_switch_noise_reduction_key:
CompressedModulusSwitchConfiguration::Standard,
},
lwe_per_glwe: classic_compression_parameters.lwe_per_glwe,
}
}
CompressionParameters::MultiBit(multi_bit_compression_parameters) => {
let multi_bit_blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
par_allocate_and_generate_new_seeded_lwe_multi_bit_bootstrap_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
multi_bit_compression_parameters.br_base_log,
multi_bit_compression_parameters.br_level,
pbs_params.glwe_noise_distribution(),
multi_bit_compression_parameters.decompression_grouping_factor,
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
CompressedDecompressionKey {
bsk: ShortintCompressedBootstrappingKey::MultiBit {
seeded_bsk: multi_bit_blind_rotate_key,
deterministic_execution: true,
},
lwe_per_glwe: multi_bit_compression_parameters.lwe_per_glwe,
}
}
}
}
}
impl ClientKey {
pub fn new_compression_private_key(
&self,
params: CompressionParameters,
) -> CompressionPrivateKeys {
ShortintEngine::with_thread_local_mut(|engine| {
self.new_compression_private_key_with_engine(params, engine)
})
}
pub(crate) fn new_compression_private_key_with_engine(
&self,
params: CompressionParameters,
engine: &mut ShortintEngine,
) -> CompressionPrivateKeys {
assert_eq!(
self.parameters().encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let post_packing_ks_key = allocate_and_generate_new_binary_glwe_secret_key(
params.packing_ks_glwe_dimension(),
params.packing_ks_polynomial_size(),
&mut engine.secret_generator,
);
CompressionPrivateKeys {
post_packing_ks_key,
params,
}
}
}
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Versionize)]
#[versionize(NoiseSquashingCompressionPrivateKeyVersions)]
pub struct NoiseSquashingCompressionPrivateKey {
pub(crate) post_packing_ks_key: GlweSecretKeyOwned<u128>,
pub(crate) params: NoiseSquashingCompressionParameters,
}
impl NoiseSquashingCompressionPrivateKey {
pub fn new(params: NoiseSquashingCompressionParameters) -> Self {
let post_packing_ks_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_binary_glwe_secret_key(
params.packing_ks_glwe_dimension,
params.packing_ks_polynomial_size,
&mut engine.secret_generator,
)
});
Self {
post_packing_ks_key,
params,
}
}
pub fn from_raw_parts(
post_packing_ks_key: GlweSecretKeyOwned<u128>,
params: NoiseSquashingCompressionParameters,
) -> Self {
assert_eq!(
post_packing_ks_key.polynomial_size(),
params.packing_ks_polynomial_size,
"Invalid polynomial size for NoiseSquashingCompressionPrivateKey, expected {}, got {}",
params.packing_ks_polynomial_size.0,
post_packing_ks_key.polynomial_size().0,
);
assert_eq!(
post_packing_ks_key.glwe_dimension(),
params.packing_ks_glwe_dimension,
"Invalid GLWE dimension for NoiseSquashingCompressionPrivateKey, expected {}, got {}",
params.packing_ks_glwe_dimension.0,
post_packing_ks_key.glwe_dimension().0,
);
Self {
post_packing_ks_key,
params,
}
}
pub fn into_raw_parts(
self,
) -> (
GlweSecretKeyOwned<u128>,
NoiseSquashingCompressionParameters,
) {
let Self {
post_packing_ks_key,
params,
} = self;
(post_packing_ks_key, params)
}
pub fn unpack_and_decrypt_squashed_noise_ciphertexts(
&self,
compressed_list: &CompressedSquashedNoiseCiphertextList,
) -> Vec<u128> {
let decryption_key = NoiseSquashingPrivateKeyView::from(self);
(0..compressed_list.len())
.map(|i| {
let ciphertext = compressed_list.unpack(i).unwrap();
decryption_key.decrypt_squashed_noise_ciphertext(&ciphertext)
})
.collect()
}
pub fn post_packing_ks_key(&self) -> &GlweSecretKeyOwned<u128> {
&self.post_packing_ks_key
}
pub fn params(&self) -> NoiseSquashingCompressionParameters {
self.params
}
}