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, ZnxZero},
 9    source::Source,
10};
11
12use crate::{
13    ScratchTakeCore,
14    layouts::{
15        GGLWEInfos, GLWESecret, GLWESwitchingKey, GLWEToLWESwitchingKey, LWEInfos, LWESecret, Rank,
16        prepared::GLWESecretPrepared,
17    },
18};
19
20impl GLWEToLWESwitchingKey<Vec<u8>> {
21    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
22    where
23        A: GGLWEInfos,
24        Module<B>: ModuleN + SvpPPolBytesOf + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes,
25    {
26        GLWESecretPrepared::bytes_of(module, infos.rank_in())
27            + (GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos) | GLWESecret::bytes_of(module, 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>: ModuleN
45            + VecZnxAutomorphismInplace<B>
46            + VecZnxAddScalarInplace
47            + VecZnxDftBytesOf
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            + SvpPPolBytesOf
63            + SvpPPolAlloc<B>,
64        Scratch<B>: ScratchAvailable + ScratchTakeCore<B>,
65    {
66        #[cfg(debug_assertions)]
67        {
68            assert!(sk_lwe.n().0 <= module.n() as u32);
69        }
70
71        let (mut sk_lwe_as_glwe, scratch_1) = scratch.take_glwe_secret(module, Rank(1));
72        sk_lwe_as_glwe.data.zero();
73        sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
74        module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_1);
75
76        self.0.encrypt_sk(
77            module,
78            sk_glwe,
79            &sk_lwe_as_glwe,
80            source_xa,
81            source_xe,
82            scratch_1,
83        );
84    }
85}