1use poulpy_hal::{
  2    api::ScratchTakeBasic,
  3    layouts::{Backend, DataMut, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWECompressedEncryptSk, GetDistribution, ScratchTakeCore,
  9    layouts::{
 10        GGLWECompressedSeedMut, GGLWECompressedToMut, GGLWEInfos, GGLWELayout, GLWEInfos, GLWESecretPrepared,
 11        GLWESecretPreparedFactory, GLWESecretTensor, GLWESecretTensorFactory, GLWESecretToRef,
 12        compressed::GLWETensorKeyCompressed,
 13    },
 14};
 15
 16impl GLWETensorKeyCompressed<Vec<u8>> {
 17    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 18    where
 19        A: GGLWEInfos,
 20        M: GLWETensorKeyCompressedEncryptSk<BE>,
 21    {
 22        module.glwe_tensor_key_compressed_encrypt_sk_tmp_bytes(infos)
 23    }
 24}
 25
 26impl<DataSelf: DataMut> GLWETensorKeyCompressed<DataSelf> {
 27    pub fn encrypt_sk<S, M, BE: Backend>(
 28        &mut self,
 29        module: &M,
 30        sk: &S,
 31        seed_xa: [u8; 32],
 32        source_xe: &mut Source,
 33        scratch: &mut Scratch<BE>,
 34    ) where
 35        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 36        M: GLWETensorKeyCompressedEncryptSk<BE>,
 37    {
 38        module.glwe_tensor_key_compressed_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
 39    }
 40}
 41
 42pub trait GLWETensorKeyCompressedEncryptSk<BE: Backend> {
 43    fn glwe_tensor_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 44    where
 45        A: GGLWEInfos;
 46
 47    fn glwe_tensor_key_compressed_encrypt_sk<R, S>(
 48        &self,
 49        res: &mut R,
 50        sk: &S,
 51        seed_xa: [u8; 32],
 52        source_xe: &mut Source,
 53        scratch: &mut Scratch<BE>,
 54    ) where
 55        R: GGLWECompressedToMut + GGLWEInfos + GGLWECompressedSeedMut,
 56        S: GLWESecretToRef + GetDistribution + GLWEInfos;
 57}
 58
 59impl<BE: Backend> GLWETensorKeyCompressedEncryptSk<BE> for Module<BE>
 60where
 61    Self: GGLWECompressedEncryptSk<BE> + GLWESecretPreparedFactory<BE> + GLWESecretTensorFactory<BE>,
 62    Scratch<BE>: ScratchTakeBasic + ScratchTakeCore<BE>,
 63{
 64    fn glwe_tensor_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 65    where
 66        A: GGLWEInfos,
 67    {
 68        let sk_prepared: usize = GLWESecretPrepared::bytes_of(self, infos.rank_out());
 69        let sk_tensor: usize = GLWESecretTensor::bytes_of_from_infos(infos);
 70
 71        let tensor_infos: GGLWELayout = GGLWELayout {
 72            n: infos.n(),
 73            base2k: infos.base2k(),
 74            k: infos.k(),
 75            rank_in: GLWESecretTensor::pairs(infos.rank().into()).into(),
 76            rank_out: infos.rank_out(),
 77            dnum: infos.dnum(),
 78            dsize: infos.dsize(),
 79        };
 80
 81        let gglwe_encrypt: usize = self.gglwe_compressed_encrypt_sk_tmp_bytes(&tensor_infos);
 82
 83        (sk_prepared + sk_tensor) + gglwe_encrypt.max(self.glwe_secret_tensor_prepare_tmp_bytes(infos.rank()))
 84    }
 85
 86    fn glwe_tensor_key_compressed_encrypt_sk<R, S>(
 87        &self,
 88        res: &mut R,
 89        sk: &S,
 90        seed_xa: [u8; 32],
 91        source_xe: &mut Source,
 92        scratch: &mut Scratch<BE>,
 93    ) where
 94        R: GGLWEInfos + GGLWECompressedToMut + GGLWECompressedSeedMut,
 95        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 96    {
 97        assert_eq!(res.rank_out(), sk.rank());
 98        assert_eq!(res.n(), sk.n());
 99
100        let (mut sk_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, res.rank());
101        let (mut sk_tensor, scratch_2) = scratch_1.take_glwe_secret_tensor(self.n().into(), res.rank());
102        sk_prepared.prepare(self, sk);
103        sk_tensor.prepare(self, sk, scratch_2);
104
105        self.gglwe_compressed_encrypt_sk(
106            res,
107            &sk_tensor.data,
108            &sk_prepared,
109            seed_xa,
110            source_xe,
111            scratch_2,
112        );
113    }
114}