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