1use poulpy_hal::{
  2    api::{
  3        SvpApplyDftToDft, SvpPPolAllocBytes, SvpPrepare, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAllocBytes, VecZnxBigNormalize,
  4        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyTmpA, VecZnxNormalizeTmpBytes,
  5    },
  6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  7    source::Source,
  8};
  9
 10use crate::{
 11    TakeGLWESecret, TakeGLWESecretPrepared,
 12    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
 13    layouts::{
 14        GGLWEInfos, GGLWETensorKey, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, Rank,
 15        compressed::{GGLWETensorKeyCompressed, GGLWETensorKeyCompressedToMut},
 16        prepared::Prepare,
 17    },
 18};
 19
 20impl GGLWETensorKeyCompressed<Vec<u8>> {
 21    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 22    where
 23        A: GGLWEInfos,
 24        Module<B>:
 25            SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigAllocBytes,
 26    {
 27        GGLWETensorKey::encrypt_sk_scratch_space(module, infos)
 28    }
 29}
 30
 31pub trait GGLWETensorKeyCompressedEncryptSk<B: Backend> {
 32    fn gglwe_tensor_key_encrypt_sk<R, S>(
 33        &self,
 34        res: &mut R,
 35        sk: &S,
 36        seed_xa: [u8; 32],
 37        source_xe: &mut Source,
 38        scratch: &mut Scratch<B>,
 39    ) where
 40        R: GGLWETensorKeyCompressedToMut,
 41        S: GLWESecretToRef;
 42}
 43
 44impl<B: Backend> GGLWETensorKeyCompressedEncryptSk<B> for Module<B>
 45where
 46    Module<B>: GGLWEKeyCompressedEncryptSk<B>
 47        + VecZnxDftApply<B>
 48        + SvpApplyDftToDft<B>
 49        + VecZnxIdftApplyTmpA<B>
 50        + VecZnxBigNormalize<B>
 51        + SvpPrepare<B>,
 52    Scratch<B>: TakeGLWESecretPrepared<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecret,
 53{
 54    fn gglwe_tensor_key_encrypt_sk<R, S>(
 55        &self,
 56        res: &mut R,
 57        sk: &S,
 58        seed_xa: [u8; 32],
 59        source_xe: &mut Source,
 60        scratch: &mut Scratch<B>,
 61    ) where
 62        R: GGLWETensorKeyCompressedToMut,
 63        S: GLWESecretToRef,
 64    {
 65        let res: &mut GGLWETensorKeyCompressed<&mut [u8]> = &mut res.to_mut();
 66        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 67
 68        #[cfg(debug_assertions)]
 69        {
 70            assert_eq!(res.rank_out(), sk.rank());
 71            assert_eq!(res.n(), sk.n());
 72        }
 73
 74        let n: usize = sk.n().into();
 75        let rank: usize = res.rank_out().into();
 76
 77        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(sk.n(), res.rank_out());
 78        sk_dft_prep.prepare(self, sk, scratch_1);
 79
 80        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(n, rank, 1);
 81
 82        for i in 0..rank {
 83            self.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 84        }
 85
 86        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(n, 1, 1);
 87        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(sk.n(), Rank(1));
 88        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(n, 1, 1);
 89
 90        let mut source_xa: Source = Source::new(seed_xa);
 91
 92        for i in 0..rank {
 93            for j in i..rank {
 94                self.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
 95
 96                self.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 97                self.vec_znx_big_normalize(
 98                    res.base2k().into(),
 99                    &mut sk_ij.data.as_vec_znx_mut(),
100                    0,
101                    res.base2k().into(),
102                    &sk_ij_big,
103                    0,
104                    scratch_5,
105                );
106
107                let (seed_xa_tmp, _) = source_xa.branch();
108
109                self.gglwe_key_compressed_encrypt_sk(
110                    res.at_mut(i, j),
111                    &sk_ij,
112                    sk,
113                    seed_xa_tmp,
114                    source_xe,
115                    scratch_5,
116                );
117            }
118        }
119    }
120}
121
122impl<DataSelf: DataMut> GGLWETensorKeyCompressed<DataSelf> {
123    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
124        &mut self,
125        module: &Module<B>,
126        sk: &GLWESecret<DataSk>,
127        seed_xa: [u8; 32],
128        source_xe: &mut Source,
129        scratch: &mut Scratch<B>,
130    ) where
131        Module<B>: GGLWETensorKeyCompressedEncryptSk<B>,
132    {
133        module.gglwe_tensor_key_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
134    }
135}