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