1use poulpy_hal::{
 2    api::{
 3        ModuleN, ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAlloc, 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::{
13    ScratchTakeCore,
14    layouts::{GGLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, LWESecret, LWEToGLWESwitchingKey, Rank},
15};
16
17impl LWEToGLWESwitchingKey<Vec<u8>> {
18    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
19    where
20        A: GGLWEInfos,
21        Module<B>: ModuleN + SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolAlloc<B>,
22    {
23        debug_assert_eq!(
24            infos.rank_in(),
25            Rank(1),
26            "rank_in != 1 is not supported for LWEToGLWESwitchingKey"
27        );
28        GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos)
29            + GLWESecret::bytes_of(module, 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>: ModuleN
47            + VecZnxAutomorphismInplace<B>
48            + VecZnxAddScalarInplace
49            + VecZnxDftBytesOf
50            + VecZnxBigNormalize<B>
51            + VecZnxDftApply<B>
52            + SvpApplyDftToDftInplace<B>
53            + VecZnxIdftApplyConsume<B>
54            + VecZnxNormalizeTmpBytes
55            + VecZnxFillUniform
56            + VecZnxSubInplace
57            + VecZnxAddInplace
58            + VecZnxNormalizeInplace<B>
59            + VecZnxAddNormal
60            + VecZnxNormalize<B>
61            + VecZnxSub
62            + SvpPrepare<B>
63            + VecZnxSwitchRing
64            + SvpPPolBytesOf
65            + SvpPPolAlloc<B>,
66        Scratch<B>: ScratchAvailable + ScratchTakeCore<B>,
67    {
68        #[cfg(debug_assertions)]
69        {
70            use crate::layouts::LWEInfos;
71
72            assert!(sk_lwe.n().0 <= module.n() as u32);
73        }
74
75        let (mut sk_lwe_as_glwe, scratch_1) = scratch.take_glwe_secret(module, Rank(1));
76        sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
77        sk_lwe_as_glwe.data.at_mut(0, 0)[sk_lwe.n().into()..].fill(0);
78        module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_1);
79
80        self.0.encrypt_sk(
81            module,
82            &sk_lwe_as_glwe,
83            sk_glwe,
84            source_xa,
85            source_xe,
86            scratch_1,
87        );
88    }
89}