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