poulpy/core/src/lwe/keyswtich.rs
core/src/lwe/keyswtich.rsRaw
1use backend::{Backend, FFT64, Module, Scratch, VecZnxOps, ZnxView, ZnxViewMut, ZnxZero};
2use sampling::source::Source;
3
4use crate::{FourierGLWESecret, GLWECiphertext, GLWESecret, GLWESwitchingKey, Infos, LWECiphertext, LWESecret, ScratchCore};
5
6/// A special [GLWESwitchingKey] required to for the conversion from [GLWECiphertext] to [LWECiphertext].
7pub struct GLWEToLWESwitchingKey<D, B: Backend>(GLWESwitchingKey<D, B>);
8
9impl GLWEToLWESwitchingKey<Vec<u8>, FFT64> {
10 pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
11 Self(GLWESwitchingKey::alloc(module, basek, k, rows, 1, rank, 1))
12 }
13
14 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
15 FourierGLWESecret::bytes_of(module, rank)
16 + (GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank) | GLWESecret::bytes_of(module, rank))
17 }
18}
19
20impl<D: AsMut<[u8]> + AsRef<[u8]>> GLWEToLWESwitchingKey<D, FFT64> {
21 pub fn encrypt_sk<DLwe, DGlwe>(
22 &mut self,
23 module: &Module<FFT64>,
24 sk_lwe: &LWESecret<DLwe>,
25 sk_glwe: &GLWESecret<DGlwe>,
26 source_xa: &mut Source,
27 source_xe: &mut Source,
28 sigma: f64,
29 scratch: &mut Scratch,
30 ) where
31 DLwe: AsRef<[u8]>,
32 DGlwe: AsRef<[u8]>,
33 {
34 #[cfg(debug_assertions)]
35 {
36 assert!(sk_lwe.n() <= module.n());
37 }
38
39 let (mut sk_lwe_as_glwe_dft, scratch1) = scratch.tmp_fourier_glwe_secret(module, 1);
40
41 {
42 let (mut sk_lwe_as_glwe, _) = scratch1.tmp_glwe_secret(module, 1);
43 sk_lwe_as_glwe.data.zero();
44 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
45 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data, 0);
46 sk_lwe_as_glwe_dft.set(module, &sk_lwe_as_glwe);
47 }
48
49 self.0.encrypt_sk(
50 module,
51 sk_glwe,
52 &sk_lwe_as_glwe_dft,
53 source_xa,
54 source_xe,
55 sigma,
56 scratch1,
57 );
58 }
59}
60
61/// A special [GLWESwitchingKey] required to for the conversion from [LWECiphertext] to [GLWECiphertext].
62pub struct LWEToGLWESwitchingKey<D, B: Backend>(GLWESwitchingKey<D, B>);
63
64impl LWEToGLWESwitchingKey<Vec<u8>, FFT64> {
65 pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
66 Self(GLWESwitchingKey::alloc(module, basek, k, rows, 1, 1, rank))
67 }
68
69 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
70 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, 1, rank) + GLWESecret::bytes_of(module, 1)
71 }
72}
73
74impl<D: AsMut<[u8]> + AsRef<[u8]>> LWEToGLWESwitchingKey<D, FFT64> {
75 pub fn encrypt_sk<DLwe, DGlwe>(
76 &mut self,
77 module: &Module<FFT64>,
78 sk_lwe: &LWESecret<DLwe>,
79 sk_glwe: &FourierGLWESecret<DGlwe, FFT64>,
80 source_xa: &mut Source,
81 source_xe: &mut Source,
82 sigma: f64,
83 scratch: &mut Scratch,
84 ) where
85 DLwe: AsRef<[u8]>,
86 DGlwe: AsRef<[u8]>,
87 {
88 #[cfg(debug_assertions)]
89 {
90 assert!(sk_lwe.n() <= module.n());
91 }
92
93 let (mut sk_lwe_as_glwe, scratch1) = scratch.tmp_glwe_secret(module, 1);
94 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
95 sk_lwe_as_glwe.data.at_mut(0, 0)[sk_lwe.n()..].fill(0);
96 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data, 0);
97
98 self.0.encrypt_sk(
99 module,
100 &sk_lwe_as_glwe,
101 &sk_glwe,
102 source_xa,
103 source_xe,
104 sigma,
105 scratch1,
106 );
107 }
108}
109
110pub struct LWESwitchingKey<D, B: Backend>(GLWESwitchingKey<D, B>);
111
112impl LWESwitchingKey<Vec<u8>, FFT64> {
113 pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize) -> Self {
114 Self(GLWESwitchingKey::alloc(module, basek, k, rows, 1, 1, 1))
115 }
116
117 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
118 GLWESecret::bytes_of(module, 1)
119 + FourierGLWESecret::bytes_of(module, 1)
120 + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, 1, 1)
121 }
122}
123
124impl<D: AsMut<[u8]> + AsRef<[u8]>> LWESwitchingKey<D, FFT64> {
125 pub fn encrypt_sk<DIn, DOut>(
126 &mut self,
127 module: &Module<FFT64>,
128 sk_lwe_in: &LWESecret<DIn>,
129 sk_lwe_out: &LWESecret<DOut>,
130 source_xa: &mut Source,
131 source_xe: &mut Source,
132 sigma: f64,
133 scratch: &mut Scratch,
134 ) where
135 DIn: AsRef<[u8]>,
136 DOut: AsRef<[u8]>,
137 {
138 #[cfg(debug_assertions)]
139 {
140 assert!(sk_lwe_in.n() <= module.n());
141 assert!(sk_lwe_out.n() <= module.n());
142 }
143
144 let (mut sk_in_glwe, scratch1) = scratch.tmp_glwe_secret(module, 1);
145 let (mut sk_out_glwe, scratch2) = scratch1.tmp_fourier_glwe_secret(module, 1);
146
147 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
148 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
149 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data, 0);
150 sk_out_glwe.set(module, &sk_in_glwe);
151 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
152 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
153 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data, 0);
154
155 self.0.encrypt_sk(
156 module,
157 &sk_in_glwe,
158 &sk_out_glwe,
159 source_xa,
160 source_xe,
161 sigma,
162 scratch2,
163 );
164 }
165}
166
167impl LWECiphertext<Vec<u8>> {
168 pub fn from_glwe_scratch_space(
169 module: &Module<FFT64>,
170 basek: usize,
171 k_lwe: usize,
172 k_glwe: usize,
173 k_ksk: usize,
174 rank: usize,
175 ) -> usize {
176 GLWECiphertext::bytes_of(module, basek, k_lwe, 1)
177 + GLWECiphertext::keyswitch_scratch_space(module, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
178 }
179
180 pub fn keyswitch_scratch_space(
181 module: &Module<FFT64>,
182 basek: usize,
183 k_lwe_out: usize,
184 k_lwe_in: usize,
185 k_ksk: usize,
186 ) -> usize {
187 GLWECiphertext::bytes_of(module, basek, k_lwe_out.max(k_lwe_in), 1)
188 + GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, k_lwe_out, k_ksk, 1, 1)
189 }
190}
191
192impl<DLwe: AsRef<[u8]> + AsMut<[u8]>> LWECiphertext<DLwe> {
193 pub fn sample_extract<DGlwe>(&mut self, a: &GLWECiphertext<DGlwe>)
194 where
195 DGlwe: AsRef<[u8]>,
196 {
197 #[cfg(debug_assertions)]
198 {
199 assert!(self.n() <= a.n());
200 }
201
202 let min_size: usize = self.size().min(a.size());
203 let n: usize = self.n();
204
205 self.data.zero();
206 (0..min_size).for_each(|i| {
207 let data_lwe: &mut [i64] = self.data.at_mut(0, i);
208 data_lwe[0] = a.data.at(0, i)[0];
209 data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
210 });
211 }
212
213 pub fn from_glwe<DGlwe, DKs>(
214 &mut self,
215 module: &Module<FFT64>,
216 a: &GLWECiphertext<DGlwe>,
217 ks: &GLWEToLWESwitchingKey<DKs, FFT64>,
218 scratch: &mut Scratch,
219 ) where
220 DGlwe: AsRef<[u8]>,
221 DKs: AsRef<[u8]>,
222 {
223 #[cfg(debug_assertions)]
224 {
225 assert_eq!(self.basek(), a.basek());
226 }
227 let (mut tmp_glwe, scratch1) = scratch.tmp_glwe_ct(module, a.basek(), self.k(), 1);
228 tmp_glwe.keyswitch(module, a, &ks.0, scratch1);
229 self.sample_extract(&tmp_glwe);
230 }
231
232 pub fn keyswitch<A, DKs>(
233 &mut self,
234 module: &Module<FFT64>,
235 a: &LWECiphertext<A>,
236 ksk: &LWESwitchingKey<DKs, FFT64>,
237 scratch: &mut Scratch,
238 ) where
239 A: AsRef<[u8]>,
240 DKs: AsRef<[u8]>,
241 {
242 #[cfg(debug_assertions)]
243 {
244 assert!(self.n() <= module.n());
245 assert!(a.n() <= module.n());
246 assert_eq!(self.basek(), a.basek());
247 }
248
249 let max_k: usize = self.k().max(a.k());
250 let basek: usize = self.basek();
251
252 let (mut glwe, scratch1) = scratch.tmp_glwe_ct(&module, basek, max_k, 1);
253 glwe.data.zero();
254
255 let n_lwe: usize = a.n();
256
257 (0..a.size()).for_each(|i| {
258 let data_lwe: &[i64] = a.data.at(0, i);
259 glwe.data.at_mut(0, i)[0] = data_lwe[0];
260 glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
261 });
262
263 glwe.keyswitch_inplace(module, &ksk.0, scratch1);
264
265 self.sample_extract(&glwe);
266 }
267}
268
269impl GLWECiphertext<Vec<u8>> {
270 pub fn from_lwe_scratch_space(
271 module: &Module<FFT64>,
272 basek: usize,
273 k_lwe: usize,
274 k_glwe: usize,
275 k_ksk: usize,
276 rank: usize,
277 ) -> usize {
278 GLWECiphertext::keyswitch_scratch_space(module, basek, k_glwe, k_lwe, k_ksk, 1, 1, rank)
279 + GLWECiphertext::bytes_of(module, basek, k_lwe, 1)
280 }
281}
282
283impl<D: AsRef<[u8]> + AsMut<[u8]>> GLWECiphertext<D> {
284 pub fn from_lwe<DLwe, DKsk>(
285 &mut self,
286 module: &Module<FFT64>,
287 lwe: &LWECiphertext<DLwe>,
288 ksk: &LWEToGLWESwitchingKey<DKsk, FFT64>,
289 scratch: &mut Scratch,
290 ) where
291 DLwe: AsRef<[u8]>,
292 DKsk: AsRef<[u8]>,
293 {
294 #[cfg(debug_assertions)]
295 {
296 assert!(lwe.n() <= self.n());
297 assert_eq!(self.basek(), self.basek());
298 }
299
300 let (mut glwe, scratch1) = scratch.tmp_glwe_ct(module, lwe.basek(), lwe.k(), 1);
301 glwe.data.zero();
302
303 let n_lwe: usize = lwe.n();
304
305 (0..lwe.size()).for_each(|i| {
306 let data_lwe: &[i64] = lwe.data.at(0, i);
307 glwe.data.at_mut(0, i)[0] = data_lwe[0];
308 glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
309 });
310
311 self.keyswitch(module, &glwe, &ksk.0, scratch1);
312 }
313}