1use poulpy_hal::{
 2    api::{
 3        ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
 4        VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf,
 5        VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
 6        VecZnxSubInplace, VecZnxSwitchRing,
 7    },
 8    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
 9    source::Source,
10};
11
12use crate::layouts::{
13    GGLWEInfos, GLWESecret, GLWESwitchingKey, GLWEToLWESwitchingKey, LWEInfos, LWESecret, Rank, prepared::GLWESecretPrepared,
14};
15
16impl GLWEToLWESwitchingKey<Vec<u8>> {
17    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
18    where
19        A: GGLWEInfos,
20        Module<B>: SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes,
21    {
22        GLWESecretPrepared::bytes_of(module, infos.rank_in())
23            + (GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos) | GLWESecret::bytes_of(infos.n(), infos.rank_in()))
24    }
25}
26
27impl<D: DataMut> GLWEToLWESwitchingKey<D> {
28    #[allow(clippy::too_many_arguments)]
29    pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
30        &mut self,
31        module: &Module<B>,
32        sk_lwe: &LWESecret<DLwe>,
33        sk_glwe: &GLWESecret<DGlwe>,
34        source_xa: &mut Source,
35        source_xe: &mut Source,
36        scratch: &mut Scratch<B>,
37    ) where
38        DLwe: DataRef,
39        DGlwe: DataRef,
40        Module<B>: VecZnxAutomorphismInplace<B>
41            + VecZnxAddScalarInplace
42            + VecZnxDftBytesOf
43            + VecZnxBigNormalize<B>
44            + VecZnxDftApply<B>
45            + SvpApplyDftToDftInplace<B>
46            + VecZnxIdftApplyConsume<B>
47            + VecZnxNormalizeTmpBytes
48            + VecZnxFillUniform
49            + VecZnxSubInplace
50            + VecZnxAddInplace
51            + VecZnxNormalizeInplace<B>
52            + VecZnxAddNormal
53            + VecZnxNormalize<B>
54            + VecZnxSub
55            + SvpPrepare<B>
56            + VecZnxSwitchRing
57            + SvpPPolBytesOf,
58        Scratch<B>: ScratchAvailable,
59    {
60        #[cfg(debug_assertions)]
61        {
62            assert!(sk_lwe.n().0 <= module.n() as u32);
63        }
64
65        let (mut sk_lwe_as_glwe, scratch_1) = scratch.take_glwe_secret(sk_glwe.n(), Rank(1));
66        sk_lwe_as_glwe.data.zero();
67        sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
68        module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_1);
69
70        self.0.encrypt_sk(
71            module,
72            sk_glwe,
73            &sk_lwe_as_glwe,
74            source_xa,
75            source_xe,
76            scratch_1,
77        );
78    }
79}