1use poulpy_hal::{
2 layouts::{
3 Backend, Data, FillUniform, HostDataMut, HostDataRef, MatZnx, MatZnxAtBackendMut, MatZnxAtBackendRef, MatZnxToBackendMut,
4 MatZnxToBackendRef, Module, ReaderFrom, TransferFrom, WriterTo,
5 },
6 source::Source,
7};
8
9use crate::api::ModuleTransfer;
10use crate::layouts::{
11 Base2K, Degree, Dnum, Dsize, GGSWAtViewRef, GLWE, GLWEInfos, GLWEViewMut, GLWEViewRef, LWEInfos, Rank, TorusPrecision,
12};
13use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
14
15use std::{
16 fmt,
17 ops::{Deref, DerefMut},
18};
19
20pub trait GGLWEInfos
21where
22 Self: GLWEInfos,
23{
24 fn dnum(&self) -> Dnum;
25 fn dsize(&self) -> Dsize;
26 fn rank_in(&self) -> Rank;
27 fn rank_out(&self) -> Rank;
28 fn gglwe_layout(&self) -> GGLWELayout {
29 GGLWELayout {
30 n: self.n(),
31 base2k: self.base2k(),
32 k: self.max_k(),
33 rank_in: self.rank_in(),
34 rank_out: self.rank_out(),
35 dsize: self.dsize(),
36 dnum: self.dnum(),
37 }
38 }
39}
40
41pub trait SetGGLWEInfos {
42 fn set_dsize(&mut self, dsize: usize);
43}
44
45#[derive(PartialEq, Eq, Copy, Clone, Debug)]
46pub struct GGLWELayout {
47 pub n: Degree,
48 pub base2k: Base2K,
49 pub k: TorusPrecision,
50 pub rank_in: Rank,
51 pub rank_out: Rank,
52 pub dnum: Dnum,
53 pub dsize: Dsize,
54}
55
56impl LWEInfos for GGLWELayout {
57 fn base2k(&self) -> Base2K {
58 self.base2k
59 }
60
61 fn n(&self) -> Degree {
62 self.n
63 }
64
65 fn size(&self) -> usize {
66 self.k.as_usize().div_ceil(self.base2k.as_usize())
67 }
68}
69
70impl GLWEInfos for GGLWELayout {
71 fn rank(&self) -> Rank {
72 self.rank_out
73 }
74}
75
76impl GGLWEInfos for GGLWELayout {
77 fn rank_in(&self) -> Rank {
78 self.rank_in
79 }
80
81 fn dsize(&self) -> Dsize {
82 self.dsize
83 }
84
85 fn rank_out(&self) -> Rank {
86 self.rank_out
87 }
88
89 fn dnum(&self) -> Dnum {
90 self.dnum
91 }
92}
93
94#[derive(PartialEq, Eq, Clone)]
95pub struct GGLWE<D: Data> {
96 pub(crate) data: MatZnx<D>,
97 pub(crate) base2k: Base2K,
98 pub(crate) dsize: Dsize,
99}
100
101pub struct GGLWEBackendRef<'a, BE: Backend + 'a> {
102 inner: GGLWE<BE::BufRef<'a>>,
103}
104
105impl<'a, BE: Backend + 'a> GGLWEBackendRef<'a, BE> {
106 pub fn from_inner(inner: GGLWE<BE::BufRef<'a>>) -> Self {
107 Self { inner }
108 }
109
110 pub fn into_inner(self) -> GGLWE<BE::BufRef<'a>> {
111 self.inner
112 }
113
114 pub fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
115 GLWEViewRef::from_inner(gglwe_at_backend_ref_from_ref::<BE>(&self.inner, row, col))
116 }
117}
118
119impl<'a, BE: Backend + 'a> Deref for GGLWEBackendRef<'a, BE> {
120 type Target = GGLWE<BE::BufRef<'a>>;
121
122 fn deref(&self) -> &Self::Target {
123 &self.inner
124 }
125}
126
127pub struct GGLWEBackendMut<'a, BE: Backend + 'a> {
128 inner: GGLWE<BE::BufMut<'a>>,
129}
130
131impl<'a, BE: Backend + 'a> GGLWEBackendMut<'a, BE> {
132 pub fn from_inner(inner: GGLWE<BE::BufMut<'a>>) -> Self {
133 Self { inner }
134 }
135
136 pub fn into_inner(self) -> GGLWE<BE::BufMut<'a>> {
137 self.inner
138 }
139
140 pub fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
141 GLWEViewRef::from_inner(gglwe_at_backend_ref_from_mut::<BE>(&self.inner, row, col))
142 }
143
144 pub fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
145 GLWEViewMut::from_inner(gglwe_at_backend_mut_from_mut::<BE>(&mut self.inner, row, col))
146 }
147}
148
149impl<'a, BE: Backend + 'a> Deref for GGLWEBackendMut<'a, BE> {
150 type Target = GGLWE<BE::BufMut<'a>>;
151
152 fn deref(&self) -> &Self::Target {
153 &self.inner
154 }
155}
156
157impl<'a, BE: Backend + 'a> DerefMut for GGLWEBackendMut<'a, BE> {
158 fn deref_mut(&mut self) -> &mut Self::Target {
159 &mut self.inner
160 }
161}
162
163impl_gglwe_infos_for_inner!(GGLWEBackendRef<'a, BE>, ['a, BE: Backend + 'a]; inner);
164impl_gglwe_infos_for_inner!(GGLWEBackendMut<'a, BE>, ['a, BE: Backend + 'a]; inner);
165
166impl<'a, BE: Backend + 'a> GGLWEToBackendRef<BE> for GGLWEBackendRef<'a, BE> {
167 fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
168 GGLWEBackendRef::from_inner(GGLWE {
169 base2k: self.inner.base2k,
170 dsize: self.inner.dsize,
171 data: poulpy_hal::layouts::mat_znx_backend_ref_from_ref::<BE>(&self.inner.data),
172 })
173 }
174}
175
176impl<'a, BE: Backend + 'a> GGSWAtViewRef<BE> for GGLWEBackendRef<'a, BE> {
177 fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
178 GGLWEBackendRef::at_view(self, row, col)
179 }
180}
181
182impl<'a, BE: Backend + 'a> GGLWEToBackendRef<BE> for GGLWEBackendMut<'a, BE> {
183 fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
184 GGLWEBackendRef::from_inner(GGLWE {
185 base2k: self.inner.base2k,
186 dsize: self.inner.dsize,
187 data: poulpy_hal::layouts::mat_znx_backend_ref_from_mut::<BE>(&self.inner.data),
188 })
189 }
190}
191
192impl<'a, BE: Backend + 'a> GGLWEToBackendMut<BE> for GGLWEBackendMut<'a, BE> {
193 fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
194 GGLWEBackendMut::from_inner(GGLWE {
195 base2k: self.inner.base2k,
196 dsize: self.inner.dsize,
197 data: poulpy_hal::layouts::mat_znx_backend_mut_from_mut::<BE>(&mut self.inner.data),
198 })
199 }
200}
201
202impl<D: Data> LWEInfos for GGLWE<D> {
203 fn base2k(&self) -> Base2K {
204 self.base2k
205 }
206
207 fn n(&self) -> Degree {
208 Degree(self.data.n() as u32)
209 }
210
211 fn size(&self) -> usize {
212 self.data.size()
213 }
214}
215
216impl<D: Data> LWEInfos for &GGLWE<D> {
217 fn base2k(&self) -> Base2K {
218 self.base2k
219 }
220
221 fn n(&self) -> Degree {
222 Degree(self.data.n() as u32)
223 }
224
225 fn size(&self) -> usize {
226 self.data.size()
227 }
228}
229
230impl<D: Data> GLWEInfos for GGLWE<D> {
231 fn rank(&self) -> Rank {
232 self.rank_out()
233 }
234}
235
236impl<D: Data> GLWEInfos for &GGLWE<D> {
237 fn rank(&self) -> Rank {
238 self.rank_out()
239 }
240}
241
242impl<D: Data> GGLWEInfos for GGLWE<D> {
243 fn rank_in(&self) -> Rank {
244 Rank(self.data.cols_in() as u32)
245 }
246
247 fn rank_out(&self) -> Rank {
248 Rank(self.data.cols_out() as u32 - 1)
249 }
250
251 fn dsize(&self) -> Dsize {
252 self.dsize
253 }
254
255 fn dnum(&self) -> Dnum {
256 Dnum(self.data.rows() as u32)
257 }
258}
259
260impl<D: Data> GGLWEInfos for &GGLWE<D> {
261 fn rank_in(&self) -> Rank {
262 Rank(self.data.cols_in() as u32)
263 }
264
265 fn rank_out(&self) -> Rank {
266 Rank(self.data.cols_out() as u32 - 1)
267 }
268
269 fn dsize(&self) -> Dsize {
270 self.dsize
271 }
272
273 fn dnum(&self) -> Dnum {
274 Dnum(self.data.rows() as u32)
275 }
276}
277
278impl<D: Data> LWEInfos for &mut GGLWE<D> {
279 fn base2k(&self) -> Base2K {
280 self.base2k
281 }
282
283 fn n(&self) -> Degree {
284 Degree(self.data.n() as u32)
285 }
286
287 fn size(&self) -> usize {
288 self.data.size()
289 }
290}
291
292impl<D: Data> GLWEInfos for &mut GGLWE<D> {
293 fn rank(&self) -> Rank {
294 self.rank_out()
295 }
296}
297
298impl<D: Data> GGLWEInfos for &mut GGLWE<D> {
299 fn rank_in(&self) -> Rank {
300 Rank(self.data.cols_in() as u32)
301 }
302
303 fn rank_out(&self) -> Rank {
304 Rank(self.data.cols_out() as u32 - 1)
305 }
306
307 fn dsize(&self) -> Dsize {
308 self.dsize
309 }
310
311 fn dnum(&self) -> Dnum {
312 Dnum(self.data.rows() as u32)
313 }
314}
315
316impl<D: HostDataRef> GGLWE<D> {
317 pub fn data(&self) -> &MatZnx<D> {
318 &self.data
319 }
320}
321
322pub(crate) trait GGLWEAtBackendRef<BE: Backend> {
323 fn at_backend(&self, row: usize, col: usize) -> GLWE<BE::BufRef<'_>>;
324}
325
326impl<BE: Backend> GGLWEAtBackendRef<BE> for GGLWE<BE::OwnedBuf> {
327 fn at_backend(&self, row: usize, col: usize) -> GLWE<BE::BufRef<'_>> {
328 let data = <MatZnx<BE::OwnedBuf> as MatZnxAtBackendRef<BE>>::at_backend(&self.data, row, col);
329 GLWE {
330 base2k: self.base2k,
331 data,
332 }
333 }
334}
335
336pub(crate) fn gglwe_at_backend_ref_from_ref<'a, 'b, BE: Backend>(
337 gglwe: &'a GGLWE<BE::BufRef<'b>>,
338 row: usize,
339 col: usize,
340) -> GLWE<BE::BufRef<'a>> {
341 let data = poulpy_hal::layouts::mat_znx_at_backend_ref_from_ref::<BE>(&gglwe.data, row, col);
342 GLWE {
343 base2k: gglwe.base2k,
344 data,
345 }
346}
347
348pub trait GGLWEAtViewRef<BE: Backend> {
349 fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE>;
350}
351
352impl<BE: Backend> GGLWEAtViewRef<BE> for GGLWE<BE::OwnedBuf> {
353 fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
354 GLWEViewRef::from_inner(<GGLWE<BE::OwnedBuf> as GGLWEAtBackendRef<BE>>::at_backend(self, row, col))
355 }
356}
357
358impl<'b, BE: Backend + 'b> GGLWEAtViewRef<BE> for &GGLWE<BE::BufRef<'b>> {
359 fn at_view(&self, row: usize, col: usize) -> GLWEViewRef<'_, BE> {
360 GLWEViewRef::from_inner(gglwe_at_backend_ref_from_ref::<BE>(self, row, col))
361 }
362}
363
364pub(crate) fn gglwe_at_backend_ref_from_mut<'a, 'b, BE: Backend>(
365 gglwe: &'a GGLWE<BE::BufMut<'b>>,
366 row: usize,
367 col: usize,
368) -> GLWE<BE::BufRef<'a>> {
369 let data = poulpy_hal::layouts::mat_znx_at_backend_ref_from_mut::<BE>(&gglwe.data, row, col);
370 GLWE {
371 base2k: gglwe.base2k,
372 data,
373 }
374}
375
376impl<D: HostDataMut> GGLWE<D> {
377 pub fn data_mut(&mut self) -> &mut MatZnx<D> {
378 &mut self.data
379 }
380}
381
382pub(crate) trait GGLWEAtBackendMut<BE: Backend> {
383 fn at_backend_mut(&mut self, row: usize, col: usize) -> GLWE<BE::BufMut<'_>>;
384}
385
386impl<BE: Backend> GGLWEAtBackendMut<BE> for GGLWE<BE::OwnedBuf> {
387 fn at_backend_mut(&mut self, row: usize, col: usize) -> GLWE<BE::BufMut<'_>> {
388 let base2k = self.base2k;
389 let data = <MatZnx<BE::OwnedBuf> as MatZnxAtBackendMut<BE>>::at_backend_mut(&mut self.data, row, col);
390 GLWE { base2k, data }
391 }
392}
393
394pub(crate) fn gglwe_at_backend_mut_from_mut<'a, 'b, BE: Backend>(
395 gglwe: &'a mut GGLWE<BE::BufMut<'b>>,
396 row: usize,
397 col: usize,
398) -> GLWE<BE::BufMut<'a>> {
399 let base2k = gglwe.base2k;
400 let data = poulpy_hal::layouts::mat_znx_at_backend_mut_from_mut::<BE>(&mut gglwe.data, row, col);
401 GLWE { base2k, data }
402}
403
404pub trait GGLWEAtViewMut<BE: Backend> {
405 fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE>;
406}
407
408impl<BE: Backend> GGLWEAtViewMut<BE> for GGLWE<BE::OwnedBuf> {
409 fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
410 GLWEViewMut::from_inner(<GGLWE<BE::OwnedBuf> as GGLWEAtBackendMut<BE>>::at_backend_mut(self, row, col))
411 }
412}
413
414impl<'b, BE: Backend + 'b> GGLWEAtViewMut<BE> for &mut GGLWE<BE::BufMut<'b>> {
415 fn at_view_mut(&mut self, row: usize, col: usize) -> GLWEViewMut<'_, BE> {
416 GLWEViewMut::from_inner(gglwe_at_backend_mut_from_mut::<BE>(*self, row, col))
417 }
418}
419
420impl<D: HostDataRef> fmt::Debug for GGLWE<D> {
421 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
422 write!(f, "{self}")
423 }
424}
425
426impl<D: HostDataMut> FillUniform for GGLWE<D> {
427 fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
428 self.data.fill_uniform(log_bound, source);
429 }
430}
431
432impl<D: HostDataRef> fmt::Display for GGLWE<D> {
433 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
434 write!(
435 f,
436 "(GGLWE: k={} base2k={} dsize={}) {}",
437 self.max_k().0,
438 self.base2k().0,
439 self.dsize().0,
440 self.data
441 )
442 }
443}
444
445impl<D: HostDataRef> GGLWE<D> {
446 pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
447 let data = self.data.at(row, col);
448 GLWE {
449 base2k: self.base2k,
450 data,
451 }
452 }
453}
454
455impl<D: HostDataMut> GGLWE<D> {
456 pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
457 let base2k = self.base2k;
458 let data = self.data.at_mut(row, col);
459 GLWE { base2k, data }
460 }
461}
462
463impl<D: HostDataRef> GGLWE<D> {
464 pub fn to_backend<BE, To>(&self, dst: &Module<To>) -> GGLWE<To::OwnedBuf>
467 where
468 BE: Backend<OwnedBuf = D>,
469 To: Backend,
470 To: TransferFrom<BE>,
471 {
472 dst.upload_gglwe(self)
473 }
474}
475
476impl<D: Data> GGLWE<D> {
477 pub fn reinterpret<To>(self) -> GGLWE<To::OwnedBuf>
479 where
480 To: Backend<OwnedBuf = D>,
481 {
482 let (n, rows, cols_in, cols_out, size) = (
483 self.data.n(),
484 self.data.rows(),
485 self.data.cols_in(),
486 self.data.cols_out(),
487 self.data.size(),
488 );
489 GGLWE {
490 data: MatZnx::from_data(self.data.into_data(), n, rows, cols_in, cols_out, size),
491 base2k: self.base2k,
492 dsize: self.dsize,
493 }
494 }
495}
496
497#[expect(
498 dead_code,
499 reason = "host-owned constructors are kept for serialization and host-only staging"
500)]
501impl GGLWE<Vec<u8>> {
502 pub(crate) fn alloc_from_infos<A>(infos: &A) -> Self
503 where
504 A: GGLWEInfos,
505 {
506 Self::alloc(
507 infos.n(),
508 infos.base2k(),
509 infos.max_k(),
510 infos.rank_in(),
511 infos.rank_out(),
512 infos.dnum(),
513 infos.dsize(),
514 )
515 }
516
517 pub(crate) fn alloc(
518 n: Degree,
519 base2k: Base2K,
520 k: TorusPrecision,
521 rank_in: Rank,
522 rank_out: Rank,
523 dnum: Dnum,
524 dsize: Dsize,
525 ) -> Self {
526 let size: usize = k.0.div_ceil(base2k.0) as usize;
527 assert!(
528 size as u32 > dsize.0,
529 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
530 dsize.0
531 );
532
533 assert!(
534 dnum.0 * dsize.0 <= size as u32,
535 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
536 dnum.0,
537 dsize.0,
538 );
539
540 GGLWE {
541 data: MatZnx::from_data(
542 poulpy_hal::layouts::HostBytesBackend::alloc_bytes(MatZnx::<Vec<u8>>::bytes_of(
543 n.into(),
544 dnum.into(),
545 rank_in.into(),
546 (rank_out + 1).into(),
547 size,
548 )),
549 n.into(),
550 dnum.into(),
551 rank_in.into(),
552 (rank_out + 1).into(),
553 size,
554 ),
555 base2k,
556 dsize,
557 }
558 }
559
560 pub fn bytes_of_from_infos<A>(infos: &A) -> usize
561 where
562 A: GGLWEInfos,
563 {
564 Self::bytes_of(
565 infos.n(),
566 infos.base2k(),
567 infos.max_k(),
568 infos.rank_in(),
569 infos.rank_out(),
570 infos.dnum(),
571 infos.dsize(),
572 )
573 }
574
575 pub fn bytes_of(
576 n: Degree,
577 base2k: Base2K,
578 k: TorusPrecision,
579 rank_in: Rank,
580 rank_out: Rank,
581 dnum: Dnum,
582 dsize: Dsize,
583 ) -> usize {
584 let size: usize = k.0.div_ceil(base2k.0) as usize;
585 assert!(
586 size as u32 > dsize.0,
587 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
588 dsize.0
589 );
590
591 assert!(
592 dnum.0 * dsize.0 <= size as u32,
593 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
594 dnum.0,
595 dsize.0,
596 );
597
598 MatZnx::bytes_of(
599 n.into(),
600 dnum.into(),
601 rank_in.into(),
602 (rank_out + 1).into(),
603 k.0.div_ceil(base2k.0) as usize,
604 )
605 }
606}
607
608pub trait GGLWEToBackendMut<BE: Backend>: GGLWEToBackendRef<BE> {
609 fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE>;
610}
611
612impl<BE: Backend, D: Data> GGLWEToBackendMut<BE> for GGLWE<D>
613where
614 MatZnx<D>: MatZnxToBackendRef<BE> + MatZnxToBackendMut<BE>,
615{
616 fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
617 GGLWEBackendMut::from_inner(GGLWE {
618 base2k: self.base2k(),
619 dsize: self.dsize(),
620 data: self.data.to_backend_mut(),
621 })
622 }
623}
624
625impl<'b, BE: Backend + 'b> GGLWEToBackendRef<BE> for &mut GGLWE<BE::BufMut<'b>> {
626 fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
627 GGLWEBackendRef::from_inner(GGLWE {
628 base2k: self.base2k(),
629 dsize: self.dsize(),
630 data: poulpy_hal::layouts::mat_znx_backend_ref_from_mut::<BE>(&self.data),
631 })
632 }
633}
634
635impl<'b, BE: Backend + 'b> GGLWEToBackendMut<BE> for &mut GGLWE<BE::BufMut<'b>> {
636 fn to_backend_mut(&mut self) -> GGLWEBackendMut<'_, BE> {
637 GGLWEBackendMut::from_inner(GGLWE {
638 base2k: self.base2k(),
639 dsize: self.dsize(),
640 data: poulpy_hal::layouts::mat_znx_backend_mut_from_mut::<BE>(&mut self.data),
641 })
642 }
643}
644
645pub trait GGLWEToBackendRef<BE: Backend> {
646 fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE>;
647}
648
649impl<BE: Backend, D: Data> GGLWEToBackendRef<BE> for GGLWE<D>
650where
651 MatZnx<D>: MatZnxToBackendRef<BE>,
652{
653 fn to_backend_ref(&self) -> GGLWEBackendRef<'_, BE> {
654 GGLWEBackendRef::from_inner(GGLWE {
655 base2k: self.base2k(),
656 dsize: self.dsize(),
657 data: self.data.to_backend_ref(),
658 })
659 }
660}
661
662impl<D: HostDataMut> ReaderFrom for GGLWE<D> {
663 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
664 self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
665 self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
666 self.data.read_from(reader)
667 }
668}
669
670impl<D: HostDataRef> WriterTo for GGLWE<D> {
671 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
672 writer.write_u32::<LittleEndian>(self.base2k.0)?;
673 writer.write_u32::<LittleEndian>(self.dsize.0)?;
674 self.data.write_to(writer)
675 }
676}