poulpy/core/src/keyswitch_key.rs
core/src/keyswitch_key.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft, ScalarZnxDftToRef,
3 ScalarZnxToRef, Scratch, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef,
4};
5use sampling::source::Source;
6
7use crate::{
8 elem::{GetRow, Infos, SetRow},
9 gglwe_ciphertext::GGLWECiphertext,
10 ggsw_ciphertext::GGSWCiphertext,
11 glwe_ciphertext_fourier::GLWECiphertextFourier,
12 keys::SecretKeyFourier,
13 vec_glwe_product::{VecGLWEProduct, VecGLWEProductScratchSpace},
14};
15
16pub struct GLWESwitchingKey<Data, B: Backend>(pub(crate) GGLWECiphertext<Data, B>);
17
18impl GLWESwitchingKey<Vec<u8>, FFT64> {
19 pub fn new(module: &Module<FFT64>, base2k: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> Self {
20 GLWESwitchingKey(GGLWECiphertext::new(
21 module, base2k, k, rows, rank_in, rank_out,
22 ))
23 }
24}
25
26impl<T, B: Backend> Infos for GLWESwitchingKey<T, B> {
27 type Inner = MatZnxDft<T, B>;
28
29 fn inner(&self) -> &Self::Inner {
30 &self.0.inner()
31 }
32
33 fn basek(&self) -> usize {
34 self.0.basek()
35 }
36
37 fn k(&self) -> usize {
38 self.0.k()
39 }
40}
41
42impl<DataSelf, B: Backend> MatZnxDftToMut<B> for GLWESwitchingKey<DataSelf, B>
43where
44 MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
45{
46 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
47 self.0.data.to_mut()
48 }
49}
50
51impl<DataSelf, B: Backend> MatZnxDftToRef<B> for GLWESwitchingKey<DataSelf, B>
52where
53 MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
54{
55 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
56 self.0.data.to_ref()
57 }
58}
59
60impl<C> GetRow<FFT64> for GLWESwitchingKey<C, FFT64>
61where
62 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
63{
64 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
65 where
66 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
67 {
68 #[cfg(debug_assertions)]
69 {
70 assert_eq!(col_j, 0);
71 }
72 module.vmp_extract_row(res, self, row_i, col_j);
73 }
74}
75
76impl<C> SetRow<FFT64> for GLWESwitchingKey<C, FFT64>
77where
78 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
79{
80 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
81 where
82 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
83 {
84 #[cfg(debug_assertions)]
85 {
86 assert_eq!(col_j, 0);
87 }
88 module.vmp_prepare_row(self, row_i, col_j, a);
89 }
90}
91
92impl GLWESwitchingKey<Vec<u8>, FFT64> {
93 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
94 GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
95 }
96
97 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
98 GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
99 }
100}
101
102impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
103where
104 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
105{
106 pub fn encrypt_sk<DataPt, DataSk>(
107 &mut self,
108 module: &Module<FFT64>,
109 pt: &ScalarZnx<DataPt>,
110 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
111 source_xa: &mut Source,
112 source_xe: &mut Source,
113 sigma: f64,
114 scratch: &mut Scratch,
115 ) where
116 ScalarZnx<DataPt>: ScalarZnxToRef,
117 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
118 {
119 self.0
120 .encrypt_sk(module, pt, sk_dft, source_xa, source_xe, sigma, scratch);
121 }
122}
123
124impl GLWESwitchingKey<Vec<u8>, FFT64> {
125 pub fn keyswitch_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
126 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
127 module, res_size, lhs, rhs,
128 )
129 }
130
131 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
132 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_inplace_scratch_space(
133 module, res_size, rhs,
134 )
135 }
136
137 pub fn external_product_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
138 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
139 module, res_size, lhs, rhs,
140 )
141 }
142
143 pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
144 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
145 module, res_size, rhs,
146 )
147 }
148}
149
150impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
151where
152 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
153{
154 pub fn keyswitch<DataLhs, DataRhs>(
155 &mut self,
156 module: &Module<FFT64>,
157 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
158 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
159 scratch: &mut base2k::Scratch,
160 ) where
161 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
162 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
163 {
164 rhs.0
165 .prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
166 }
167
168 pub fn keyswitch_inplace<DataRhs>(
169 &mut self,
170 module: &Module<FFT64>,
171 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
172 scratch: &mut base2k::Scratch,
173 ) where
174 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
175 {
176 rhs.0
177 .prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
178 }
179
180 pub fn external_product<DataLhs, DataRhs>(
181 &mut self,
182 module: &Module<FFT64>,
183 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
184 rhs: &GGSWCiphertext<DataRhs, FFT64>,
185 scratch: &mut Scratch,
186 ) where
187 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
188 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
189 {
190 rhs.prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
191 }
192
193 pub fn external_product_inplace<DataRhs>(
194 &mut self,
195 module: &Module<FFT64>,
196 rhs: &GGSWCiphertext<DataRhs, FFT64>,
197 scratch: &mut Scratch,
198 ) where
199 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
200 {
201 rhs.prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
202 }
203}