poulpy/core/src/gglwe/layouts_compressed.rs
core/src/gglwe/layouts_compressed.rsRaw
1use backend::hal::{
2 api::{MatZnxAlloc, MatZnxAllocBytes, VecZnxCopy, VecZnxFillUniform},
3 layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
4};
5
6use crate::{AutomorphismKey, GGLWECiphertext, GLWECiphertextCompressed, GLWESwitchingKey, GLWETensorKey, Infos};
7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
8
9#[derive(PartialEq, Eq)]
10pub struct GGLWECiphertextCompressed<D: Data> {
11 pub(crate) data: MatZnx<D>,
12 pub(crate) basek: usize,
13 pub(crate) k: usize,
14 pub(crate) rank_out: usize,
15 pub(crate) digits: usize,
16 pub(crate) seed: Vec<[u8; 32]>,
17}
18
19impl GGLWECiphertextCompressed<Vec<u8>> {
20 pub fn alloc<B: Backend>(
21 module: &Module<B>,
22 basek: usize,
23 k: usize,
24 rows: usize,
25 digits: usize,
26 rank_in: usize,
27 rank_out: usize,
28 ) -> Self
29 where
30 Module<B>: MatZnxAlloc,
31 {
32 let size: usize = k.div_ceil(basek);
33 debug_assert!(
34 size > digits,
35 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
36 size,
37 digits
38 );
39
40 assert!(
41 rows * digits <= size,
42 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
43 rows,
44 digits,
45 size
46 );
47
48 Self {
49 data: module.mat_znx_alloc(rows, rank_in, 1, size),
50 basek: basek,
51 k,
52 rank_out,
53 digits,
54 seed: vec![[0u8; 32]; rows * rank_in],
55 }
56 }
57
58 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize
59 where
60 Module<B>: MatZnxAllocBytes,
61 {
62 let size: usize = k.div_ceil(basek);
63 debug_assert!(
64 size > digits,
65 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
66 size,
67 digits
68 );
69
70 assert!(
71 rows * digits <= size,
72 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
73 rows,
74 digits,
75 size
76 );
77
78 module.mat_znx_alloc_bytes(rows, rank_in, 1, rows)
79 }
80}
81
82impl<D: Data> Infos for GGLWECiphertextCompressed<D> {
83 type Inner = MatZnx<D>;
84
85 fn inner(&self) -> &Self::Inner {
86 &self.data
87 }
88
89 fn basek(&self) -> usize {
90 self.basek
91 }
92
93 fn k(&self) -> usize {
94 self.k
95 }
96}
97
98impl<D: Data> GGLWECiphertextCompressed<D> {
99 pub fn rank(&self) -> usize {
100 self.rank_out
101 }
102
103 pub fn digits(&self) -> usize {
104 self.digits
105 }
106
107 pub fn rank_in(&self) -> usize {
108 self.data.cols_in()
109 }
110
111 pub fn rank_out(&self) -> usize {
112 self.rank_out
113 }
114}
115
116impl<D: DataRef> GGLWECiphertextCompressed<D> {
117 pub(crate) fn at(&self, row: usize, col: usize) -> GLWECiphertextCompressed<&[u8]> {
118 GLWECiphertextCompressed {
119 data: self.data.at(row, col),
120 basek: self.basek,
121 k: self.k,
122 rank: self.rank_out,
123 seed: self.seed[self.rank_in() * row + col],
124 }
125 }
126}
127
128impl<D: DataMut> GGLWECiphertextCompressed<D> {
129 pub(crate) fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertextCompressed<&mut [u8]> {
130 let rank_in: usize = self.rank_in();
131 GLWECiphertextCompressed {
132 data: self.data.at_mut(row, col),
133 basek: self.basek,
134 k: self.k,
135 rank: self.rank_out,
136 seed: self.seed[rank_in * row + col], // Warning: value is copied and not borrow mut
137 }
138 }
139}
140
141impl<D: DataMut> ReaderFrom for GGLWECiphertextCompressed<D> {
142 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
143 self.k = reader.read_u64::<LittleEndian>()? as usize;
144 self.basek = reader.read_u64::<LittleEndian>()? as usize;
145 self.digits = reader.read_u64::<LittleEndian>()? as usize;
146 self.rank_out = reader.read_u64::<LittleEndian>()? as usize;
147 let seed_len = reader.read_u64::<LittleEndian>()? as usize;
148 if seed_len != self.seed.len() {
149 } else {
150 for s in &mut self.seed {
151 reader.read_exact(s)?;
152 }
153 }
154 self.data.read_from(reader)
155 }
156}
157
158impl<D: DataRef> WriterTo for GGLWECiphertextCompressed<D> {
159 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
160 writer.write_u64::<LittleEndian>(self.k as u64)?;
161 writer.write_u64::<LittleEndian>(self.basek as u64)?;
162 writer.write_u64::<LittleEndian>(self.digits as u64)?;
163 writer.write_u64::<LittleEndian>(self.rank_out as u64)?;
164 writer.write_u64::<LittleEndian>(self.seed.len() as u64)?;
165 for s in &self.seed {
166 writer.write_all(s)?;
167 }
168 self.data.write_to(writer)
169 }
170}
171
172impl<D: DataMut> GGLWECiphertext<D> {
173 pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &GGLWECiphertextCompressed<DataOther>)
174 where
175 Module<B>: VecZnxFillUniform + VecZnxCopy,
176 {
177 #[cfg(debug_assertions)]
178 {
179 use backend::hal::api::ZnxInfos;
180
181 assert_eq!(
182 self.n(),
183 other.data.n(),
184 "invalid receiver: self.n()={} != other.n()={}",
185 self.n(),
186 other.data.n()
187 );
188 assert_eq!(
189 self.size(),
190 other.size(),
191 "invalid receiver: self.size()={} != other.size()={}",
192 self.size(),
193 other.size()
194 );
195 assert_eq!(
196 self.rank_in(),
197 other.rank_in(),
198 "invalid receiver: self.rank_in()={} != other.rank_in()={}",
199 self.rank_in(),
200 other.rank_in()
201 );
202 assert_eq!(
203 self.rank_out(),
204 other.rank_out(),
205 "invalid receiver: self.rank_out()={} != other.rank_out()={}",
206 self.rank_out(),
207 other.rank_out()
208 );
209
210 assert_eq!(
211 self.rows(),
212 other.rows(),
213 "invalid receiver: self.rows()={} != other.rows()={}",
214 self.rows(),
215 other.rows()
216 );
217 }
218
219 let rank_in: usize = self.rank_in();
220 let rows: usize = self.rows();
221
222 (0..rank_in).for_each(|col_i| {
223 (0..rows).for_each(|row_i| {
224 self.at_mut(row_i, col_i)
225 .decompress(module, &other.at(row_i, col_i));
226 });
227 });
228 }
229}
230
231#[derive(PartialEq, Eq)]
232pub struct GLWESwitchingKeyCompressed<D: Data> {
233 pub(crate) key: GGLWECiphertextCompressed<D>,
234 pub(crate) sk_in_n: usize, // Degree of sk_in
235 pub(crate) sk_out_n: usize, // Degree of sk_out
236}
237
238impl<D: Data> Infos for GLWESwitchingKeyCompressed<D> {
239 type Inner = MatZnx<D>;
240
241 fn inner(&self) -> &Self::Inner {
242 &self.key.inner()
243 }
244
245 fn basek(&self) -> usize {
246 self.key.basek()
247 }
248
249 fn k(&self) -> usize {
250 self.key.k()
251 }
252}
253
254impl<D: Data> GLWESwitchingKeyCompressed<D> {
255 pub fn rank(&self) -> usize {
256 self.key.rank()
257 }
258
259 pub fn digits(&self) -> usize {
260 self.key.digits()
261 }
262
263 pub fn rank_in(&self) -> usize {
264 self.key.rank_in()
265 }
266
267 pub fn rank_out(&self) -> usize {
268 self.key.rank_out()
269 }
270}
271
272impl GLWESwitchingKeyCompressed<Vec<u8>> {
273 pub fn alloc<B: Backend>(
274 module: &Module<B>,
275 basek: usize,
276 k: usize,
277 rows: usize,
278 digits: usize,
279 rank_in: usize,
280 rank_out: usize,
281 ) -> Self
282 where
283 Module<B>: MatZnxAlloc,
284 {
285 GLWESwitchingKeyCompressed {
286 key: GGLWECiphertextCompressed::alloc(module, basek, k, rows, digits, rank_in, rank_out),
287 sk_in_n: 0,
288 sk_out_n: 0,
289 }
290 }
291
292 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize
293 where
294 Module<B>: MatZnxAllocBytes,
295 {
296 GGLWECiphertextCompressed::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, rank_in)
297 }
298}
299
300impl<D: DataMut> ReaderFrom for GLWESwitchingKeyCompressed<D> {
301 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
302 self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
303 self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
304 self.key.read_from(reader)
305 }
306}
307
308impl<D: DataRef> WriterTo for GLWESwitchingKeyCompressed<D> {
309 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
310 writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
311 writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
312 self.key.write_to(writer)
313 }
314}
315
316impl<D: DataMut> GLWESwitchingKey<D> {
317 pub fn decompress<DataOther: DataRef, B: Backend>(
318 &mut self,
319 module: &Module<B>,
320 other: &GLWESwitchingKeyCompressed<DataOther>,
321 ) where
322 Module<B>: VecZnxFillUniform + VecZnxCopy,
323 {
324 self.key.decompress(module, &other.key);
325 self.sk_in_n = other.sk_in_n;
326 self.sk_out_n = other.sk_out_n;
327 }
328}
329
330#[derive(PartialEq, Eq)]
331pub struct AutomorphismKeyCompressed<D: Data> {
332 pub(crate) key: GLWESwitchingKeyCompressed<D>,
333 pub(crate) p: i64,
334}
335
336impl AutomorphismKeyCompressed<Vec<u8>> {
337 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
338 where
339 Module<B>: MatZnxAlloc,
340 {
341 AutomorphismKeyCompressed {
342 key: GLWESwitchingKeyCompressed::alloc(module, basek, k, rows, digits, rank, rank),
343 p: 0,
344 }
345 }
346
347 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
348 where
349 Module<B>: MatZnxAllocBytes,
350 {
351 GLWESwitchingKeyCompressed::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, rank)
352 }
353}
354
355impl<D: Data> Infos for AutomorphismKeyCompressed<D> {
356 type Inner = MatZnx<D>;
357
358 fn inner(&self) -> &Self::Inner {
359 &self.key.inner()
360 }
361
362 fn basek(&self) -> usize {
363 self.key.basek()
364 }
365
366 fn k(&self) -> usize {
367 self.key.k()
368 }
369}
370
371impl<D: Data> AutomorphismKeyCompressed<D> {
372 pub fn rank(&self) -> usize {
373 self.key.rank()
374 }
375
376 pub fn digits(&self) -> usize {
377 self.key.digits()
378 }
379
380 pub fn rank_in(&self) -> usize {
381 self.key.rank_in()
382 }
383
384 pub fn rank_out(&self) -> usize {
385 self.key.rank_out()
386 }
387}
388
389impl<D: DataMut> ReaderFrom for AutomorphismKeyCompressed<D> {
390 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
391 self.p = reader.read_u64::<LittleEndian>()? as i64;
392 self.key.read_from(reader)
393 }
394}
395
396impl<D: DataRef> WriterTo for AutomorphismKeyCompressed<D> {
397 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
398 writer.write_u64::<LittleEndian>(self.p as u64)?;
399 self.key.write_to(writer)
400 }
401}
402
403impl<D: DataMut> AutomorphismKey<D> {
404 pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &AutomorphismKeyCompressed<DataOther>)
405 where
406 Module<B>: VecZnxFillUniform + VecZnxCopy,
407 {
408 self.key.decompress(module, &other.key);
409 self.p = other.p;
410 }
411}
412
413#[derive(PartialEq, Eq)]
414pub struct GLWETensorKeyCompressed<D: Data> {
415 pub(crate) keys: Vec<GLWESwitchingKeyCompressed<D>>,
416}
417
418impl GLWETensorKeyCompressed<Vec<u8>> {
419 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
420 where
421 Module<B>: MatZnxAlloc,
422 {
423 let mut keys: Vec<GLWESwitchingKeyCompressed<Vec<u8>>> = Vec::new();
424 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
425 (0..pairs).for_each(|_| {
426 keys.push(GLWESwitchingKeyCompressed::alloc(
427 module, basek, k, rows, digits, 1, rank,
428 ));
429 });
430 Self { keys: keys }
431 }
432
433 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
434 where
435 Module<B>: MatZnxAllocBytes,
436 {
437 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
438 pairs * GLWESwitchingKeyCompressed::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, 1)
439 }
440}
441
442impl<D: Data> Infos for GLWETensorKeyCompressed<D> {
443 type Inner = MatZnx<D>;
444
445 fn inner(&self) -> &Self::Inner {
446 &self.keys[0].inner()
447 }
448
449 fn basek(&self) -> usize {
450 self.keys[0].basek()
451 }
452
453 fn k(&self) -> usize {
454 self.keys[0].k()
455 }
456}
457
458impl<D: Data> GLWETensorKeyCompressed<D> {
459 pub fn rank(&self) -> usize {
460 self.keys[0].rank()
461 }
462
463 pub fn digits(&self) -> usize {
464 self.keys[0].digits()
465 }
466
467 pub fn rank_in(&self) -> usize {
468 self.keys[0].rank_in()
469 }
470
471 pub fn rank_out(&self) -> usize {
472 self.keys[0].rank_out()
473 }
474}
475
476impl<D: DataMut> ReaderFrom for GLWETensorKeyCompressed<D> {
477 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
478 let len: usize = reader.read_u64::<LittleEndian>()? as usize;
479 if self.keys.len() != len {
480 return Err(std::io::Error::new(
481 std::io::ErrorKind::InvalidData,
482 format!("self.keys.len()={} != read len={}", self.keys.len(), len),
483 ));
484 }
485 for key in &mut self.keys {
486 key.read_from(reader)?;
487 }
488 Ok(())
489 }
490}
491
492impl<D: DataRef> WriterTo for GLWETensorKeyCompressed<D> {
493 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
494 writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
495 for key in &self.keys {
496 key.write_to(writer)?;
497 }
498 Ok(())
499 }
500}
501
502impl<D: DataMut> GLWETensorKeyCompressed<D> {
503 pub(crate) fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKeyCompressed<D> {
504 if i > j {
505 std::mem::swap(&mut i, &mut j);
506 };
507 let rank: usize = self.rank();
508 &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
509 }
510}
511
512impl<D: DataMut> GLWETensorKey<D> {
513 pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &GLWETensorKeyCompressed<DataOther>)
514 where
515 Module<B>: VecZnxFillUniform + VecZnxCopy,
516 {
517 #[cfg(debug_assertions)]
518 {
519 assert_eq!(
520 self.keys.len(),
521 other.keys.len(),
522 "invalid receiver: self.keys.len()={} != other.keys.len()={}",
523 self.keys.len(),
524 other.keys.len()
525 );
526 }
527
528 self.keys
529 .iter_mut()
530 .zip(other.keys.iter())
531 .for_each(|(a, b)| {
532 a.decompress(module, b);
533 });
534 }
535}