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