Skip to main content

tfhe/high_level_api/keys/
key_switching_key.rs

1use tfhe_versionable::Versionize;
2
3use crate::backward_compatibility::keys::KeySwitchingKeyVersions;
4use crate::high_level_api::integers::{FheIntId, FheUintId};
5use crate::high_level_api::re_randomization::ReRandomizationMetadata;
6use crate::integer::BooleanBlock;
7use crate::named::Named;
8use crate::prelude::FheKeyswitch;
9pub use crate::shortint::parameters::key_switching::ShortintKeySwitchingParameters;
10use crate::{ClientKey, FheBool, FheInt, FheUint, ServerKey, Tag};
11use std::fmt::{Display, Formatter};
12
13#[derive(Copy, Clone, Debug)]
14pub struct IncompatibleParameters;
15
16impl Display for IncompatibleParameters {
17    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
18        write!(f, "{self:?}")
19    }
20}
21
22impl std::error::Error for IncompatibleParameters {}
23
24#[derive(serde::Deserialize, serde::Serialize, Versionize)]
25#[versionize(KeySwitchingKeyVersions)]
26pub struct KeySwitchingKey {
27    key: crate::integer::key_switching_key::KeySwitchingKey,
28    tag_in: Tag,
29    tag_out: Tag,
30}
31
32impl KeySwitchingKey {
33    pub fn new(
34        key_pair_from: (&ClientKey, &ServerKey),
35        key_pair_to: (&ClientKey, &ServerKey),
36    ) -> Result<Self, IncompatibleParameters> {
37        let params_from = key_pair_from.0.key.block_parameters();
38        let params_to = key_pair_to.0.key.block_parameters();
39
40        if params_to != params_from {
41            return Err(IncompatibleParameters);
42        }
43
44        // params_to == params_from, so we can use the parameters from params_to
45        let params = ShortintKeySwitchingParameters {
46            ks_base_log: params_to.ks_base_log(),
47            ks_level: params_to.ks_level(),
48            destination_key: params_to.encryption_key_choice(),
49        };
50
51        Ok(Self::with_parameters(key_pair_from, key_pair_to, params))
52    }
53
54    pub fn with_parameters(
55        key_pair_from: (&ClientKey, &ServerKey),
56        key_pair_to: (&ClientKey, &ServerKey),
57        params: ShortintKeySwitchingParameters,
58    ) -> Self {
59        Self {
60            key: crate::integer::key_switching_key::KeySwitchingKey::new(
61                (&key_pair_from.0.key.key, Some(&key_pair_from.1.key.key)),
62                (&key_pair_to.0.key.key, &key_pair_to.1.key.key),
63                params,
64            ),
65            tag_in: key_pair_from.0.tag.clone(),
66            tag_out: key_pair_to.0.tag.clone(),
67        }
68    }
69
70    pub fn tag_in(&self) -> &Tag {
71        &self.tag_in
72    }
73
74    pub fn tag_out(&self) -> &Tag {
75        &self.tag_out
76    }
77}
78
79impl<Id> FheKeyswitch<FheUint<Id>> for KeySwitchingKey
80where
81    Id: FheUintId,
82{
83    fn keyswitch(&self, input: &FheUint<Id>) -> FheUint<Id> {
84        let radix = input.ciphertext.on_cpu();
85        let casted = self.key.cast(&*radix);
86        FheUint::new(
87            casted,
88            self.tag_out.clone(),
89            ReRandomizationMetadata::default(),
90        )
91    }
92}
93
94impl<Id> FheKeyswitch<FheInt<Id>> for KeySwitchingKey
95where
96    Id: FheIntId,
97{
98    fn keyswitch(&self, input: &FheInt<Id>) -> FheInt<Id> {
99        let radix = input.ciphertext.on_cpu();
100        let casted = self.key.cast(&*radix);
101        FheInt::new(
102            casted,
103            self.tag_out.clone(),
104            ReRandomizationMetadata::default(),
105        )
106    }
107}
108
109impl FheKeyswitch<FheBool> for KeySwitchingKey {
110    fn keyswitch(&self, input: &FheBool) -> FheBool {
111        let boolean_block = input.ciphertext.on_cpu();
112        let casted = self.key.key.cast(boolean_block.as_ref());
113        FheBool::new(
114            BooleanBlock::new_unchecked(casted),
115            self.tag_out.clone(),
116            ReRandomizationMetadata::default(),
117        )
118    }
119}
120
121impl Named for KeySwitchingKey {
122    const NAME: &'static str = "high_level_api::KeySwitchingKey";
123}