1use poulpy_hal::{
  2    api::{
  3        Convolution, ModuleN, VecZnxAdd, VecZnxAddInplace, VecZnxCopy, VecZnxMulXpMinusOne, VecZnxMulXpMinusOneInplace, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSubNegateInplace, VecZnxZero
  4    },
  5    layouts::{Backend, Module, Scratch, VecZnx, ZnxZero},
  6    reference::vec_znx::vec_znx_rotate_inplace_tmp_bytes,
  7};
  8
  9use crate::{
 10    layouts::{GLWEInfos, GLWETensor, GLWETensorToMut, GLWEToMut, GLWEToRef, LWEInfos, SetGLWEInfos, TorusPrecision, GLWE}, ScratchTakeCore
 11};
 12
 13pub trait GLWETensoring<BE: Backend> where Self: Convolution<BE>, Scratch<BE>: ScratchTakeCore<BE> {
 14    fn glwe_tensor<R, A, B>(&self, res: &mut R, res_scale: i64, a: &A, b: &B, scratch: &mut Scratch<BE>) where R: GLWETensorToMut, A: GLWEToRef, B: GLWEToRef{
 15
 16        let res: &mut GLWETensor<&mut [u8]> = &mut res.to_mut();
 17        let a: &mut GLWE<&[u8]> = &mut a.to_ref();
 18        let b: &GLWE<&[u8]> = &b.to_ref();
 19
 20        self.bivariate_convolution(res.data_mut(), res_scale, a, b, scratch);
 21    }
 22}
 23
 24pub trait GLWEAdd
 25where
 26    Self: ModuleN + VecZnxAdd + VecZnxCopy + VecZnxAddInplace,
 27{
 28    fn glwe_add<R, A, B>(&self, res: &mut R, a: &A, b: &B)
 29    where
 30        R: GLWEToMut,
 31        A: GLWEToRef,
 32        B: GLWEToRef,
 33    {
 34        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
 35        let a: &mut GLWE<&[u8]> = &mut a.to_ref();
 36        let b: &GLWE<&[u8]> = &b.to_ref();
 37
 38        assert_eq!(a.n(), self.n() as u32);
 39        assert_eq!(b.n(), self.n() as u32);
 40        assert_eq!(res.n(), self.n() as u32);
 41        assert_eq!(a.base2k(), b.base2k());
 42        assert!(res.rank() >= a.rank().max(b.rank()));
 43
 44        let min_col: usize = (a.rank().min(b.rank()) + 1).into();
 45        let max_col: usize = (a.rank().max(b.rank() + 1)).into();
 46        let self_col: usize = (res.rank() + 1).into();
 47
 48        (0..min_col).for_each(|i| {
 49            self.vec_znx_add(res.data_mut(), i, a.data(), i, b.data(), i);
 50        });
 51
 52        if a.rank() > b.rank() {
 53            (min_col..max_col).for_each(|i| {
 54                self.vec_znx_copy(res.data_mut(), i, a.data(), i);
 55            });
 56        } else {
 57            (min_col..max_col).for_each(|i| {
 58                self.vec_znx_copy(res.data_mut(), i, b.data(), i);
 59            });
 60        }
 61
 62        let size: usize = res.size();
 63        (max_col..self_col).for_each(|i| {
 64            (0..size).for_each(|j| {
 65                res.data.zero_at(i, j);
 66            });
 67        });
 68
 69        res.set_base2k(a.base2k());
 70        res.set_k(set_k_binary(res, a, b));
 71    }
 72
 73    fn glwe_add_inplace<R, A>(&self, res: &mut R, a: &A)
 74    where
 75        R: GLWEToMut,
 76        A: GLWEToRef,
 77    {
 78        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
 79        let a: &GLWE<&[u8]> = &a.to_ref();
 80
 81        assert_eq!(res.n(), self.n() as u32);
 82        assert_eq!(a.n(), self.n() as u32);
 83        assert_eq!(res.base2k(), a.base2k());
 84        assert!(res.rank() >= a.rank());
 85
 86        (0..(a.rank() + 1).into()).for_each(|i| {
 87            self.vec_znx_add_inplace(res.data_mut(), i, a.data(), i);
 88        });
 89
 90        res.set_k(set_k_unary(res, a))
 91    }
 92}
 93
 94impl<BE: Backend> GLWEAdd for Module<BE> where Self: ModuleN + VecZnxAdd + VecZnxCopy + VecZnxAddInplace {}
 95
 96impl<BE: Backend> GLWESub for Module<BE> where
 97    Self: ModuleN + VecZnxSub + VecZnxCopy + VecZnxNegateInplace + VecZnxSubInplace + VecZnxSubNegateInplace
 98{
 99}
