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