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