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