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