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