1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxBig,
  4        TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAllocBytes, VecZnxBigNormalize,
  5        VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA, VecZnxFillUniform,
  6        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    TakeGLWESecret, TakeGLWESecretPrepared,
 14    layouts::{
 15        GGLWESwitchingKey, GGLWETensorKey, GLWESecret, Infos,
 16        prepared::{GLWESecretPrepared, Prepare},
 17    },
 18};
 19
 20impl GGLWETensorKey<Vec<u8>> {
 21    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
 22    where
 23        Module<B>:
 24            SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigAllocBytes,
 25    {
 26        GLWESecretPrepared::bytes_of(module, n, rank)
 27            + module.vec_znx_dft_alloc_bytes(n, rank, 1)
 28            + module.vec_znx_big_alloc_bytes(n, 1, 1)
 29            + module.vec_znx_dft_alloc_bytes(n, 1, 1)
 30            + GLWESecret::bytes_of(n, 1)
 31            + GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank, rank)
 32    }
 33}
 34
 35impl<DataSelf: DataMut> GGLWETensorKey<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        Module<B>: SvpApply<B>
 45            + VecZnxDftToVecZnxBigTmpA<B>
 46            + VecZnxAddScalarInplace
 47            + VecZnxDftAllocBytes
 48            + VecZnxBigNormalize<B>
 49            + VecZnxDftFromVecZnx<B>
 50            + SvpApplyInplace<B>
 51            + VecZnxDftToVecZnxBigConsume<B>
 52            + VecZnxNormalizeTmpBytes
 53            + VecZnxFillUniform
 54            + VecZnxSubABInplace
 55            + VecZnxAddInplace
 56            + VecZnxNormalizeInplace<B>
 57            + VecZnxAddNormal
 58            + VecZnxNormalize<B>
 59            + VecZnxSub
 60            + SvpPrepare<B>
 61            + VecZnxSwithcDegree
 62            + SvpPPolAllocBytes,
 63        Scratch<B>:
 64            TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B> + TakeVecZnxBig<B>,
 65    {
 66        #[cfg(debug_assertions)]
 67        {
 68            assert_eq!(self.rank(), sk.rank());
 69            assert_eq!(self.n(), sk.n());
 70        }
 71
 72        let n: usize = sk.n();
 73
 74        let rank: usize = self.rank();
 75
 76        let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_prepared(n, rank);
 77        sk_dft_prep.prepare(module, sk, scratch1);
 78
 79        let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(n, rank, 1);
 80
 81        (0..rank).for_each(|i| {
 82            module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 83        });
 84
 85        let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(n, 1, 1);
 86        let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(n, 1);
 87        let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(n, 1, 1);
 88
 89        (0..rank).for_each(|i| {
 90            (i..rank).for_each(|j| {
 91                module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
 92
 93                module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 94                module.vec_znx_big_normalize(
 95                    self.basek(),
 96                    &mut sk_ij.data.as_vec_znx_mut(),
 97                    0,
 98                    &sk_ij_big,
 99                    0,
100                    scratch5,
101                );
102
103                self.at_mut(i, j)
104                    .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, scratch5);
105            });
106        })
107    }
108}