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_binary(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 self.set_k(set_k_unary(self, a))
73 }
74
75 fn sub<A, B>(&mut self, module: &Module<FFT64>, a: &A, b: &B)
76 where
77 A: GLWECiphertextToRef,
78 B: GLWECiphertextToRef,
79 {
80 #[cfg(debug_assertions)]
81 {
82 assert_eq!(a.n(), module.n());
83 assert_eq!(b.n(), module.n());
84 assert_eq!(self.n(), module.n());
85 assert_eq!(a.basek(), b.basek());
86 assert!(self.rank() >= a.rank().max(b.rank()));
87 }
88
89 let min_col: usize = a.rank().min(b.rank()) + 1;
90 let max_col: usize = a.rank().max(b.rank() + 1);
91 let self_col: usize = self.rank() + 1;
92
93 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
94 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
95 let b_ref: &GLWECiphertext<&[u8]> = &b.to_ref();
96
97 (0..min_col).for_each(|i| {
98 module.vec_znx_sub(&mut self_mut.data, i, &a_ref.data, i, &b_ref.data, i);
99 });
100
101 if a.rank() > b.rank() {
102 (min_col..max_col).for_each(|i| {
103 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
104 });
105 } else {
106 (min_col..max_col).for_each(|i| {
107 module.vec_znx_copy(&mut self_mut.data, i, &b_ref.data, i);
108 module.vec_znx_negate_inplace(&mut self_mut.data, i);
109 });
110 }
111
112 let size: usize = self_mut.size();
113 (max_col..self_col).for_each(|i| {
114 (0..size).for_each(|j| {
115 self_mut.data.zero_at(i, j);
116 });
117 });
118
119 self.set_basek(a.basek());
120 self.set_k(set_k_binary(self, a, b));
121 }
122
123 fn sub_inplace_ab<A>(&mut self, module: &Module<FFT64>, a: &A)
124 where
125 A: GLWECiphertextToRef + Infos,
126 {
127 #[cfg(debug_assertions)]
128 {
129 assert_eq!(a.n(), module.n());
130 assert_eq!(self.n(), module.n());
131 assert_eq!(self.basek(), a.basek());
132 assert!(self.rank() >= a.rank())
133 }
134
135 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
136 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
137
138 (0..a.rank() + 1).for_each(|i| {
139 module.vec_znx_sub_ab_inplace(&mut self_mut.data, i, &a_ref.data, i);
140 });
141
142 self.set_k(set_k_unary(self, a))
143 }
144
145 fn sub_inplace_ba<A>(&mut self, module: &Module<FFT64>, a: &A)
146 where
147 A: GLWECiphertextToRef + Infos,
148 {
149 #[cfg(debug_assertions)]
150 {
151 assert_eq!(a.n(), module.n());
152 assert_eq!(self.n(), module.n());
153 assert_eq!(self.basek(), a.basek());
154 assert!(self.rank() >= a.rank())
155 }
156
157 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
158 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
159
160 (0..a.rank() + 1).for_each(|i| {
161 module.vec_znx_sub_ba_inplace(&mut self_mut.data, i, &a_ref.data, i);
162 });
163
164 self.set_k(set_k_unary(self, a))
165 }
166
167 fn rotate<A>(&mut self, module: &Module<FFT64>, k: i64, a: &A)
168 where
169 A: GLWECiphertextToRef + Infos,
170 {
171 #[cfg(debug_assertions)]
172 {
173 assert_eq!(a.n(), module.n());
174 assert_eq!(self.n(), module.n());
175 assert_eq!(self.rank(), a.rank())
176 }
177
178 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
179 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
180
181 (0..a.rank() + 1).for_each(|i| {
182 module.vec_znx_rotate(k, &mut self_mut.data, i, &a_ref.data, i);
183 });
184
185 self.set_basek(a.basek());
186 self.set_k(set_k_unary(self, a))
187 }
188
189 fn rotate_inplace(&mut self, module: &Module<FFT64>, k: i64) {
190 #[cfg(debug_assertions)]
191 {
192 assert_eq!(self.n(), module.n());
193 }
194
195 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
196
197 (0..self_mut.rank() + 1).for_each(|i| {
198 module.vec_znx_rotate_inplace(k, &mut self_mut.data, i);
199 });
200 }
201
202 fn copy<A>(&mut self, module: &Module<FFT64>, a: &A)
203 where
204 A: GLWECiphertextToRef + Infos,
205 {
206 #[cfg(debug_assertions)]
207 {
208 assert_eq!(self.n(), module.n());
209 assert_eq!(a.n(), module.n());
210 assert_eq!(self.rank(), a.rank());
211 }
212
213 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
214 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
215
216 (0..self_mut.rank() + 1).for_each(|i| {
217 module.vec_znx_copy(&mut self_mut.data, i, &a_ref.data, i);
218 });
219
220 self.set_k(a.k().min(self.size() * self.basek()));
221 self.set_basek(a.basek());
222 }
223
224 fn rsh(&mut self, k: usize, scratch: &mut Scratch) {
225 let basek: usize = self.basek();
226 let mut self_mut: GLWECiphertext<&mut [u8]> = self.to_mut();
227 self_mut.data.rsh(basek, k, scratch);
228 }
229
230 fn normalize<A>(&mut self, module: &Module<FFT64>, a: &A, scratch: &mut Scratch)
231 where
232 A: GLWECiphertextToRef,
233 {
234 #[cfg(debug_assertions)]
235 {
236 assert_eq!(self.n(), module.n());
237 assert_eq!(a.n(), module.n());
238 assert_eq!(self.rank(), a.rank());
239 }
240
241 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
242 let a_ref: &GLWECiphertext<&[u8]> = &a.to_ref();
243
244 (0..self_mut.rank() + 1).for_each(|i| {
245 module.vec_znx_normalize(a.basek(), &mut self_mut.data, i, &a_ref.data, i, scratch);
246 });
247 self.set_basek(a.basek());
248 self.set_k(a.k().min(self.k()));
249 }
250
251 fn normalize_inplace(&mut self, module: &Module<FFT64>, scratch: &mut Scratch) {
252 #[cfg(debug_assertions)]
253 {
254 assert_eq!(self.n(), module.n());
255 }
256 let self_mut: &mut GLWECiphertext<&mut [u8]> = &mut self.to_mut();
257 (0..self_mut.rank() + 1).for_each(|i| {
258 module.vec_znx_normalize_inplace(self_mut.basek(), &mut self_mut.data, i, scratch);
259 });
260 }
261}
262
263impl GLWECiphertext<Vec<u8>> {
264 pub fn rsh_scratch_space(module: &Module<FFT64>) -> usize {
265 VecZnx::rsh_scratch_space(module.n())
266 }
267}
268
269// c = op(a, b)
270fn set_k_binary(c: &impl Infos, a: &impl Infos, b: &impl Infos) -> usize {
271 // If either operands is a ciphertext
272 if a.rank() != 0 || b.rank() != 0 {
273 // If a is a plaintext (but b ciphertext)
274 let k = if a.rank() == 0 {
275 b.k()
276 // If b is a plaintext (but a ciphertext)
277 } else if b.rank() == 0 {
278 a.k()
279 // If a & b are both ciphertexts
280 } else {
281 a.k().min(b.k())
282 };
283 k.min(c.k())
284 // If a & b are both plaintexts
285 } else {
286 c.k()
287 }
288}
289
290// a = op(a, b)
291fn set_k_unary(a: &impl Infos, b: &impl Infos) -> usize {
292 if a.rank() != 0 || b.rank() != 0 {
293 a.k().min(b.k())
294 } else {
295 a.k()
296 }
297}