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
182pub trait GLWERotate<BE: Backend>
183where
184    Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE>,
185{
186    fn glwe_rotate<R, A>(&self, k: i64, res: &mut R, a: &A)
187    where
188        R: GLWEToMut,
189        A: GLWEToRef,
190    {
191        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
192        let a: &GLWE<&[u8]> = &a.to_ref();
193
194        assert_eq!(a.n(), self.n() as u32);
195        assert_eq!(res.rank(), a.rank());
196
197        (0..(a.rank() + 1).into()).for_each(|i| {
198            self.vec_znx_rotate(k, res.data_mut(), i, a.data(), i);
199        });
200
201        res.set_base2k(a.base2k());
202        res.set_k(set_k_unary(res, a))
203    }
204
205    fn glwe_rotate_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
206    where
207        R: GLWEToMut,
208        Scratch<BE>: ScratchTakeCore<BE>,
209    {
210        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
211
212        (0..(res.rank() + 1).into()).for_each(|i| {
213            self.vec_znx_rotate_inplace(k, res.data_mut(), i, scratch);
214        });
215    }
216}
217
218pub trait GLWEMulXpMinusOne<BE: Backend>
219where
220    Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE>,
221{
222    fn glwe_mul_xp_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A)
223    where
224        R: GLWEToMut,
225        A: GLWEToRef,
226    {
227        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
228        let a: &GLWE<&[u8]> = &a.to_ref();
229
230        assert_eq!(res.n(), self.n() as u32);
231        assert_eq!(a.n(), self.n() as u32);
232        assert_eq!(res.rank(), a.rank());
233
234        for i in 0..res.rank().as_usize() + 1 {
235            self.vec_znx_mul_xp_minus_one(k, res.data_mut(), i, a.data(), i);
236        }
237
238        res.set_base2k(a.base2k());
239        res.set_k(set_k_unary(res, a))
240    }
241
242    fn glwe_mul_xp_minus_one_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
243    where
244        R: GLWEToMut,
245    {
246        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
247
248        assert_eq!(res.n(), self.n() as u32);
249
250        for i in 0..res.rank().as_usize() + 1 {
251            self.vec_znx_mul_xp_minus_one_inplace(k, res.data_mut(), i, scratch);
252        }
253    }
254}
255
256pub trait GLWECopy
257where
258    Self: ModuleN + VecZnxCopy,
259{
260    fn glwe_copy<R, A>(&self, res: &mut R, a: &A)
261    where
262        R: GLWEToMut,
263        A: GLWEToRef,
264    {
265        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
266        let a: &GLWE<&[u8]> = &a.to_ref();
267
268        assert_eq!(res.n(), self.n() as u32);
269        assert_eq!(a.n(), self.n() as u32);
270        assert_eq!(res.rank(), a.rank());
271
272        for i in 0..res.rank().as_usize() + 1 {
273            self.vec_znx_copy(res.data_mut(), i, a.data(), i);
274        }
275
276        res.set_k(a.k().min(res.max_k()));
277        res.set_base2k(a.base2k());
278    }
279}
280
281pub trait GLWEShift<BE: Backend>
282where
283    Self: ModuleN + VecZnxRshInplace<BE>,
284{
285    fn glwe_rsh_tmp_byte(&self) -> usize {
286        VecZnx::rsh_tmp_bytes(self.n())
287    }
288
289    fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
290    where
291        R: GLWEToMut,
292        Scratch<BE>: ScratchTakeCore<BE>,
293    {
294        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
295        let base2k: usize = res.base2k().into();
296        for i in 0..res.rank().as_usize() + 1 {
297            self.vec_znx_rsh_inplace(base2k, k, res.data_mut(), i, scratch);
298        }
299    }
300}
301
302impl GLWE<Vec<u8>> {
303    pub fn rsh_tmp_bytes<M, BE: Backend>(module: &M) -> usize
304    where
305        M: GLWEShift<BE>,
306    {
307        module.glwe_rsh_tmp_byte()
308    }
309}
310
311pub trait GLWENormalize<BE: Backend>
312where
313    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE>,
314{
315    fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
316    where
317        R: GLWEToMut,
318        A: GLWEToRef,
319        Scratch<BE>: ScratchTakeCore<BE>,
320    {
321        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
322        let a: &GLWE<&[u8]> = &a.to_ref();
323
324        assert_eq!(res.n(), self.n() as u32);
325        assert_eq!(a.n(), self.n() as u32);
326        assert_eq!(res.rank(), a.rank());
327
328        for i in 0..res.rank().as_usize() + 1 {
329            self.vec_znx_normalize(
330                res.base2k().into(),
331                res.data_mut(),
332                i,
333                a.base2k().into(),
334                a.data(),
335                i,
336                scratch,
337            );
338        }
339
340        res.set_k(a.k().min(res.k()));
341    }
342
343    fn glwe_normalize_inplace<R>(&self, res: &mut R, scratch: &mut Scratch<BE>)
344    where
345        R: GLWEToMut,
346        Scratch<BE>: ScratchTakeCore<BE>,
347    {
348        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
349        for i in 0..res.rank().as_usize() + 1 {
350            self.vec_znx_normalize_inplace(res.base2k().into(), res.data_mut(), i, scratch);
351        }
352    }
353}
354
355// c = op(a, b)
356fn set_k_binary(c: &impl GLWEInfos, a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
357    // If either operands is a ciphertext
358    if a.rank() != 0 || b.rank() != 0 {
359        // If a is a plaintext (but b ciphertext)
360        let k = if a.rank() == 0 {
361            b.k()
362        // If b is a plaintext (but a ciphertext)
363        } else if b.rank() == 0 {
364            a.k()
365        // If a & b are both ciphertexts
366        } else {
367            a.k().min(b.k())
368        };
369        k.min(c.k())
370    // If a & b are both plaintexts
371    } else {
372        c.k()
373    }
374}
375
376// a = op(a, b)
377fn set_k_unary(a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
378    if a.rank() != 0 || b.rank() != 0 {
379        a.k().min(b.k())
380    } else {
381        a.k()
382    }
383}