1use backend::hal::{
 2    api::{
 3        ScratchAvailable, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxAutomorphismInplace,
 4        VecZnxSwithcDegree, ZnxView, ZnxViewMut, ZnxZero,
 5    },
 6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 7};
 8use sampling::source::Source;
 9
10use crate::{
11    TakeGLWESecret, TakeGLWESecretPrepared,
12    layouts::{GGLWESwitchingKey, GLWESecret, GLWEToLWESwitchingKey, LWESecret, prepared::GLWESecretPrepared},
13};
14
15use crate::trait_families::{GGLWEEncryptSkFamily, GGLWESwitchingKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
16
17impl GLWEToLWESwitchingKey<Vec<u8>> {
18    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank_in: usize) -> usize
19    where
20        Module<B>: GGLWEEncryptSkFamily<B> + GLWESecretPreparedModuleFamily<B>,
21    {
22        GLWESecretPrepared::bytes_of(module, n, rank_in)
23            + (GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank_in, 1) | GLWESecret::bytes_of(n, rank_in))
24    }
25}
26
27impl<D: DataMut> GLWEToLWESwitchingKey<D> {
28    pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
29        &mut self,
30        module: &Module<B>,
31        sk_lwe: &LWESecret<DLwe>,
32        sk_glwe: &GLWESecret<DGlwe>,
33        source_xa: &mut Source,
34        source_xe: &mut Source,
35        sigma: f64,
36        scratch: &mut Scratch<B>,
37    ) where
38        DLwe: DataRef,
39        DGlwe: DataRef,
40        Module<B>: GGLWESwitchingKeyEncryptSkFamily<B>
41            + VecZnxAutomorphismInplace
42            + VecZnxSwithcDegree
43            + VecZnxAddScalarInplace
44            + GLWESecretPreparedModuleFamily<B>,
45        Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretPrepared<B> + TakeVecZnx,
46    {
47        #[cfg(debug_assertions)]
48        {
49            assert!(sk_lwe.n() <= module.n());
50        }
51
52        let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(sk_glwe.n(), 1);
53        sk_lwe_as_glwe.data.zero();
54        sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
55        module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0);
56
57        self.0.encrypt_sk(
58            module,
59            sk_glwe,
60            &sk_lwe_as_glwe,
61            source_xa,
62            source_xe,
63            sigma,
64            scratch1,
65        );
66    }
67}