1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchTakeBasic, SvpApplyDftToDft, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf,
  4        VecZnxIdftApplyTmpA,
  5    },
  6    layouts::{Backend, DataMut, Module, Scratch},
  7    source::Source,
  8};
  9
 10use crate::{
 11    ScratchTakeCore,
 12    encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
 13    layouts::{
 14        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, GetDist, LWEInfos, Rank, TensorKey, TensorKeyToMut,
 15        prepared::{GLWESecretPrepare, GLWESecretPrepared, GLWESecretPreparedAlloc},
 16    },
 17};
 18
 19impl TensorKey<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: TensorKeyEncryptSk<BE>,
 24    {
 25        module.tensor_key_encrypt_sk_tmp_bytes(infos)
 26    }
 27}
 28
 29impl<DataSelf: DataMut> TensorKey<DataSelf> {
 30    pub fn encrypt_sk<M, S, BE: Backend>(
 31        &mut self,
 32        module: &M,
 33        sk: &S,
 34        source_xa: &mut Source,
 35        source_xe: &mut Source,
 36        scratch: &mut Scratch<BE>,
 37    ) where
 38        M: TensorKeyEncryptSk<BE>,
 39        S: GLWESecretToRef + GetDist,
 40        Scratch<BE>: ScratchTakeCore<BE>,
 41    {
 42        module.tensor_key_encrypt_sk(self, sk, source_xa, source_xe, scratch);
 43    }
 44}
 45
 46pub trait TensorKeyEncryptSk<BE: Backend> {
 47    fn tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 48    where
 49        A: GGLWEInfos;
 50
 51    fn tensor_key_encrypt_sk<R, S>(
 52        &self,
 53        res: &mut R,
 54        sk: &S,
 55        source_xa: &mut Source,
 56        source_xe: &mut Source,
 57        scratch: &mut Scratch<BE>,
 58    ) where
 59        R: TensorKeyToMut,
 60        S: GLWESecretToRef + GetDist;
 61}
 62
 63impl<BE: Backend> TensorKeyEncryptSk<BE> for Module<BE>
 64where
 65    Self: ModuleN
 66        + GLWESwitchingKeyEncryptSk<BE>
 67        + VecZnxDftBytesOf
 68        + VecZnxBigBytesOf
 69        + GLWESecretPreparedAlloc<BE>
 70        + GLWESecretPrepare<BE>
 71        + VecZnxDftApply<BE>
 72        + SvpApplyDftToDft<BE>
 73        + VecZnxIdftApplyTmpA<BE>
 74        + VecZnxBigNormalize<BE>,
 75    Scratch<BE>: ScratchTakeCore<BE>,
 76{
 77    fn tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 78    where
 79        A: GGLWEInfos,
 80    {
 81        GLWESecretPrepared::bytes_of(self, infos.rank_out())
 82            + self.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
 83            + self.bytes_of_vec_znx_big(1, 1)
 84            + self.bytes_of_vec_znx_dft(1, 1)
 85            + GLWESecret::bytes_of(self, Rank(1))
 86            + GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos)
 87    }
 88
 89    fn tensor_key_encrypt_sk<R, S>(
 90        &self,
 91        res: &mut R,
 92        sk: &S,
 93        source_xa: &mut Source,
 94        source_xe: &mut Source,
 95        scratch: &mut Scratch<BE>,
 96    ) where
 97        R: TensorKeyToMut,
 98        S: GLWESecretToRef + GetDist,
 99    {
100        let res: &mut TensorKey<&mut [u8]> = &mut res.to_mut();
101
102        // let n: RingDegree = sk.n();
103        let rank: Rank = res.rank_out();
104
105        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, rank);
106        sk_dft_prep.prepare(self, sk);
107
108        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
109
110        assert_eq!(res.rank_out(), sk.rank());
111        assert_eq!(res.n(), sk.n());
112
113        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, rank.into(), 1);
114
115        (0..rank.into()).for_each(|i| {
116            self.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
117        });
118
119        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(self, 1, 1);
120        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(self, Rank(1));
121        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(self, 1, 1);
122
123        (0..rank.into()).for_each(|i| {
124            (i..rank.into()).for_each(|j| {
125                self.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
126
127                self.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
128                self.vec_znx_big_normalize(
129                    res.base2k().into(),
130                    &mut sk_ij.data.as_vec_znx_mut(),
131                    0,
132                    res.base2k().into(),
133                    &sk_ij_big,
134                    0,
135                    scratch_5,
136                );
137
138                res.at_mut(i, j)
139                    .encrypt_sk(self, &sk_ij, sk, source_xa, source_xe, scratch_5);
140            });
141        })
142    }
143}