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