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