poulpy/core/src/glwe/ops.rs
core/src/glwe/ops.rsRaw
1use backend::{FFT64, Module, Scratch, VecZnx, VecZnxOps, ZnxZero};
2
3use crate::{GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, Infos, SetMetaData};
4
5pub trait GLWEOps: GLWECiphertextToMut + SetMetaData + Sized {
6 fn add<A, B>(&mut self, module: &Module<FFT64>, a: &A, b: &B)
7 where
8 A: GLWECiphertextToRef,
9 B: GLWECiphertextToRef,
10 {
11 #[cfg(debug_assertions)]
12 {
13 assert_eq!(a.n(), module.n());
14 assert_eq!(b.n(), module.n());
15 assert_eq!(self.n(), module.n());
16 assert_eq!(a.basek(), b.basek());
17 assert!(self.rank() >= a.rank().max(b.rank()));
18 }
19
20 let min_col: usize = a.rank().min(b.rank()) + 1;
21 let max_col: usize = a.rank().max(b.rank() + 1);
22 let self_col: usize = self.rank() + 1;
23
24 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
25 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
26 let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
27
28 (0..min_col).for_each(|i| {
29 module.vec_znx_add(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
30 });
31
32 if a.rank() > b.rank() {
33 (min_col..max_col).for_each(|i| {
34 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
35 });
36 } else {
37 (min_col..max_col).for_each(|i| {
38 module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
39 });
40 }
41
42 let size: usize = self_mut.size();
43 (max_col..self_col).for_each(|i| {
44 (0..size).for_each(|j| {
45 self_mut.data.zero_at(i, j);
46 });
47 });
48
49 self.set_basek(a.basek());
50 self.set_k(set_k(self, a, b));
51 }
52
53 fn add_inplace<A>(&mut self, module: &Module<FFT64>, a: &A)
54 where
55 A: GLWECiphertextToRef + Infos,
56 {
57 #[cfg(debug_assertions)]
58 {
59 assert_eq!(a.n(), module.n());
60 assert_eq!(self.n(), module.n());
61 assert_eq!(self.basek(), a.basek());
62 assert!(self.rank() >= a.rank())
63 }
64
65 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
66 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
67
68 (0..a.rank() + 1).for_each(|i| {
69 module.vec_znx_add_inplace(&mut self_mut.data, i, &a_ref.data, i);
70 });
71
72 if a.rank() != 0 {
73 self.set_k(a.k().min(self.k()));
74 }
75 }
76
77 fn sub<A, B>(&mut self, module: &Module<FFT64>, a: &A, b: &B)
78 where
79 A: GLWECiphertextToRef,
80 B: GLWECiphertextToRef,
81 {
82 #[cfg(debug_assertions)]
83 {
84 assert_eq!(a.n(), module.n());
85 assert_eq!(b.n(), module.n());
86 assert_eq!(self.n(), module.n());
87 assert_eq!(a.basek(), b.basek());
88 assert!(self.rank() >= a.rank().max(b.rank()));
89 }
90
91 let min_col: usize = a.rank().min(b.rank()) + 1;
92 let max_col: usize = a.rank().max(b.rank() + 1);
93 let self_col: usize = self.rank() + 1;
94
95 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
96 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
97 let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
98
99 (0..min_col).for_each(|i| {
100 module.vec_znx_sub(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
101 });
102
103 if a.rank() > b.rank() {
104 (min_col..max_col).for_each(|i| {
105 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
106 });
107 } else {
108 (min_col..max_col).for_each(|i| {
109 module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
110 module.vec_znx_negate_inplace(&mut self_mut.data, i);
111 });
112 }
113
114 let size: usize = self_mut.size();
115 (max_col..self_col).for_each(|i| {
116 (0..size).for_each(|j| {
117 self_mut.data.zero_at(i, j);
118 });
119 });
120
121 self.set_basek(a.basek());
122 self.set_k(set_k(self, a, b));
123 }
124
125 fn sub_inplace_ab<A>(&mut self, module: &Module<FFT64>, a: &A)
126 where
127 A: GLWECiphertextToRef + Infos,
128 {
129 #[cfg(debug_assertions)]
130 {
131 assert_eq!(a.n(), module.n());
132 assert_eq!(self.n(), module.n());
133 assert_eq!(self.basek(), a.basek());
134 assert!(self.rank() >= a.rank())
135 }
136
137 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
138 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
139
140 (0..a.rank() + 1).for_each(|i| {
141 module.vec_znx_sub_ab_inplace(&mut self_mut.data, i, &a_ref.data, i);
142 });
143
144 if a.rank() != 0 {
145 self.set_k(a.k().min(self.k()));
146 }
147 }
148
149 fn sub_inplace_ba<A>(&mut self, module: &Module<FFT64>, a: &A)
150 where
151 A: GLWECiphertextToRef + Infos,
152 {
153 #[cfg(debug_assertions)]
154 {
155 assert_eq!(a.n(), module.n());
156 assert_eq!(self.n(), module.n());
157 assert_eq!(self.basek(), a.basek());
158 assert!(self.rank() >= a.rank())
159 }
160
161 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
162 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
163
164 (0..a.rank() + 1).for_each(|i| {
165 module.vec_znx_sub_ba_inplace(&mut self_mut.data, i, &a_ref.data, i);
166 });
167
168 if a.rank() != 0 {
169 self.set_k(a.k().min(self.k()));
170 }
171 }
172
173 fn rotate<A>(&mut self, module: &Module<FFT64>, k: i64, a: &A)
174 where
175 A: GLWECiphertextToRef + Infos,
176 {
177 #[cfg(debug_assertions)]
178 {
179 assert_eq!(a.n(), module.n());
180 assert_eq!(self.n(), module.n());
181 assert_eq!(self.rank(), a.rank())
182 }
183
184 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
185 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
186
187 (0..a.rank() + 1).for_each(|i| {
188 module.vec_znx_rotate(k, &mut self_mut.data, i, &a_ref.data, i);
189 });
190
191 self.set_basek(a.basek());
192 if a.rank() != 0 {
193 self.set_k(a.k().min(self.k()));
194 }
195 }
196
197 fn rotate_inplace(&mut self, module: &Module<FFT64>, k: i64) {
198 #[cfg(debug_assertions)]
199 {
200 assert_eq!(self.n(), module.n());
201 }
202
203 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
204
205 (0..self_mut.rank() + 1).for_each(|i| {
206 module.vec_znx_rotate_inplace(k, &mut self_mut.data, i);
207 });
208 }
209
210 fn copy<A>(&mut self, module: &Module<FFT64>, a: &A)
211 where
212 A: GLWECiphertextToRef + Infos,
213 {
214 #[cfg(debug_assertions)]
215 {
216 assert_eq!(self.n(), module.n());
217 assert_eq!(a.n(), module.n());
218 assert_eq!(self.rank(), a.rank());
219 assert_eq!(self.k(), a.k());
220 assert_eq!(self.basek(), a.basek());
221 }
222
223 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
224 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
225
226 (0..self_mut.rank() + 1).for_each(|i| {
227 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
228 });
229
230 self.set_k(a.k());
231 self.set_basek(a.basek());
232 }
233
234 fn rsh(&mut self, k: usize, scratch: &mut Scratch) {
235 let basek: usize = self.basek();
236 let mut self_mut: GLWECiphertext<&mut [u8]> = self.to_mut();
237 self_mut.data.rsh(basek, k, scratch);
238 }
239
240 fn normalize<A>(&mut self, module: &Module<FFT64>, a: &A, scratch: &mut Scratch)
241 where
242 A: GLWECiphertextToRef,
243 {
244 #[cfg(debug_assertions)]
245 {
246 assert_eq!(self.n(), module.n());
247 assert_eq!(a.n(), module.n());
248 assert_eq!(self.rank(), a.rank());
249 }
250
251 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
252 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
253
254 (0..self_mut.rank() + 1).for_each(|i| {
255 module.vec_znx_normalize(a.basek(), &mut self_mut.data, i, &a_ref.data, i, scratch);
256 });
257 self.set_basek(a.basek());
258 self.set_k(a.k().min(self.k()));
259 }
260
261 fn normalize_inplace(&mut self, module: &Module<FFT64>, scratch: &mut Scratch) {
262 #[cfg(debug_assertions)]
263 {
264 assert_eq!(self.n(), module.n());
265 }
266 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
267 (0..self_mut.rank() + 1).for_each(|i| {
268 module.vec_znx_normalize_inplace(self_mut.basek(), &mut self_mut.data, i, scratch);
269 });
270 }
271}
272
273impl GLWECiphertext<Vec<u8>> {
274 pub fn rsh_scratch_space(module: &Module<FFT64>) -> usize {
275 VecZnx::rsh_scratch_space(module.n())
276 }
277}
278
279// c = op(a, b)
280fn set_k(c: &impl Infos, a: &impl Infos, b: &impl Infos) -> usize {
281 if a.rank() != 0 || b.rank() != 0 {
282 let k = if a.rank() == 0 {
283 b.k()
284 } else if b.rank() == 0 {
285 a.k()
286 } else {
287 a.k().min(b.k())
288 };
289 k.min(c.k())
290 } else {
291 c.k()
292 }
293}