1use poulpy_hal::{
  2    api::{
  3        SvpApplyDftToDft, SvpPPolBytesOf, SvpPrepare, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigBytesOf, VecZnxBigNormalize,
  4        VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyTmpA, VecZnxNormalizeTmpBytes,
  5    },
  6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  7    source::Source,
  8};
  9
 10use crate::{
 11    TakeGLWESecret, TakeGLWESecretPrepared,
 12    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
 13    layouts::{
 14        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, Rank, TensorKey,
 15        compressed::{TensorKeyCompressed, TensorKeyCompressedToMut},
 16    },
 17};
 18
 19impl TensorKeyCompressed<Vec<u8>> {
 20    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 21    where
 22        A: GGLWEInfos,
 23        Module<B>: SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigBytesOf,
 24    {
 25        TensorKey::encrypt_sk_scratch_space(module, infos)
 26    }
 27}
 28
 29pub trait GGLWETensorKeyCompressedEncryptSk<B: Backend> {
 30    fn gglwe_tensor_key_encrypt_sk<R, S>(
 31        &self,
 32        res: &mut R,
 33        sk: &S,
 34        seed_xa: [u8; 32],
 35        source_xe: &mut Source,
 36        scratch: &mut Scratch<B>,
 37    ) where
 38        R: TensorKeyCompressedToMut,
 39        S: GLWESecretToRef;
 40}
 41
 42impl<B: Backend> GGLWETensorKeyCompressedEncryptSk<B> for Module<B>
 43where
 44    Module<B>: GGLWEKeyCompressedEncryptSk<B>
 45        + VecZnxDftApply<B>
 46        + SvpApplyDftToDft<B>
 47        + VecZnxIdftApplyTmpA<B>
 48        + VecZnxBigNormalize<B>
 49        + SvpPrepare<B>,
 50    Scratch<B>: TakeGLWESecretPrepared<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecret,
 51{
 52    fn gglwe_tensor_key_encrypt_sk<R, S>(
 53        &self,
 54        res: &mut R,
 55        sk: &S,
 56        seed_xa: [u8; 32],
 57        source_xe: &mut Source,
 58        scratch: &mut Scratch<B>,
 59    ) where
 60        R: TensorKeyCompressedToMut,
 61        S: GLWESecretToRef,
 62    {
 63        let res: &mut TensorKeyCompressed<&mut [u8]> = &mut res.to_mut();
 64        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 65
 66        #[cfg(debug_assertions)]
 67        {
 68            assert_eq!(res.rank_out(), sk.rank());
 69            assert_eq!(res.n(), sk.n());
 70        }
 71
 72        let n: usize = sk.n().into();
 73        let rank: usize = res.rank_out().into();
 74
 75        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(sk.n(), res.rank_out());
 76        sk_dft_prep.prepare(self, sk, scratch_1);
 77
 78        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(n, rank, 1);
 79
 80        for i in 0..rank {
 81            self.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
 82        }
 83
 84        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(n, 1, 1);
 85        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(sk.n(), Rank(1));
 86        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(n, 1, 1);
 87
 88        let mut source_xa: Source = Source::new(seed_xa);
 89
 90        for i in 0..rank {
 91            for j in i..rank {
 92                self.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
 93
 94                self.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 95                self.vec_znx_big_normalize(
 96                    res.base2k().into(),
 97                    &mut sk_ij.data.as_vec_znx_mut(),
 98                    0,
 99                    res.base2k().into(),
100                    &sk_ij_big,
101                    0,
102                    scratch_5,
103                );
104
105                let (seed_xa_tmp, _) = source_xa.branch();
106
107                self.gglwe_key_compressed_encrypt_sk(
108                    res.at_mut(i, j),
109                    &sk_ij,
110                    sk,
111                    seed_xa_tmp,
112                    source_xe,
113                    scratch_5,
114                );
115            }
116        }
117    }
118}
119
120impl<DataSelf: DataMut> TensorKeyCompressed<DataSelf> {
121    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
122        &mut self,
123        module: &Module<B>,
124        sk: &GLWESecret<DataSk>,
125        seed_xa: [u8; 32],
126        source_xe: &mut Source,
127        scratch: &mut Scratch<B>,
128    ) where
129        Module<B>: GGLWETensorKeyCompressedEncryptSk<B>,
130    {
131        module.gglwe_tensor_key_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
132    }
133}