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