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