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