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