1use backend::hal::{
 2    api::{
 3        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxAllocBytes, VecZnxNormalizeInplace, ZnxZero,
 4    },
 5    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
 6};
 7use sampling::source::Source;
 8
 9use crate::{GGSWCiphertext, GLWECiphertext, GLWEEncryptSkFamily, GLWESecretExec, Infos, TakeGLWEPt};
10
11pub trait GGSWEncryptSkFamily<B: Backend> = GLWEEncryptSkFamily<B>;
12
13impl GGSWCiphertext<Vec<u8>> {
14    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
15    where
16        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAllocBytes,
17    {
18        let size = k.div_ceil(basek);
19        GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
20            + module.vec_znx_alloc_bytes(rank + 1, size)
21            + module.vec_znx_alloc_bytes(1, size)
22            + module.vec_znx_dft_alloc_bytes(rank + 1, size)
23    }
24}
25
26impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
27    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
28        &mut self,
29        module: &Module<B>,
30        pt: &ScalarZnx<DataPt>,
31        sk: &GLWESecretExec<DataSk, B>,
32        source_xa: &mut Source,
33        source_xe: &mut Source,
34        sigma: f64,
35        scratch: &mut Scratch<B>,
36    ) where
37        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAddScalarInplace,
38        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
39    {
40        #[cfg(debug_assertions)]
41        {
42            use backend::hal::api::ZnxInfos;
43
44            assert_eq!(self.rank(), sk.rank());
45            assert_eq!(self.n(), module.n());
46            assert_eq!(pt.n(), module.n());
47            assert_eq!(sk.n(), module.n());
48        }
49
50        let basek: usize = self.basek();
51        let k: usize = self.k();
52        let rank: usize = self.rank();
53        let digits: usize = self.digits();
54
55        let (mut tmp_pt, scratch1) = scratch.take_glwe_pt(module, basek, k);
56
57        (0..self.rows()).for_each(|row_i| {
58            tmp_pt.data.zero();
59
60            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
61            module.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (digits - 1) + row_i * digits, pt, 0);
62            module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch1);
63
64            (0..rank + 1).for_each(|col_j| {
65                // rlwe encrypt of vec_znx_pt into vec_znx_ct
66
67                self.at_mut(row_i, col_j).encrypt_sk_internal(
68                    module,
69                    Some((&tmp_pt, col_j)),
70                    sk,
71                    source_xa,
72                    source_xe,
73                    sigma,
74                    scratch1,
75                );
76            });
77        });
78    }
79}