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