1use poulpy_hal::{
  2    api::{ModuleN, ScratchTakeBasic, VecZnxCopy},
  3    layouts::{Backend, DataMut, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEEncryptSk, GetDistribution, ScratchTakeCore,
  9    layouts::{
 10        GGLWEInfos, GGLWEToGGSWKey, GGLWEToGGSWKeyToMut, GLWEInfos, GLWESecret, GLWESecretTensor, GLWESecretTensorFactory,
 11        GLWESecretToRef,
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
 13    },
 14};
 15
 16impl GGLWEToGGSWKey<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: GGLWEToGGSWKeyEncryptSk<BE>,
 21    {
 22        module.gglwe_to_ggsw_key_encrypt_sk_tmp_bytes(infos)
 23    }
 24}
 25
 26impl<DataSelf: DataMut> GGLWEToGGSWKey<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: GGLWEToGGSWKeyEncryptSk<BE>,
 36        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 37        Scratch<BE>: ScratchTakeCore<BE>,
 38    {
 39        module.gglwe_to_ggsw_key_encrypt_sk(self, sk, source_xa, source_xe, scratch);
 40    }
 41}
 42
 43pub trait GGLWEToGGSWKeyEncryptSk<BE: Backend> {
 44    fn gglwe_to_ggsw_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 45    where
 46        A: GGLWEInfos;
 47
 48    fn gglwe_to_ggsw_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: GGLWEToGGSWKeyToMut,
 57        S: GLWESecretToRef + GetDistribution + GLWEInfos;
 58}
 59
 60impl<BE: Backend> GGLWEToGGSWKeyEncryptSk<BE> for Module<BE>
 61where
 62    Self: ModuleN + GGLWEEncryptSk<BE> + GLWESecretTensorFactory<BE> + GLWESecretPreparedFactory<BE> + VecZnxCopy,
 63    Scratch<BE>: ScratchTakeCore<BE>,
 64{
 65    fn gglwe_to_ggsw_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());
 70        let sk_tensor: usize = GLWESecretTensor::bytes_of_from_infos(infos);
 71        let gglwe_encrypt: usize = self.gglwe_encrypt_sk_tmp_bytes(infos);
 72        let sk_ij = GLWESecret::bytes_of(self.n().into(), infos.rank());
 73        (sk_prepared + sk_tensor + sk_ij) + gglwe_encrypt.max(self.glwe_secret_tensor_prepare_tmp_bytes(infos.rank()))
 74    }
 75
 76    fn gglwe_to_ggsw_key_encrypt_sk<R, S>(
 77        &self,
 78        res: &mut R,
 79        sk: &S,
 80        source_xa: &mut Source,
 81        source_xe: &mut Source,
 82        scratch: &mut Scratch<BE>,
 83    ) where
 84        R: GGLWEToGGSWKeyToMut,
 85        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 86    {
 87        let res: &mut GGLWEToGGSWKey<&mut [u8]> = &mut res.to_mut();
 88
 89        let rank: usize = res.rank_out().as_usize();
 90
 91        let (mut sk_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, res.rank());
 92        let (mut sk_tensor, scratch_2) = scratch_1.take_glwe_secret_tensor(self.n().into(), res.rank());
 93        sk_prepared.prepare(self, sk);
 94        sk_tensor.prepare(self, sk, scratch_2);
 95
 96        let (mut sk_ij, scratch_3) = scratch_2.take_scalar_znx(self.n(), rank);
 97
 98        for i in 0..rank {
 99            for j in 0..rank {
100                self.vec_znx_copy(
101                    &mut sk_ij.as_vec_znx_mut(),
102                    j,
103                    &sk_tensor.at(i, j).as_vec_znx(),
104                    0,
105                );
106            }
107
108            res.at_mut(i)
109                .encrypt_sk(self, &sk_ij, &sk_prepared, source_xa, source_xe, scratch_3);
110        }
111    }
112}