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, Infos, LWESecret, LWESwitchingKey, prepared::GLWESecretPrepared},
13};
14
15use crate::trait_families::{GGLWEEncryptSkFamily, GGLWESwitchingKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
16
17impl LWESwitchingKey<Vec<u8>> {
18    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
19    where
20        Module<B>: GGLWEEncryptSkFamily<B> + GLWESecretPreparedModuleFamily<B>,
21    {
22        GLWESecret::bytes_of(n, 1)
23            + GLWESecretPrepared::bytes_of(module, n, 1)
24            + GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, 1, 1)
25    }
26}
27
28impl<D: DataMut> LWESwitchingKey<D> {
29    pub fn encrypt_sk<DIn, DOut, B: Backend>(
30        &mut self,
31        module: &Module<B>,
32        sk_lwe_in: &LWESecret<DIn>,
33        sk_lwe_out: &LWESecret<DOut>,
34        source_xa: &mut Source,
35        source_xe: &mut Source,
36        sigma: f64,
37        scratch: &mut Scratch<B>,
38    ) where
39        DIn: DataRef,
40        DOut: DataRef,
41        Module<B>: GGLWESwitchingKeyEncryptSkFamily<B>
42            + VecZnxAutomorphismInplace
43            + VecZnxSwithcDegree
44            + VecZnxAddScalarInplace
45            + GLWESecretPreparedModuleFamily<B>,
46        Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretPrepared<B> + TakeVecZnx,
47    {
48        #[cfg(debug_assertions)]
49        {
50            assert!(sk_lwe_in.n() <= self.n());
51            assert!(sk_lwe_out.n() <= self.n());
52            assert!(self.n() <= module.n());
53        }
54
55        let (mut sk_in_glwe, scratch1) = scratch.take_glwe_secret(self.n(), 1);
56        let (mut sk_out_glwe, scratch2) = scratch1.take_glwe_secret(self.n(), 1);
57
58        sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
59        sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
60        module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0);
61
62        sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
63        sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
64        module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0);
65
66        self.0.encrypt_sk(
67            module,
68            &sk_in_glwe,
69            &sk_out_glwe,
70            source_xa,
71            source_xe,
72            sigma,
73            scratch2,
74        );
75    }
76}