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