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