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    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 {}
266
267pub trait GLWECopy
268where
269    Self: ModuleN + VecZnxCopy,
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        assert_eq!(res.rank(), a.rank());
282
283        for i in 0..res.rank().as_usize() + 1 {
284            self.vec_znx_copy(res.data_mut(), i, a.data(), i);
285        }
286
287        res.set_k(a.k().min(res.max_k()));
288        res.set_base2k(a.base2k());
289    }
290}
291
292impl<BE: Backend> GLWEShift<BE> for Module<BE> where Self: ModuleN + VecZnxRshInplace<BE> {}
293
294pub trait GLWEShift<BE: Backend>
295where
296    Self: ModuleN + VecZnxRshInplace<BE>,
297{
298    fn glwe_rsh_tmp_byte(&self) -> usize {
299        VecZnx::rsh_tmp_bytes(self.n())
300    }
301
302    fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
303    where
304        R: GLWEToMut,
305        Scratch<BE>: ScratchTakeCore<BE>,
306    {
307        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
308        let base2k: usize = res.base2k().into();
309        for i in 0..res.rank().as_usize() + 1 {
310            self.vec_znx_rsh_inplace(base2k, k, res.data_mut(), i, scratch);
311        }
312    }
313}
314
315impl GLWE<Vec<u8>> {
316    pub fn rsh_tmp_bytes<M, BE: Backend>(module: &M) -> usize
317    where
318        M: GLWEShift<BE>,
319    {
320        module.glwe_rsh_tmp_byte()
321    }
322}
323
324impl<BE: Backend> GLWENormalize<BE> for Module<BE> where Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> {}
325
326pub trait GLWENormalize<BE: Backend>
327where
328    Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE>,
329{
330    fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
331    where
332        R: GLWEToMut,
333        A: GLWEToRef,
334        Scratch<BE>: ScratchTakeCore<BE>,
335    {
336        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
337        let a: &GLWE<&[u8]> = &a.to_ref();
338
339        assert_eq!(res.n(), self.n() as u32);
340        assert_eq!(a.n(), self.n() as u32);
341        assert_eq!(res.rank(), a.rank());
342
343        for i in 0..res.rank().as_usize() + 1 {
344            self.vec_znx_normalize(
345                res.base2k().into(),
346                res.data_mut(),
347                i,
348                a.base2k().into(),
349                a.data(),
350                i,
351                scratch,
352            );
353        }
354
355        res.set_k(a.k().min(res.k()));
356    }
357
358    fn glwe_normalize_inplace<R>(&self, res: &mut R, scratch: &mut Scratch<BE>)
359    where
360        R: GLWEToMut,
361        Scratch<BE>: ScratchTakeCore<BE>,
362    {
363        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
364        for i in 0..res.rank().as_usize() + 1 {
365            self.vec_znx_normalize_inplace(res.base2k().into(), res.data_mut(), i, scratch);
366        }
367    }
368}
369
370// c = op(a, b)
371fn set_k_binary(c: &impl GLWEInfos, a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
372    // If either operands is a ciphertext
373    if a.rank() != 0 || b.rank() != 0 {
374        // If a is a plaintext (but b ciphertext)
375        let k = if a.rank() == 0 {
376            b.k()
377        // If b is a plaintext (but a ciphertext)
378        } else if b.rank() == 0 {
379            a.k()
380        // If a & b are both ciphertexts
381        } else {
382            a.k().min(b.k())
383        };
384        k.min(c.k())
385    // If a & b are both plaintexts
386    } else {
387        c.k()
388    }
389}
390
391// a = op(a, b)
392fn set_k_unary(a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
393    if a.rank() != 0 || b.rank() != 0 {
394        a.k().min(b.k())
395    } else {
396        a.k()
397    }
398}