1use poulpy_hal::{
 2    api::{
 3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddNormal, VecZnxFillUniform, VecZnxNormalizeInplace, ZnxView, ZnxViewMut,
 4    },
 5    layouts::{Backend, DataMut, DataRef, Module, ScratchOwned, VecZnx},
 6    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 7    source::Source,
 8};
 9
10use crate::{
11    encryption::{SIGMA, SIGMA_BOUND},
12    layouts::{Infos, LWECiphertext, LWEPlaintext, LWESecret},
13};
14
15impl<DataSelf: DataMut> LWECiphertext<DataSelf> {
16    pub fn encrypt_sk<DataPt, DataSk, B>(
17        &mut self,
18        module: &Module<B>,
19        pt: &LWEPlaintext<DataPt>,
20        sk: &LWESecret<DataSk>,
21        source_xa: &mut Source,
22        source_xe: &mut Source,
23    ) where
24        DataPt: DataRef,
25        DataSk: DataRef,
26        Module<B>: VecZnxFillUniform + VecZnxAddNormal + VecZnxNormalizeInplace<B>,
27        B: Backend + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
28    {
29        #[cfg(debug_assertions)]
30        {
31            assert_eq!(self.n(), sk.n())
32        }
33
34        let basek: usize = self.basek();
35        let k: usize = self.k();
36
37        module.vec_znx_fill_uniform(basek, &mut self.data, 0, k, source_xa);
38
39        let mut tmp_znx: VecZnx<Vec<u8>> = VecZnx::alloc(1, 1, self.size());
40
41        let min_size = self.size().min(pt.size());
42
43        (0..min_size).for_each(|i| {
44            tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0]
45                - self.data.at(0, i)[1..]
46                    .iter()
47                    .zip(sk.data.at(0, 0))
48                    .map(|(x, y)| x * y)
49                    .sum::<i64>();
50        });
51
52        (min_size..self.size()).for_each(|i| {
53            tmp_znx.at_mut(0, i)[0] -= self.data.at(0, i)[1..]
54                .iter()
55                .zip(sk.data.at(0, 0))
56                .map(|(x, y)| x * y)
57                .sum::<i64>();
58        });
59
60        module.vec_znx_add_normal(basek, &mut self.data, 0, k, source_xe, SIGMA, SIGMA_BOUND);
61
62        module.vec_znx_normalize_inplace(
63            basek,
64            &mut tmp_znx,
65            0,
66            ScratchOwned::alloc(size_of::<i64>()).borrow(),
67        );
68
69        (0..self.size()).for_each(|i| {
70            self.data.at_mut(0, i)[0] = tmp_znx.at(0, i)[0];
71        });
72    }
73}