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