1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolBytesOf, SvpPrepare, TakeScalarZnx, TakeVecZnx,
  4        TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigBytesOf,
  5        VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA,
  6        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    TakeGLWESecret, TakeGLWESecretPrepared,
 14    layouts::{
 15        Degree, GGLWEInfos, GLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, Rank, TensorKey, prepared::GLWESecretPrepared,
 16    },
 17};
 18
 19impl TensorKey<Vec<u8>> {
 20    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 21    where
 22        A: GGLWEInfos,
 23        Module<B>: SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigBytesOf,
 24    {
 25        GLWESecretPrepared::bytes_of(module, infos.rank_out())
 26            + module.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
 27            + module.bytes_of_vec_znx_big(1, 1)
 28            + module.bytes_of_vec_znx_dft(1, 1)
 29            + GLWESecret::bytes_of(Degree(module.n() as u32), Rank(1))
 30            + GLWESwitchingKey::encrypt_sk_scratch_space(module, infos)
 31    }
 32}
 33
 34impl<DataSelf: DataMut> TensorKey<DataSelf> {
 35    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
 36        &mut self,
 37        module: &Module<B>,
 38        sk: &GLWESecret<DataSk>,
 39        source_xa: &mut Source,
 40        source_xe: &mut Source,
 41        scratch: &mut Scratch<B>,
 42    ) where
 43        Module<B>: SvpApplyDftToDft<B>
 44            + VecZnxIdftApplyTmpA<B>
 45            + VecZnxAddScalarInplace
 46            + VecZnxDftBytesOf
 47            + VecZnxBigNormalize<B>
 48            + VecZnxDftApply<B>
 49            + SvpApplyDftToDftInplace<B>
 50            + VecZnxIdftApplyConsume<B>
 51            + VecZnxNormalizeTmpBytes
 52            + VecZnxFillUniform
 53            + VecZnxSubInplace
 54            + VecZnxAddInplace
 55            + VecZnxNormalizeInplace<B>
 56            + VecZnxAddNormal
 57            + VecZnxNormalize<B>
 58            + VecZnxSub
 59            + SvpPrepare<B>
 60            + VecZnxSwitchRing
 61            + SvpPPolBytesOf,
 62        Scratch<B>:
 63            TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B> + TakeVecZnxBig<B>,
 64    {
 65        #[cfg(debug_assertions)]
 66        {
 67            assert_eq!(self.rank_out(), sk.rank());
 68            assert_eq!(self.n(), sk.n());
 69        }
 70
 71        let n: Degree = sk.n();
 72        let rank: Rank = self.rank_out();
 73
 74        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(n, rank);
 75        sk_dft_prep.prepare(module, sk, scratch_1);
 76
 77        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(n.into(), rank.into(), 1);
 78
 79        (0..rank.into()).for_each(|i| {
 80            module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 81        });
 82
 83        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(n.into(), 1, 1);
 84        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(n, Rank(1));
 85        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(n.into(), 1, 1);
 86
 87        (0..rank.into()).for_each(|i| {
 88            (i..rank.into()).for_each(|j| {
 89                module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
 90
 91                module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 92                module.vec_znx_big_normalize(
 93                    self.base2k().into(),
 94                    &mut sk_ij.data.as_vec_znx_mut(),
 95                    0,
 96                    self.base2k().into(),
 97                    &sk_ij_big,
 98                    0,
 99                    scratch_5,
100                );
101
102                self.at_mut(i, j)
103                    .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, scratch_5);
104            });
105        })
106    }
107}