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,
 12    encryption::gglwe_tsk::TensorKeyEncryptSk,
 13    layouts::{
 14        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretPrepared, GLWESecretPreparedAlloc, GLWESecretToRef, LWEInfos, Rank,
 15        TensorKeyCompressedAtMut, compressed::TensorKeyCompressed,
 16    },
 17};
 18
 19impl TensorKeyCompressed<Vec<u8>> {
 20    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 21    where
 22        A: GGLWEInfos,
 23        M: GGLWETensorKeyCompressedEncryptSk<BE>,
 24    {
 25        module.tensor_key_compressed_encrypt_sk_tmp_bytes(infos)
 26    }
 27}
 28
 29impl<DataSelf: DataMut> TensorKeyCompressed<DataSelf> {
 30    pub fn encrypt_sk<S, M, BE: Backend>(
 31        &mut self,
 32        module: &M,
 33        sk: &S,
 34        seed_xa: [u8; 32],
 35        source_xe: &mut Source,
 36        scratch: &mut Scratch<BE>,
 37    ) where
 38        S: GLWESecretToRef + GetDistribution,
 39        M: GGLWETensorKeyCompressedEncryptSk<BE>,
 40    {
 41        module.tensor_key_compressed_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
 42    }
 43}
 44
 45pub trait GGLWETensorKeyCompressedEncryptSk<BE: Backend> {
 46    fn tensor_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 47    where
 48        A: GGLWEInfos;
 49
 50    fn tensor_key_compressed_encrypt_sk<R, S, D>(
 51        &self,
 52        res: &mut R,
 53        sk: &S,
 54        seed_xa: [u8; 32],
 55        source_xe: &mut Source,
 56        scratch: &mut Scratch<BE>,
 57    ) where
 58        D: DataMut,
 59        R: TensorKeyCompressedAtMut<D> + GGLWEInfos,
 60        S: GLWESecretToRef + GetDistribution;
 61}
 62
 63impl<BE: Backend> GGLWETensorKeyCompressedEncryptSk<BE> for Module<BE>
 64where
 65    Self: ModuleN
 66        + GGLWECompressedEncryptSk<BE>
 67        + TensorKeyEncryptSk<BE>
 68        + VecZnxDftApply<BE>
 69        + SvpApplyDftToDft<BE>
 70        + VecZnxIdftApplyTmpA<BE>
 71        + VecZnxBigNormalize<BE>
 72        + SvpPrepare<BE>
 73        + SvpPPolBytesOf
 74        + VecZnxDftBytesOf
 75        + VecZnxBigBytesOf
 76        + GLWESecretPreparedAlloc<BE>,
 77    Scratch<BE>: ScratchTakeBasic + ScratchTakeCore<BE>,
 78{
 79    fn tensor_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 80    where
 81        A: GGLWEInfos,
 82    {
 83        GLWESecretPrepared::bytes_of(self, infos.rank_out())
 84            + self.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
 85            + self.bytes_of_vec_znx_big(1, 1)
 86            + self.bytes_of_vec_znx_dft(1, 1)
 87            + GLWESecret::bytes_of(self.n().into(), Rank(1))
 88            + self.gglwe_compressed_encrypt_sk_tmp_bytes(infos)
 89    }
 90
 91    fn tensor_key_compressed_encrypt_sk<R, S, D>(
 92        &self,
 93        res: &mut R,
 94        sk: &S,
 95        seed_xa: [u8; 32],
 96        source_xe: &mut Source,
 97        scratch: &mut Scratch<BE>,
 98    ) where
 99        D: DataMut,
100        R: GGLWEInfos + TensorKeyCompressedAtMut<D>,
101        S: GLWESecretToRef + GetDistribution,
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.gglwe_compressed_encrypt_sk(
147                    res.at_mut(i, j),
148                    &sk_ij.data,
149                    &sk_dft_prep,
150                    seed_xa_tmp,
151                    source_xe,
152                    scratch_5,
153                );
154            }
155        }
156    }
157}