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
 87pub trait GLWESub
 88where
 89    Self: ModuleN + VecZnxSub + VecZnxCopy + VecZnxNegateInplace + VecZnxSubInplace + VecZnxSubNegateInplace,
 90{
 91    fn glwe_sub<R, A, B>(&self, res: &mut R, a: &A, b: &B)
 92    where
 93        R: GLWEToMut,
 94        A: GLWEToRef,
 95        B: GLWEToRef,
 96    {
 97        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
 98        let a: &GLWE<&[u8]> = &a.to_ref();
 99        let b: &GLWE<&[u8]> = &b.to_ref();
100
101        assert_eq!(a.n(), self.n() as u32);
102        assert_eq!(b.n(), self.n() as u32);
103        assert_eq!(a.base2k(), b.base2k());
104        assert!(res.rank() >= a.rank().max(b.rank()));
105
106        let min_col: usize = (a.rank().min(b.rank()) + 1).into();
107        let max_col: usize = (a.rank().max(b.rank() + 1)).into();
108        let self_col: usize = (res.rank() + 1).into();
109
110        (0..min_col).for_each(|i| {
111            self.vec_znx_sub(res.data_mut(), i, a.data(), i, b.data(), i);
112        });
113
114        if a.rank() > b.rank() {
115            (min_col..max_col).for_each(|i| {
116                self.vec_znx_copy(res.data_mut(), i, a.data(), i);
117            });
118        } else {
119            (min_col..max_col).for_each(|i| {
120                self.vec_znx_copy(res.data_mut(), i, b.data(), i);
121                self.vec_znx_negate_inplace(res.data_mut(), i);
122            });
123        }
124
125        let size: usize = res.size();
126        (max_col..self_col).for_each(|i| {
127            (0..size).for_each(|j| {
128                res.data.zero_at(i, j);
129            });
130        });
131
132        res.set_base2k(a.base2k());
133        res.set_k(set_k_binary(res, a, b));
134    }
135
136    fn glwe_sub_inplace<R, A>(&self, res: &mut R, a: &A)
137    where
138        R: GLWEToMut,
139        A: GLWEToRef,
140    {
141        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
142        let a: &GLWE<&[u8]> = &a.to_ref();
143
144        assert_eq!(res.n(), self.n() as u32);
145        assert_eq!(a.n(), self.n() as u32);
146        assert_eq!(res.base2k(), a.base2k());
147        assert!(res.rank() >= a.rank());
148
149        (0..(a.rank() + 1).into()).for_each(|i| {
150            self.vec_znx_sub_inplace(res.data_mut(), i, a.data(), i);
151        });
152
153        res.set_k(set_k_unary(res, a))
154    }
155
156    fn glwe_sub_negate_inplace<R, A>(&self, res: &mut R, a: &A)
157    where
158        R: GLWEToMut,
159        A: GLWEToRef,
160    {
161        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
162        let a: &GLWE<&[u8]> = &a.to_ref();
163
164        assert_eq!(res.n(), self.n() as u32);
165        assert_eq!(a.n(), self.n() as u32);
166        assert_eq!(res.base2k(), a.base2k());
167        assert!(res.rank() >= a.rank());
168
169        (0..(a.rank() + 1).into()).for_each(|i| {
170            self.vec_znx_sub_negate_inplace(res.data_mut(), i, a.data(), i);
171        });
172
173        res.set_k(set_k_unary(res, a))
174    }
175}
176
177pub trait GLWERotate<BE: Backend>
178where
179    Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE>,
180{
181    fn glwe_rotate<R, A>(&self, k: i64, res: &mut R, a: &A)
182    where
183        R: GLWEToMut,
184        A: GLWEToRef,
185    {
186        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
187        let a: &GLWE<&[u8]> = &a.to_ref();
188
189        assert_eq!(a.n(), self.n() as u32);
190        assert_eq!(res.rank(), a.rank());
191
192        (0..(a.rank() + 1).into()).for_each(|i| {
193            self.vec_znx_rotate(k, res.data_mut(), i, a.data(), i);
194        });
195
196        res.set_base2k(a.base2k());
197        res.set_k(set_k_unary(res, a))
198    }
199
200    fn glwe_rotate_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
201    where
202        R: GLWEToMut,
203        Scratch<BE>: ScratchTakeCore<BE>,
204    {
205        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
206
207        (0..(res.rank() + 1).into()).for_each(|i| {
208            self.vec_znx_rotate_inplace(k, res.data_mut(), i, scratch);
209        });
210    }
211}
212
213pub trait GLWEMulXpMinusOne<BE: Backend>
214where
215    Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE>,
216{
217    fn glwe_mul_xp_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A)
218    where
219        R: GLWEToMut,
220        A: GLWEToRef,
221    {
222        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
223        let a: &GLWE<&[u8]> = &a.to_ref();
224
225        assert_eq!(res.n(), self.n() as u32);
226        assert_eq!(a.n(), self.n() as u32);
227        assert_eq!(res.rank(), a.rank());
228
229        for i in 0..res.rank().as_usize() + 1 {
230            self.vec_znx_mul_xp_minus_one(k, res.data_mut(), i, a.data(), i);
231        }
232
233        res.set_base2k(a.base2k());
234        res.set_k(set_k_unary(res, a))
235    }
236
237    fn glwe_mul_xp_minus_one_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
238    where
239        R: GLWEToMut,
240    {
241        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
242
243        assert_eq!(res.n(), self.n() as u32);
244
245        for i in 0..res.rank().as_usize() + 1 {
246            self.vec_znx_mul_xp_minus_one_inplace(k, res.data_mut(), i, scratch);
247        }
248    }
249}
250
251pub trait GLWECopy
252where
253    Self: ModuleN + VecZnxCopy,
254{
255    fn glwe_copy<R, A>(&self, res: &mut R, a: &A)
256    where
257        R: GLWEToMut,
258        A: GLWEToRef,
259    {
260        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
261        let a: &GLWE<&[u8]> = &a.to_ref();
262
263        assert_eq!(res.n(), self.n() as u32);
264        assert_eq!(a.n(), self.n() as u32);
265        assert_eq!(res.rank(), a.rank());
266
267        for i in 0..res.rank().as_usize() + 1 {
268            self.vec_znx_copy(res.data_mut(), i, a.data(), i);
269        }
270
271        res.set_k(a.k().min(res.max_k()));
272        res.set_base2k(a.base2k());
273    }
274}
275
276pub trait GLWEShift<BE: Backend>
277where
278    Self: ModuleN + VecZnxRshInplace<BE>,
279{
280    fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
281    where
282        R: GLWEToMut,
283        Scratch<BE>: ScratchTakeCore<BE>,
284    {
285        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
286        let base2k: usize = res.base2k().into();
287        for i in 0..res.rank().as_usize() + 1 {
288            self.vec_znx_rsh_inplace(base2k, k, res.data_mut(), i, scratch);
289        }
290    }
291}
292
293impl GLWE<Vec<u8>> {
294    pub fn rsh_tmp_bytes(n: usize) -> usize {
295        VecZnx::rsh_tmp_bytes(n)
296    }
297}
298
299pub trait GLWENormalize<BE: Backend>
300where
301    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE>,
302{
303    fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
304    where
305        R: GLWEToMut,
306        A: GLWEToRef,
307        Scratch<BE>: ScratchTakeCore<BE>,
308    {
309        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
310        let a: &GLWE<&[u8]> = &a.to_ref();
311
312        assert_eq!(res.n(), self.n() as u32);
313        assert_eq!(a.n(), self.n() as u32);
314        assert_eq!(res.rank(), a.rank());
315
316        for i in 0..res.rank().as_usize() + 1 {
317            self.vec_znx_normalize(
318                res.base2k().into(),
319                res.data_mut(),
320                i,
321                a.base2k().into(),
322                a.data(),
323                i,
324                scratch,
325            );
326        }
327
328        res.set_k(a.k().min(res.k()));
329    }
330
331    fn glwe_normalize_inplace<R>(&mut self, res: &mut R, scratch: &mut Scratch<BE>)
332    where
333        R: GLWEToMut,
334        Scratch<BE>: ScratchTakeCore<BE>,
335    {
336        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
337        for i in 0..res.rank().as_usize() + 1 {
338            self.vec_znx_normalize_inplace(res.base2k().into(), res.data_mut(), i, scratch);
339        }
340    }
341}
342
343// c = op(a, b)
344fn set_k_binary(c: &impl GLWEInfos, a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
345    // If either operands is a ciphertext
346    if a.rank() != 0 || b.rank() != 0 {
347        // If a is a plaintext (but b ciphertext)
348        let k = if a.rank() == 0 {
349            b.k()
350        // If b is a plaintext (but a ciphertext)
351        } else if b.rank() == 0 {
352            a.k()
353        // If a & b are both ciphertexts
354        } else {
355            a.k().min(b.k())
356        };
357        k.min(c.k())
358    // If a & b are both plaintexts
359    } else {
360        c.k()
361    }
362}
363
364// a = op(a, b)
365fn set_k_unary(a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
366    if a.rank() != 0 || b.rank() != 0 {
367        a.k().min(b.k())
368    } else {
369        a.k()
370    }
371}