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