1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::ModuleLogN,
5 layouts::{Backend, GaloisElement, Module, Scratch},
6};
7
8use crate::{
9 GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, GLWETrace, ScratchTakeCore,
10 layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GetGaloisElement},
11};
12pub trait GLWEPacking<BE: Backend> {
13 fn glwe_pack_galois_elements(&self) -> Vec<i64>;
14
15 fn glwe_pack_tmp_bytes<R, K>(&self, res: &R, key: &K) -> usize
16 where
17 R: GLWEInfos,
18 K: GGLWEInfos;
19
20 fn glwe_pack<R, A, K, H>(
23 &self,
24 res: &mut R,
25 a: HashMap<usize, &mut A>,
26 log_gap_out: usize,
27 keys: &H,
28 scratch: &mut Scratch<BE>,
29 ) where
30 R: GLWEToMut + GLWEInfos,
31 A: GLWEToMut + GLWEInfos,
32 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
33 H: GLWEAutomorphismKeyHelper<K, BE>;
34}
35
36impl<BE: Backend> GLWEPacking<BE> for Module<BE>
37where
38 Self: GLWEAutomorphism<BE>
39 + GaloisElement
40 + ModuleLogN
41 + GLWERotate<BE>
42 + GLWESub
43 + GLWEShift<BE>
44 + GLWEAdd
45 + GLWENormalize<BE>
46 + GLWECopy
47 + GLWETrace<BE>,
48 Scratch<BE>: ScratchTakeCore<BE>,
49{
50 fn glwe_pack_galois_elements(&self) -> Vec<i64> {
51 self.glwe_trace_galois_elements()
52 }
53
54 fn glwe_pack_tmp_bytes<R, K>(&self, res: &R, key: &K) -> usize
55 where
56 R: GLWEInfos,
57 K: GGLWEInfos,
58 {
59 self.glwe_rotate_tmp_bytes()
60 .max(self.glwe_rsh_tmp_byte())
61 .max(self.glwe_normalize_tmp_bytes())
62 .max(self.glwe_automorphism_tmp_bytes(res, res, key))
63 + GLWE::bytes_of_from_infos(res)
64 }
65
66 fn glwe_pack<R, A, K, H>(
69 &self,
70 res: &mut R,
71 mut a: HashMap<usize, &mut A>,
72 log_gap_out: usize,
73 keys: &H,
74 scratch: &mut Scratch<BE>,
75 ) where
76 R: GLWEToMut + GLWEInfos,
77 A: GLWEToMut + GLWEInfos,
78 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
79 H: GLWEAutomorphismKeyHelper<K, BE>,
80 {
81 assert!(*a.keys().max().unwrap() < self.n());
82
83 let log_n: usize = self.log_n();
84
85 for i in 0..(log_n - log_gap_out) {
86 let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
87
88 let key: &K = if i == 0 {
89 keys.get_automorphism_key(-1).unwrap()
90 } else {
91 keys.get_automorphism_key(self.galois_element(1 << (i - 1)))
92 .unwrap()
93 };
94
95 for j in 0..t {
96 let mut lo: Option<&mut A> = a.remove(&j);
97 let mut hi: Option<&mut A> = a.remove(&(j + t));
98
99 pack_internal(self, &mut lo, &mut hi, i, key, scratch);
100
101 if let Some(lo) = lo {
102 a.insert(j, lo);
103 } else if let Some(hi) = hi {
104 a.insert(j, hi);
105 }
106 }
107 }
108
109 self.glwe_trace(res, log_n - log_gap_out, *a.get(&0).unwrap(), keys, scratch);
110 }
111}
112
113#[allow(clippy::too_many_arguments)]
114fn pack_internal<M, A, B, K, BE: Backend>(
115 module: &M,
116 a: &mut Option<&mut A>,
117 b: &mut Option<&mut B>,
118 i: usize,
119 auto_key: &K,
120 scratch: &mut Scratch<BE>,
121) where
122 M: GLWEAutomorphism<BE> + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
123 A: GLWEToMut + GLWEInfos,
124 B: GLWEToMut + GLWEInfos,
125 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
126 Scratch<BE>: ScratchTakeCore<BE>,
127{
128 if let Some(a) = a.as_deref_mut() {
139 let t: i64 = 1 << (a.n().log2() - i - 1);
140
141 if let Some(b) = b.as_deref_mut() {
142 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
143
144 module.glwe_rotate_inplace(-t, a, scratch_1);
146
147 module.glwe_sub(&mut tmp_b, a, b);
149 module.glwe_rsh(1, &mut tmp_b, scratch_1);
150
151 module.glwe_add_inplace(a, b);
153 module.glwe_rsh(1, a, scratch_1);
154
155 module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
156
157 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
159
160 module.glwe_sub_inplace(a, &tmp_b);
162 module.glwe_normalize_inplace(a, scratch_1);
163
164 module.glwe_rotate_inplace(t, a, scratch_1);
168 } else {
169 module.glwe_rsh(1, a, scratch);
170 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
172 }
173 } else if let Some(b) = b.as_deref_mut() {
174 let t: i64 = 1 << (b.n().log2() - i - 1);
175
176 let (mut tmp_b, scratch_1) = scratch.take_glwe(b);
177 module.glwe_rotate(t, &mut tmp_b, b);
178 module.glwe_rsh(1, &mut tmp_b, scratch_1);
179
180 module.glwe_automorphism_sub_negate(b, &tmp_b, auto_key, scratch_1);
182 }
183}