1use backend::hal::{
 2    api::{
 3        ScratchAvailable, SvpApply, TakeScalarZnx, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddScalarInplace,
 4        VecZnxDftToVecZnxBigTmpA, VecZnxSwithcDegree,
 5    },
 6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 7};
 8use sampling::source::Source;
 9
10use crate::{
11    TakeGLWESecret, TakeGLWESecretExec,
12    layouts::{GGLWETensorKey, GLWESecret, Infos, compressed::GGLWETensorKeyCompressed},
13    trait_families::GLWEDecryptFamily,
14};
15
16use crate::trait_families::{GGLWETensorKeyEncryptSkFamily, GLWESecretExecModuleFamily};
17
18impl GGLWETensorKeyCompressed<Vec<u8>> {
19    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
20    where
21        Module<B>: GGLWETensorKeyEncryptSkFamily<B> + GLWESecretExecModuleFamily<B>,
22    {
23        GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k, rank)
24    }
25}
26
27impl<DataSelf: DataMut> GGLWETensorKeyCompressed<DataSelf> {
28    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
29        &mut self,
30        module: &Module<B>,
31        sk: &GLWESecret<DataSk>,
32        seed_xa: [u8; 32],
33        source_xe: &mut Source,
34        sigma: f64,
35        scratch: &mut Scratch<B>,
36    ) where
37        Module<B>: GGLWETensorKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace + GLWESecretExecModuleFamily<B>,
38        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx + TakeVecZnx,
39    {
40        #[cfg(debug_assertions)]
41        {
42            assert_eq!(self.rank(), sk.rank());
43            assert_eq!(self.n(), sk.n());
44        }
45
46        let n: usize = sk.n();
47        let rank: usize = self.rank();
48
49        let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(n, rank);
50        sk_dft_prep.prepare(module, &sk);
51
52        let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(n, rank, 1);
53
54        (0..rank).for_each(|i| {
55            module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
56        });
57
58        let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(n, 1, 1);
59        let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(n, 1);
60        let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(n, 1, 1);
61
62        let mut source_xa: Source = Source::new(seed_xa);
63
64        (0..rank).for_each(|i| {
65            (i..rank).for_each(|j| {
66                module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
67
68                module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
69                module.vec_znx_big_normalize(
70                    self.basek(),
71                    &mut sk_ij.data.as_vec_znx_mut(),
72                    0,
73                    &sk_ij_big,
74                    0,
75                    scratch5,
76                );
77
78                let (seed_xa_tmp, _) = source_xa.branch();
79
80                self.at_mut(i, j)
81                    .encrypt_sk(module, &sk_ij, sk, seed_xa_tmp, source_xe, sigma, scratch5);
82            });
83        })
84    }
85}