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::{
 15        Degree, GGLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, LWESecret, LWESwitchingKey, Rank,
 16        prepared::GLWESecretPrepared,
 17    },
 18};
 19
 20impl LWESwitchingKey<Vec<u8>> {
 21    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 22    where
 23        A: GGLWEInfos,
 24        Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
 25    {
 26        debug_assert_eq!(
 27            infos.dsize().0,
 28            1,
 29            "dsize > 1 is not supported for LWESwitchingKey"
 30        );
 31        debug_assert_eq!(
 32            infos.rank_in().0,
 33            1,
 34            "rank_in > 1 is not supported for LWESwitchingKey"
 35        );
 36        debug_assert_eq!(
 37            infos.rank_out().0,
 38            1,
 39            "rank_out > 1 is not supported for LWESwitchingKey"
 40        );
 41        GLWESecret::alloc_bytes_with(Degree(module.n() as u32), Rank(1))
 42            + GLWESecretPrepared::alloc_bytes(module, Rank(1))
 43            + GLWESwitchingKey::encrypt_sk_scratch_space(module, infos)
 44    }
 45}
 46
 47impl<D: DataMut> LWESwitchingKey<D> {
 48    #[allow(clippy::too_many_arguments)]
 49    pub fn encrypt_sk<DIn, DOut, B: Backend>(
 50        &mut self,
 51        module: &Module<B>,
 52        sk_lwe_in: &LWESecret<DIn>,
 53        sk_lwe_out: &LWESecret<DOut>,
 54        source_xa: &mut Source,
 55        source_xe: &mut Source,
 56        scratch: &mut Scratch<B>,
 57    ) where
 58        DIn: DataRef,
 59        DOut: DataRef,
 60        Module<B>: VecZnxAutomorphismInplace<B>
 61            + VecZnxAddScalarInplace
 62            + VecZnxDftAllocBytes
 63            + VecZnxBigNormalize<B>
 64            + VecZnxDftApply<B>
 65            + SvpApplyDftToDftInplace<B>
 66            + VecZnxIdftApplyConsume<B>
 67            + VecZnxNormalizeTmpBytes
 68            + VecZnxFillUniform
 69            + VecZnxSubInplace
 70            + VecZnxAddInplace
 71            + VecZnxNormalizeInplace<B>
 72            + VecZnxAddNormal
 73            + VecZnxNormalize<B>
 74            + VecZnxSub
 75            + SvpPrepare<B>
 76            + VecZnxSwitchRing
 77            + SvpPPolAllocBytes,
 78        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
 79    {
 80        #[cfg(debug_assertions)]
 81        {
 82            assert!(sk_lwe_in.n().0 <= self.n().0);
 83            assert!(sk_lwe_out.n().0 <= self.n().0);
 84            assert!(self.n().0 <= module.n() as u32);
 85        }
 86
 87        let (mut sk_in_glwe, scratch_1) = scratch.take_glwe_secret(self.n(), Rank(1));
 88        let (mut sk_out_glwe, scratch_2) = scratch_1.take_glwe_secret(self.n(), Rank(1));
 89
 90        sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n().into()].copy_from_slice(sk_lwe_out.data.at(0, 0));
 91        sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n().into()..].fill(0);
 92        module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 93
 94        sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n().into()].copy_from_slice(sk_lwe_in.data.at(0, 0));
 95        sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n().into()..].fill(0);
 96        module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 97
 98        self.0.encrypt_sk(
 99            module,
100            &sk_in_glwe,
101            &sk_out_glwe,
102            source_xa,
103            source_xe,
104            scratch_2,
105        );
106    }
107}