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