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