poulpy/core/src/lwe/keyswtich_layouts.rs
core/src/lwe/keyswtich_layouts.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxAutomorphismInplace,
4 VecZnxSwithcDegree, ZnxView, ZnxViewMut, ZnxZero,
5 },
6 layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, Scratch, WriterTo},
7};
8use sampling::source::Source;
9
10use crate::{
11 GGLWEEncryptSkFamily, GLWECiphertext, GLWEKeyswitchFamily, GLWESecret, GLWESecretExec, GLWESwitchingKey, Infos,
12 LWECiphertext, LWESecret, TakeGLWESecret, TakeGLWESecretExec,
13};
14
15/// A special [GLWESwitchingKey] required to for the conversion from [GLWECiphertext] to [LWECiphertext].
16#[derive(PartialEq, Eq)]
17pub struct GLWEToLWESwitchingKey<D: Data>(pub(crate) GLWESwitchingKey<D>);
18
19impl<D: Data> Infos for GLWEToLWESwitchingKey<D> {
20 type Inner = MatZnx<D>;
21
22 fn inner(&self) -> &Self::Inner {
23 &self.0.inner()
24 }
25
26 fn basek(&self) -> usize {
27 self.0.basek()
28 }
29
30 fn k(&self) -> usize {
31 self.0.k()
32 }
33}
34
35impl<D: Data> GLWEToLWESwitchingKey<D> {
36 pub fn digits(&self) -> usize {
37 self.0.digits()
38 }
39
40 pub fn rank(&self) -> usize {
41 self.0.rank()
42 }
43
44 pub fn rank_in(&self) -> usize {
45 self.0.rank_in()
46 }
47
48 pub fn rank_out(&self) -> usize {
49 self.0.rank_out()
50 }
51}
52
53impl<D: DataMut> ReaderFrom for GLWEToLWESwitchingKey<D> {
54 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
55 self.0.read_from(reader)
56 }
57}
58
59impl<D: DataRef> WriterTo for GLWEToLWESwitchingKey<D> {
60 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
61 self.0.write_to(writer)
62 }
63}
64
65impl GLWEToLWESwitchingKey<Vec<u8>> {
66 pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, rank_in: usize) -> Self {
67 Self(GLWESwitchingKey::alloc(n, basek, k, rows, 1, rank_in, 1))
68 }
69
70 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank_in: usize) -> usize
71 where
72 Module<B>: GGLWEEncryptSkFamily<B>,
73 {
74 GLWESecretExec::bytes_of(module, n, rank_in)
75 + (GLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank_in, 1) | GLWESecret::bytes_of(n, rank_in))
76 }
77}
78
79impl<D: DataMut> GLWEToLWESwitchingKey<D> {
80 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
81 &mut self,
82 module: &Module<B>,
83 sk_lwe: &LWESecret<DLwe>,
84 sk_glwe: &GLWESecret<DGlwe>,
85 source_xa: &mut Source,
86 source_xe: &mut Source,
87 sigma: f64,
88 scratch: &mut Scratch<B>,
89 ) where
90 DLwe: DataRef,
91 DGlwe: DataRef,
92 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
93 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
94 {
95 #[cfg(debug_assertions)]
96 {
97 assert!(sk_lwe.n() <= module.n());
98 }
99
100 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(sk_glwe.n(), 1);
101 sk_lwe_as_glwe.data.zero();
102 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
103 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0);
104
105 self.0.encrypt_sk(
106 module,
107 sk_glwe,
108 &sk_lwe_as_glwe,
109 source_xa,
110 source_xe,
111 sigma,
112 scratch1,
113 );
114 }
115}
116
117#[derive(PartialEq, Eq)]
118pub struct LWEToGLWESwitchingKey<D: Data>(pub(crate) GLWESwitchingKey<D>);
119
120impl<D: Data> Infos for LWEToGLWESwitchingKey<D> {
121 type Inner = MatZnx<D>;
122
123 fn inner(&self) -> &Self::Inner {
124 &self.0.inner()
125 }
126
127 fn basek(&self) -> usize {
128 self.0.basek()
129 }
130
131 fn k(&self) -> usize {
132 self.0.k()
133 }
134}
135
136impl<D: Data> LWEToGLWESwitchingKey<D> {
137 pub fn digits(&self) -> usize {
138 self.0.digits()
139 }
140
141 pub fn rank(&self) -> usize {
142 self.0.rank()
143 }
144
145 pub fn rank_in(&self) -> usize {
146 self.0.rank_in()
147 }
148
149 pub fn rank_out(&self) -> usize {
150 self.0.rank_out()
151 }
152}
153
154impl<D: DataMut> ReaderFrom for LWEToGLWESwitchingKey<D> {
155 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
156 self.0.read_from(reader)
157 }
158}
159
160impl<D: DataRef> WriterTo for LWEToGLWESwitchingKey<D> {
161 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
162 self.0.write_to(writer)
163 }
164}
165
166impl LWEToGLWESwitchingKey<Vec<u8>> {
167 pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, rank_out: usize) -> Self {
168 Self(GLWESwitchingKey::alloc(n, basek, k, rows, 1, 1, rank_out))
169 }
170
171 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank_out: usize) -> usize
172 where
173 Module<B>: GGLWEEncryptSkFamily<B>,
174 {
175 GLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, 1, rank_out) + GLWESecret::bytes_of(n, 1)
176 }
177}
178
179impl<D: DataMut> LWEToGLWESwitchingKey<D> {
180 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
181 &mut self,
182 module: &Module<B>,
183 sk_lwe: &LWESecret<DLwe>,
184 sk_glwe: &GLWESecret<DGlwe>,
185 source_xa: &mut Source,
186 source_xe: &mut Source,
187 sigma: f64,
188 scratch: &mut Scratch<B>,
189 ) where
190 DLwe: DataRef,
191 DGlwe: DataRef,
192 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
193 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
194 {
195 #[cfg(debug_assertions)]
196 {
197 assert!(sk_lwe.n() <= module.n());
198 }
199
200 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(sk_glwe.n(), 1);
201 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
202 sk_lwe_as_glwe.data.at_mut(0, 0)[sk_lwe.n()..].fill(0);
203 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0);
204
205 self.0.encrypt_sk(
206 module,
207 &sk_lwe_as_glwe,
208 &sk_glwe,
209 source_xa,
210 source_xe,
211 sigma,
212 scratch1,
213 );
214 }
215}
216
217#[derive(PartialEq, Eq)]
218pub struct LWESwitchingKey<D: Data>(pub(crate) GLWESwitchingKey<D>);
219
220impl<D: Data> Infos for LWESwitchingKey<D> {
221 type Inner = MatZnx<D>;
222
223 fn inner(&self) -> &Self::Inner {
224 &self.0.inner()
225 }
226
227 fn basek(&self) -> usize {
228 self.0.basek()
229 }
230
231 fn k(&self) -> usize {
232 self.0.k()
233 }
234}
235
236impl<D: Data> LWESwitchingKey<D> {
237 pub fn digits(&self) -> usize {
238 self.0.digits()
239 }
240
241 pub fn rank(&self) -> usize {
242 self.0.rank()
243 }
244
245 pub fn rank_in(&self) -> usize {
246 self.0.rank_in()
247 }
248
249 pub fn rank_out(&self) -> usize {
250 self.0.rank_out()
251 }
252}
253
254impl<D: DataMut> ReaderFrom for LWESwitchingKey<D> {
255 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
256 self.0.read_from(reader)
257 }
258}
259
260impl<D: DataRef> WriterTo for LWESwitchingKey<D> {
261 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
262 self.0.write_to(writer)
263 }
264}
265
266impl LWESwitchingKey<Vec<u8>> {
267 pub fn alloc(n: usize, basek: usize, k: usize, rows: usize) -> Self {
268 Self(GLWESwitchingKey::alloc(n, basek, k, rows, 1, 1, 1))
269 }
270
271 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
272 where
273 Module<B>: GGLWEEncryptSkFamily<B>,
274 {
275 GLWESecret::bytes_of(n, 1)
276 + GLWESecretExec::bytes_of(module, n, 1)
277 + GLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, 1, 1)
278 }
279}
280
281impl<D: DataMut> LWESwitchingKey<D> {
282 pub fn encrypt_sk<DIn, DOut, B: Backend>(
283 &mut self,
284 module: &Module<B>,
285 sk_lwe_in: &LWESecret<DIn>,
286 sk_lwe_out: &LWESecret<DOut>,
287 source_xa: &mut Source,
288 source_xe: &mut Source,
289 sigma: f64,
290 scratch: &mut Scratch<B>,
291 ) where
292 DIn: DataRef,
293 DOut: DataRef,
294 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
295 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
296 {
297 #[cfg(debug_assertions)]
298 {
299 assert!(sk_lwe_in.n() <= self.n());
300 assert!(sk_lwe_out.n() <= self.n());
301 assert!(self.n() <= module.n());
302 }
303
304 let (mut sk_in_glwe, scratch1) = scratch.take_glwe_secret(self.n(), 1);
305 let (mut sk_out_glwe, scratch2) = scratch1.take_glwe_secret(self.n(), 1);
306
307 sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
308 sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
309 module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0);
310
311 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
312 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
313 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0);
314
315 self.0.encrypt_sk(
316 module,
317 &sk_in_glwe,
318 &sk_out_glwe,
319 source_xa,
320 source_xe,
321 sigma,
322 scratch2,
323 );
324 }
325}
326
327impl LWECiphertext<Vec<u8>> {
328 pub fn from_glwe_scratch_space<B: Backend>(
329 module: &Module<B>,
330 n: usize,
331 basek: usize,
332 k_lwe: usize,
333 k_glwe: usize,
334 k_ksk: usize,
335 rank: usize,
336 ) -> usize
337 where
338 Module<B>: GLWEKeyswitchFamily<B>,
339 {
340 GLWECiphertext::bytes_of(n, basek, k_lwe, 1)
341 + GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
342 }
343
344 pub fn keyswitch_scratch_space<B: Backend>(
345 module: &Module<B>,
346 n: usize,
347 basek: usize,
348 k_lwe_out: usize,
349 k_lwe_in: usize,
350 k_ksk: usize,
351 ) -> usize
352 where
353 Module<B>: GLWEKeyswitchFamily<B>,
354 {
355 GLWECiphertext::bytes_of(n, basek, k_lwe_out.max(k_lwe_in), 1)
356 + GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_lwe_out, k_ksk, 1, 1)
357 }
358}