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