1use poulpy_hal::{
  2    api::{
  3        ModuleN, SvpApplyDftToDftInplace, SvpPPolAlloc, 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::{
 13    ScratchTakeCore,
 14    layouts::{
 15    GGLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, LWESecret, LWESwitchingKey, Rank,
 16    prepared::GLWESecretPrepared,
 17    },
 18};
 19
 20impl LWESwitchingKey<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        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::bytes_of(module, Rank(1))
 42            + GLWESecretPrepared::bytes_of(module, Rank(1))
 43            + GLWESwitchingKey::encrypt_sk_tmp_bytes(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>: ModuleN
 61            + VecZnxAutomorphismInplace<B>
 62            + VecZnxAddScalarInplace
 63            + VecZnxDftBytesOf
 64            + VecZnxBigNormalize<B>
 65            + VecZnxDftApply<B>
 66            + SvpApplyDftToDftInplace<B>
 67            + VecZnxIdftApplyConsume<B>
 68            + VecZnxNormalizeTmpBytes
 69            + VecZnxFillUniform
 70            + VecZnxSubInplace
 71            + VecZnxAddInplace
 72            + VecZnxNormalizeInplace<B>
 73            + VecZnxAddNormal
 74            + VecZnxNormalize<B>
 75            + VecZnxSub
 76            + SvpPrepare<B>
 77            + VecZnxSwitchRing
 78            + SvpPPolBytesOf
 79            + SvpPPolAlloc<B>,
 80        Scratch<B>: ScratchTakeCore<B>,
 81    {
 82        #[cfg(debug_assertions)]
 83        {
 84            assert!(sk_lwe_in.n().0 <= self.n().0);
 85            assert!(sk_lwe_out.n().0 <= self.n().0);
 86            assert!(self.n().0 <= module.n() as u32);
 87        }
 88
 89        let (mut sk_in_glwe, scratch_1) = scratch.take_glwe_secret(module, Rank(1));
 90        let (mut sk_out_glwe, scratch_2) = scratch_1.take_glwe_secret(module, Rank(1));
 91
 92        sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n().into()].copy_from_slice(sk_lwe_out.data.at(0, 0));
 93        sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n().into()..].fill(0);
 94        module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 95
 96        sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n().into()].copy_from_slice(sk_lwe_in.data.at(0, 0));
 97        sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n().into()..].fill(0);
 98        module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0, scratch_2);
 99
100        self.0.encrypt_sk(
101            module,
102            &sk_in_glwe,
103            &sk_out_glwe,
104            source_xa,
105            source_xe,
106            scratch_2,
107        );
108    }
109}