1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchTakeBasic, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
  4        VecZnxAddScalarInplace, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
  5        VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
  6        VecZnxSubInplace, VecZnxSwitchRing,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9    oep::VecZnxBigAllocBytesImpl,
 10    source::Source,
 11};
 12
 13use crate::{
 14    ScratchTakeCore,
 15    layouts::{
 16        GetDist, GGLWEInfos, GLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, Rank, TensorKey, prepared::GLWESecretPrepared,
 17    },
 18};
 19
 20impl TensorKey<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 + VecZnxBigBytesOf,
 25    {
 26        GLWESecretPrepared::bytes_of(module, infos.rank_out())
 27            + module.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
 28            + module.bytes_of_vec_znx_big(1, 1)
 29            + module.bytes_of_vec_znx_dft(1, 1)
 30            + GLWESecret::bytes_of(module, Rank(1))
 31            + GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos)
 32    }
 33}
 34
 35impl<DataSelf: DataMut> TensorKey<DataSelf> {
 36    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
 37        &mut self,
 38        module: &Module<B>,
 39        sk: &GLWESecret<DataSk>,
 40        source_xa: &mut Source,
 41        source_xe: &mut Source,
 42        scratch: &mut Scratch<B>,
 43    ) where
 44        GLWESecret<DataSk>: GetDist,
 45        Module<B>: ModuleN
 46            + SvpApplyDftToDft<B>
 47            + VecZnxIdftApplyTmpA<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            + VecZnxBigAllocBytesImpl<B>
 66            + VecZnxBigBytesOf
 67            + SvpPPolAlloc<B>,
 68        Scratch<B>: ScratchTakeBasic + ScratchTakeCore<B>,
 69    {
 70        #[cfg(debug_assertions)]
 71        {
 72            assert_eq!(self.rank_out(), sk.rank());
 73            assert_eq!(self.n(), sk.n());
 74        }
 75
 76        // let n: RingDegree = sk.n();
 77        let rank: Rank = self.rank_out();
 78
 79        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(module, rank);
 80        sk_dft_prep.prepare(module, sk);
 81
 82        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(module, rank.into(), 1);
 83
 84        (0..rank.into()).for_each(|i| {
 85            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 86        });
 87
 88        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(module, 1, 1);
 89        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(module, Rank(1));
 90        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(module, 1, 1);
 91
 92        (0..rank.into()).for_each(|i| {
 93            (i..rank.into()).for_each(|j| {
 94                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
 95
 96                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 97                module.vec_znx_big_normalize(
 98                    self.base2k().into(),
 99                    &mut sk_ij.data.as_vec_znx_mut(),
100                    0,
101                    self.base2k().into(),
102                    &sk_ij_big,
103                    0,
104                    scratch_5,
105                );
106
107                self.at_mut(i, j)
108                    .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, scratch_5);
109            });
110        })
111    }
112}