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, TakeGLWESecretPrepared,
12    layouts::{GGLWETensorKey, GLWESecret, Infos, compressed::GGLWETensorKeyCompressed, prepared::Prepare},
13    trait_families::GLWEDecryptFamily,
14};
15
16use crate::trait_families::{GGLWETensorKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
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> + GLWESecretPreparedModuleFamily<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>:
38            GGLWETensorKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace + GLWESecretPreparedModuleFamily<B>,
39        Scratch<B>:
40            ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretPrepared<B> + TakeScalarZnx + TakeVecZnx,
41    {
42        #[cfg(debug_assertions)]
43        {
44            assert_eq!(self.rank(), sk.rank());
45            assert_eq!(self.n(), sk.n());
46        }
47
48        let n: usize = sk.n();
49        let rank: usize = self.rank();
50
51        let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_prepared(n, rank);
52        sk_dft_prep.prepare(module, &sk, scratch1);
53
54        let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(n, rank, 1);
55
56        (0..rank).for_each(|i| {
57            module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
58        });
59
60        let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(n, 1, 1);
61        let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(n, 1);
62        let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(n, 1, 1);
63
64        let mut source_xa: Source = Source::new(seed_xa);
65
66        (0..rank).for_each(|i| {
67            (i..rank).for_each(|j| {
68                module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
69
70                module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
71                module.vec_znx_big_normalize(
72                    self.basek(),
73                    &mut sk_ij.data.as_vec_znx_mut(),
74                    0,
75                    &sk_ij_big,
76                    0,
77                    scratch5,
78                );
79
80                let (seed_xa_tmp, _) = source_xa.branch();
81
82                self.at_mut(i, j)
83                    .encrypt_sk(module, &sk_ij, sk, seed_xa_tmp, source_xe, sigma, scratch5);
84            });
85        })
86    }
87}