use crate::core_crypto::prelude::*;
use crate::integer::ciphertext::NoiseSquashingCompressionKey;
use crate::integer::compression_keys::CompressionKey;
use crate::shortint::noise_squashing::atomic_pattern::compressed::CompressedAtomicPatternNoiseSquashingKey;
use crate::high_level_api::keys::cpk_re_randomization::ReRandomizationKey;
#[cfg(feature = "gpu")]
use crate::high_level_api::keys::inner::CudaReRandomizationKey;
use crate::integer::compression_keys::CompressedDecompressionKey;
use crate::integer::noise_squashing::CompressedNoiseSquashingKey;
use crate::integer::oprf::ExpandedOprfServerKey;
use crate::shortint::atomic_pattern::expanded::{
ExpandedAtomicPatternServerKey, ExpandedKS32AtomicPatternServerKey,
ExpandedStandardAtomicPatternServerKey,
};
pub use crate::shortint::noise_squashing::atomic_pattern::ExpandedAtomicPatternNoiseSquashingKey;
pub use crate::shortint::noise_squashing::ExpandedNoiseSquashingKey;
pub use crate::shortint::server_key::expanded::{
ShortintExpandedBootstrappingKey, ShortintExpandedServerKey,
};
use crate::transciphering::ExpandedTranscipheringServerKey;
#[derive(PartialEq)]
pub struct ExpandedDecompressionKey {
pub bsk: ShortintExpandedBootstrappingKey<u64, u64>,
pub lwe_per_glwe: LweCiphertextCount,
}
#[derive(PartialEq)]
pub struct IntegerExpandedServerKey {
pub compute_key: ShortintExpandedServerKey,
pub cpk_key_switching_key_material:
Option<crate::integer::key_switching_key::KeySwitchingKeyMaterial>,
pub compression_key: Option<CompressionKey>,
pub decompression_key: Option<ExpandedDecompressionKey>,
pub noise_squashing_key: Option<ExpandedNoiseSquashingKey>,
pub noise_squashing_compression_key: Option<NoiseSquashingCompressionKey>,
pub cpk_re_randomization_key: Option<ReRandomizationKey>,
pub oprf_key: Option<ExpandedOprfServerKey>,
pub transciphering_key: Option<ExpandedTranscipheringServerKey>,
}
pub(crate) fn compute_key_to_cpu(
compute_key: ShortintExpandedServerKey,
) -> crate::integer::ServerKey {
use crate::shortint::atomic_pattern::{
AtomicPatternServerKey, KS32AtomicPatternServerKey, StandardAtomicPatternServerKey,
};
let atomic_pattern_key = match compute_key.atomic_pattern {
ExpandedAtomicPatternServerKey::Standard(std_ap) => {
let ExpandedStandardAtomicPatternServerKey {
key_switching_key,
bootstrapping_key,
pbs_order,
} = std_ap;
AtomicPatternServerKey::Standard(StandardAtomicPatternServerKey {
key_switching_key,
bootstrapping_key: bootstrapping_key.into_fourier(),
pbs_order,
})
}
ExpandedAtomicPatternServerKey::KeySwitch32(ks32_ap) => {
let ExpandedKS32AtomicPatternServerKey {
key_switching_key,
bootstrapping_key,
ciphertext_modulus,
} = ks32_ap;
AtomicPatternServerKey::KeySwitch32(KS32AtomicPatternServerKey {
key_switching_key,
bootstrapping_key: bootstrapping_key.into_fourier(),
ciphertext_modulus,
})
}
};
crate::integer::ServerKey::from_raw_parts(crate::shortint::ServerKey::from_raw_parts(
atomic_pattern_key,
compute_key.message_modulus,
compute_key.carry_modulus,
compute_key.max_degree,
compute_key.max_noise_level,
))
}
pub(crate) fn expanded_noise_squashing_key_to_cpu(
ns_key: ExpandedNoiseSquashingKey,
) -> crate::integer::noise_squashing::NoiseSquashingKey {
use crate::shortint::noise_squashing::atomic_pattern::ks32::KS32AtomicPatternNoiseSquashingKey;
use crate::shortint::noise_squashing::atomic_pattern::standard::StandardAtomicPatternNoiseSquashingKey;
use crate::shortint::noise_squashing::atomic_pattern::AtomicPatternNoiseSquashingKey;
let (ap, msg_mod, carry_mod, ct_mod) = ns_key.into_raw_parts();
let ap = match ap {
ExpandedAtomicPatternNoiseSquashingKey::Standard(bootstrapping_key) => {
AtomicPatternNoiseSquashingKey::Standard(
StandardAtomicPatternNoiseSquashingKey::from_raw_parts(
bootstrapping_key.into_fourier(),
),
)
}
ExpandedAtomicPatternNoiseSquashingKey::KeySwitch32(bootstrapping_key) => {
AtomicPatternNoiseSquashingKey::KeySwitch32(
KS32AtomicPatternNoiseSquashingKey::from_raw_parts(
bootstrapping_key.into_fourier(),
),
)
}
};
crate::integer::noise_squashing::NoiseSquashingKey::from_raw_parts(
crate::shortint::noise_squashing::NoiseSquashingKey::from_raw_parts(
ap, msg_mod, carry_mod, ct_mod,
),
)
}
pub(crate) fn expanded_decompression_key_to_cpu(
key: ExpandedDecompressionKey,
) -> crate::integer::compression_keys::DecompressionKey {
let ExpandedDecompressionKey { bsk, lwe_per_glwe } = key;
let bsk = bsk.into_fourier();
crate::integer::compression_keys::DecompressionKey::from_raw_parts(
crate::shortint::list_compression::DecompressionKey { bsk, lwe_per_glwe },
)
}
impl IntegerExpandedServerKey {
pub fn convert_to_cpu(self) -> crate::high_level_api::keys::IntegerServerKey {
use crate::high_level_api::keys::IntegerServerKey;
let Self {
compute_key,
cpk_key_switching_key_material,
compression_key,
decompression_key,
noise_squashing_key,
noise_squashing_compression_key,
cpk_re_randomization_key,
oprf_key,
transciphering_key,
} = self;
IntegerServerKey {
key: compute_key_to_cpu(compute_key),
cpk_key_switching_key_material,
compression_key,
decompression_key: decompression_key.map(expanded_decompression_key_to_cpu),
noise_squashing_key: noise_squashing_key.map(expanded_noise_squashing_key_to_cpu),
noise_squashing_compression_key,
cpk_re_randomization_key,
oprf_key: oprf_key.map(|oprf_key| oprf_key.to_fourier()),
transciphering_key: transciphering_key.map(|key| key.to_fourier()),
}
}
}
#[cfg(feature = "gpu")]
impl IntegerExpandedServerKey {
pub fn convert_to_gpu(
&self,
streams: &crate::core_crypto::gpu::CudaStreams,
) -> crate::Result<crate::high_level_api::keys::inner::IntegerCudaServerKey> {
use crate::high_level_api::keys::cpk_re_randomization::{
CudaReRandomizationKeySwitchingKey, ReRandomizationKeySwitchingKey,
};
use crate::integer::gpu::key_switching_key::CudaKeySwitchingKeyMaterial;
use crate::integer::gpu::list_compression::server_keys::{
CudaCompressionKey, CudaDecompressionKey, CudaNoiseSquashingCompressionKey,
};
use crate::integer::gpu::noise_squashing::keys::CudaNoiseSquashingKey;
use crate::integer::gpu::CudaServerKey;
let Self {
compute_key,
cpk_key_switching_key_material,
compression_key,
decompression_key,
noise_squashing_key,
noise_squashing_compression_key,
cpk_re_randomization_key,
oprf_key,
transciphering_key: _,
} = self;
let key = CudaServerKey::from_expanded_server_key(compute_key, streams)?;
let cpk_key_switching_key_material = cpk_key_switching_key_material.as_ref().map(|ksk| {
CudaKeySwitchingKeyMaterial::from_key_switching_key_material(&ksk.as_view(), streams)
});
let compression_key = compression_key
.as_ref()
.map(|ck| CudaCompressionKey::from_compression_key(ck, streams));
let decompression_key = decompression_key
.as_ref()
.map(|dk| {
CudaDecompressionKey::from_expanded_decompression_key(
dk,
compute_key.glwe_dimension(),
compute_key.polynomial_size(),
compute_key.message_modulus,
compute_key.carry_modulus,
compute_key.ciphertext_modulus,
streams,
)
})
.transpose()?;
let noise_squashing_key = noise_squashing_key
.as_ref()
.map(|nsk| CudaNoiseSquashingKey::from_expanded_noise_squashing_key(nsk, streams));
let noise_squashing_compression_key =
noise_squashing_compression_key.as_ref().map(|nsck| {
CudaNoiseSquashingCompressionKey::from_noise_squashing_compression_key(
nsck, streams,
)
});
let cpk_re_randomization_key =
cpk_re_randomization_key
.as_ref()
.map(|re_rand_key| match re_rand_key {
ReRandomizationKey::LegacyDedicatedCPK { ksk } => match ksk {
ReRandomizationKeySwitchingKey::UseCPKEncryptionKSK => {
CudaReRandomizationKey::LegacyDedicatedCPK {
ksk: CudaReRandomizationKeySwitchingKey::UseCPKEncryptionKSK,
}
}
ReRandomizationKeySwitchingKey::DedicatedKSK(ksk) => {
CudaReRandomizationKey::LegacyDedicatedCPK {
ksk: CudaReRandomizationKeySwitchingKey::DedicatedKSK(
CudaKeySwitchingKeyMaterial::from_key_switching_key_material(
&ksk.as_view(),
streams,
),
),
}
}
},
ReRandomizationKey::DerivedCPKWithoutKeySwitch { cpk } => {
CudaReRandomizationKey::DerivedCPKWithoutKeySwitch { cpk: cpk.clone() }
}
});
let oprf_key = oprf_key.as_ref().map(|expanded_oprf_key| {
crate::integer::gpu::CudaOprfServerKey::from_expanded_cpu(expanded_oprf_key, streams)
});
Ok(crate::high_level_api::keys::inner::IntegerCudaServerKey {
key,
cpk_key_switching_key_material,
compression_key,
decompression_key,
noise_squashing_key,
noise_squashing_compression_key,
cpk_re_randomization_key,
oprf_key,
})
}
}
impl CompressedDecompressionKey {
pub fn expand(&self) -> ExpandedDecompressionKey {
ExpandedDecompressionKey {
bsk: self.key.bsk.expand(),
lwe_per_glwe: self.key.lwe_per_glwe,
}
}
}
impl CompressedNoiseSquashingKey {
pub fn expand(&self) -> ExpandedNoiseSquashingKey {
let expanded_ap = match self.key.atomic_pattern() {
CompressedAtomicPatternNoiseSquashingKey::Standard(compressed_std) => {
ExpandedAtomicPatternNoiseSquashingKey::Standard(
compressed_std.bootstrapping_key().expand(),
)
}
CompressedAtomicPatternNoiseSquashingKey::KeySwitch32(compressed_ks32) => {
ExpandedAtomicPatternNoiseSquashingKey::KeySwitch32(
compressed_ks32.bootstrapping_key().expand(),
)
}
};
ExpandedNoiseSquashingKey::from_raw_parts(
expanded_ap,
self.key.message_modulus(),
self.key.carry_modulus(),
self.key.output_ciphertext_modulus(),
)
}
}