100
101pub trait GLWESub
102where
103    Self: ModuleN + VecZnxSub + VecZnxCopy + VecZnxNegateInplace + VecZnxSubInplace + VecZnxSubNegateInplace,
104{
105    fn glwe_sub<R, A, B>(&self, res: &mut R, a: &A, b: &B)
106    where
107        R: GLWEToMut,
108        A: GLWEToRef,
109        B: GLWEToRef,
110    {
111        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
112        let a: &GLWE<&[u8]> = &a.to_ref();
113        let b: &GLWE<&[u8]> = &b.to_ref();
114
115        assert_eq!(a.n(), self.n() as u32);
116        assert_eq!(b.n(), self.n() as u32);
117        assert_eq!(a.base2k(), b.base2k());
118        assert!(res.rank() >= a.rank().max(b.rank()));
119
120        let min_col: usize = (a.rank().min(b.rank()) + 1).into();
121        let max_col: usize = (a.rank().max(b.rank() + 1)).into();
122        let self_col: usize = (res.rank() + 1).into();
123
124        (0..min_col).for_each(|i| {
125            self.vec_znx_sub(res.data_mut(), i, a.data(), i, b.data(), i);
126        });
127
128        if a.rank() > b.rank() {
129            (min_col..max_col).for_each(|i| {
130                self.vec_znx_copy(res.data_mut(), i, a.data(), i);
131            });
132        } else {
133            (min_col..max_col).for_each(|i| {
134                self.vec_znx_copy(res.data_mut(), i, b.data(), i);
135                self.vec_znx_negate_inplace(res.data_mut(), i);
136            });
137        }
138
139        let size: usize = res.size();
140        (max_col..self_col).for_each(|i| {
141            (0..size).for_each(|j| {
142                res.data.zero_at(i, j);
143            });
144        });
145
146        res.set_base2k(a.base2k());
147        res.set_k(set_k_binary(res, a, b));
148    }
149
150    fn glwe_sub_inplace<R, A>(&self, res: &mut R, a: &A)
151    where
152        R: GLWEToMut,
153        A: GLWEToRef,
154    {
155        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
156        let a: &GLWE<&[u8]> = &a.to_ref();
157
158        assert_eq!(res.n(), self.n() as u32);
159        assert_eq!(a.n(), self.n() as u32);
160        assert_eq!(res.base2k(), a.base2k());
161        assert!(res.rank() >= a.rank());
162
163        (0..(a.rank() + 1).into()).for_each(|i| {
164            self.vec_znx_sub_inplace(res.data_mut(), i, a.data(), i);
165        });
166
167        res.set_k(set_k_unary(res, a))
168    }
169
170    fn glwe_sub_negate_inplace<R, A>(&self, res: &mut R, a: &A)
171    where
172        R: GLWEToMut,
173        A: GLWEToRef,
174    {
175        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
176        let a: &GLWE<&[u8]> = &a.to_ref();
177
178        assert_eq!(res.n(), self.n() as u32);
179        assert_eq!(a.n(), self.n() as u32);
180        assert_eq!(res.base2k(), a.base2k());
181        assert!(res.rank() >= a.rank());
182
183        (0..(a.rank() + 1).into()).for_each(|i| {
184            self.vec_znx_sub_negate_inplace(res.data_mut(), i, a.data(), i);
185        });
186
187        res.set_k(set_k_unary(res, a))
188    }
189}
190
191impl<BE: Backend> GLWERotate<BE> for Module<BE> where Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE> {}
192
193pub trait GLWERotate<BE: Backend>
194where
195    Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE>,
196{
197    fn glwe_rotate_tmp_bytes(&self) -> usize {
198        vec_znx_rotate_inplace_tmp_bytes(self.n())
199    }
200
201    fn glwe_rotate<R, A>(&self, k: i64, res: &mut R, a: &A)
202    where
203        R: GLWEToMut,
204        A: GLWEToRef,
205    {
206        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
207        let a: &GLWE<&[u8]> = &a.to_ref();
208
209        assert_eq!(a.n(), self.n() as u32);
210        assert_eq!(res.rank(), a.rank());
211
212        (0..(a.rank() + 1).into()).for_each(|i| {
213            self.vec_znx_rotate(k, res.data_mut(), i, a.data(), i);
214        });
215
216        res.set_base2k(a.base2k());
217        res.set_k(set_k_unary(res, a))
218    }
219
220    fn glwe_rotate_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
221    where
222        R: GLWEToMut,
223        Scratch<BE>: ScratchTakeCore<BE>,
224    {
225        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
226
227        (0..(res.rank() + 1).into()).for_each(|i| {
228            self.vec_znx_rotate_inplace(k, res.data_mut(), i, scratch);
229        });
230    }
231}
232
233impl<BE: Backend> GLWEMulXpMinusOne<BE> for Module<BE> where Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE> {}
234
235pub trait GLWEMulXpMinusOne<BE: Backend>
236where
237    Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE>,
238{
239    fn glwe_mul_xp_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A)
240    where
241        R: GLWEToMut,
242        A: GLWEToRef,
243    {
244        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
245        let a: &GLWE<&[u8]> = &a.to_ref();
246
247        assert_eq!(res.n(), self.n() as u32);
248        assert_eq!(a.n(), self.n() as u32);
249        assert_eq!(res.rank(), a.rank());
250
251        for i in 0..res.rank().as_usize() + 1 {
252            self.vec_znx_mul_xp_minus_one(k, res.data_mut(), i, a.data(), i);
253        }
254
255        res.set_base2k(a.base2k());
256        res.set_k(set_k_unary(res, a))
257    }
258
259    fn glwe_mul_xp_minus_one_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
260    where
261        R: GLWEToMut,
262    {
263        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
264
265        assert_eq!(res.n(), self.n() as u32);
266
267        for i in 0..res.rank().as_usize() + 1 {
268            self.vec_znx_mul_xp_minus_one_inplace(k, res.data_mut(), i, scratch);
269        }
270    }
271}
272
273impl<BE: Backend> GLWECopy for Module<BE> where Self: ModuleN + VecZnxCopy + VecZnxZero {}
274
275pub trait GLWECopy
276where
277    Self: ModuleN + VecZnxCopy + VecZnxZero,
278{
279    fn glwe_copy<R, A>(&self, res: &mut R, a: &A)
280    where
281        R: GLWEToMut,
282        A: GLWEToRef,
283    {
284        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
285        let a: &GLWE<&[u8]> = &a.to_ref();
286
287        assert_eq!(res.n(), self.n() as u32);
288        assert_eq!(a.n(), self.n() as u32);
289
290        let min_rank: usize = res.rank().min(a.rank()).as_usize() + 1;
291
292        for i in 0..min_rank {
293            self.vec_znx_copy(res.data_mut(), i, a.data(), i);
294        }
295
296        for i in min_rank..(res.rank() + 1).into() {
297            self.vec_znx_zero(res.data_mut(), i);
298        }
299
300        res.set_k(a.k().min(res.max_k()));
301        res.set_base2k(a.base2k());
302    }
303}
304
305impl<BE: Backend> GLWEShift<BE> for Module<BE> where Self: ModuleN + VecZnxRshInplace<BE> {}
306
307pub trait GLWEShift<BE: Backend>
308where
309    Self: ModuleN + VecZnxRshInplace<BE>,
310{
311    fn glwe_rsh_tmp_byte(&self) -> usize {
312        VecZnx::rsh_tmp_bytes(self.n())
313    }
314
315    fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
316    where
317        R: GLWEToMut,
318        Scratch<BE>: ScratchTakeCore<BE>,
319    {
320        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
321        let base2k: usize = res.base2k().into();
322        for i in 0..res.rank().as_usize() + 1 {
323            self.vec_znx_rsh_inplace(base2k, k, res.data_mut(), i, scratch);
324        }
325    }
326}
327
328impl GLWE<Vec<u8>> {
329    pub fn rsh_tmp_bytes<M, BE: Backend>(module: &M) -> usize
330    where
331        M: GLWEShift<BE>,
332    {
333        module.glwe_rsh_tmp_byte()
334    }
335}
336
337impl<BE: Backend> GLWENormalize<BE> for Module<BE> where Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> {}
338
339pub trait GLWENormalize<BE: Backend>
340where
341    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE>,
342{
343    fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
344    where
345        R: GLWEToMut,
346        A: GLWEToRef,
347        Scratch<BE>: ScratchTakeCore<BE>,
348    {
349        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
350        let a: &GLWE<&[u8]> = &a.to_ref();
351
352        assert_eq!(res.n(), self.n() as u32);
353        assert_eq!(a.n(), self.n() as u32);
354        assert_eq!(res.rank(), a.rank());
355
356        for i in 0..res.rank().as_usize() + 1 {
357            self.vec_znx_normalize(
358                res.base2k().into(),
359                res.data_mut(),
360                i,
361                a.base2k().into(),
362                a.data(),
363                i,
364                scratch,
365            );
366        }
367
368        res.set_k(a.k().min(res.k()));
369    }
370
371    fn glwe_normalize_inplace<R>(&self, res: &mut R, scratch: &mut Scratch<BE>)
372    where
373        R: GLWEToMut,
374        Scratch<BE>: ScratchTakeCore<BE>,
375    {
376        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
377        for i in 0..res.rank().as_usize() + 1 {
378            self.vec_znx_normalize_inplace(res.base2k().into(), res.data_mut(), i, scratch);
379        }
380    }
381}
382
383// c = op(a, b)
384fn set_k_binary(c: &impl GLWEInfos, a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
385    // If either operands is a ciphertext
386    if a.rank() != 0 || b.rank() != 0 {
387        // If a is a plaintext (but b ciphertext)
388        let k = if a.rank() == 0 {
389            b.k()
390        // If b is a plaintext (but a ciphertext)
391        } else if b.rank() == 0 {
392            a.k()
393        // If a & b are both ciphertexts
394        } else {
395            a.k().min(b.k())
396        };
397        k.min(c.k())
398    // If a & b are both plaintexts
399    } else {
400        c.k()
401    }
402}
403
404// a = op(a, b)
405fn set_k_unary(a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
406    if a.rank() != 0 || b.rank() != 0 {
407        a.k().min(b.k())
408    } else {
409        a.k()
410    }
411}