1use poulpy_hal::{
  2    api::{ModuleN, ScratchTakeBasic, VecZnxCopy},
  3    layouts::{Backend, DataMut, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWECompressedEncryptSk, GetDistribution, ScratchTakeCore,
  9    layouts::{
 10        GGLWEInfos, GGLWEToGGSWKeyCompressed, GGLWEToGGSWKeyCompressedToMut, GLWEInfos, GLWESecret, GLWESecretTensor,
 11        GLWESecretTensorFactory, GLWESecretToRef,
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
 13    },
 14};
 15
 16impl GGLWEToGGSWKeyCompressed<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: GGLWEToGGSWKeyCompressedEncryptSk<BE>,
 21    {
 22        module.gglwe_to_ggsw_key_encrypt_sk_tmp_bytes(infos)
 23    }
 24}
 25
 26impl<DataSelf: DataMut> GGLWEToGGSWKeyCompressed<DataSelf> {
 27    pub fn encrypt_sk<M, S, 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        M: GGLWEToGGSWKeyCompressedEncryptSk<BE>,
 36        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 37        Scratch<BE>: ScratchTakeCore<BE>,
 38    {
 39        module.gglwe_to_ggsw_key_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
 40    }
 41}
 42
 43pub trait GGLWEToGGSWKeyCompressedEncryptSk<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        seed_xa: [u8; 32],
 53        source_xe: &mut Source,
 54        scratch: &mut Scratch<BE>,
 55    ) where
 56        R: GGLWEToGGSWKeyCompressedToMut + GGLWEInfos,
 57        S: GLWESecretToRef + GetDistribution + GLWEInfos;
 58}
 59
 60impl<BE: Backend> GGLWEToGGSWKeyCompressedEncryptSk<BE> for Module<BE>
 61where
 62    Self: ModuleN + GGLWECompressedEncryptSk<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_compressed_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        seed_xa: [u8; 32],
 81        source_xe: &mut Source,
 82        scratch: &mut Scratch<BE>,
 83    ) where
 84        R: GGLWEToGGSWKeyCompressedToMut + GGLWEInfos,
 85        S: GLWESecretToRef + GetDistribution + GLWEInfos,
 86    {
 87        assert_eq!(res.rank(), sk.rank());
 88        assert_eq!(res.n(), sk.n());
 89
 90        let res: &mut GGLWEToGGSWKeyCompressed<&mut [u8]> = &mut res.to_mut();
 91        let rank: usize = res.rank_out().as_usize();
 92
 93        let (mut sk_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, res.rank());
 94        let (mut sk_tensor, scratch_2) = scratch_1.take_glwe_secret_tensor(self.n().into(), res.rank());
 95        sk_prepared.prepare(self, sk);
 96        sk_tensor.prepare(self, sk, scratch_2);
 97
 98        let (mut sk_ij, scratch_3) = scratch_2.take_scalar_znx(self.n(), rank);
 99
100        let mut source_xa = Source::new(seed_xa);
101
102        for i in 0..rank {
103            for j in 0..rank {
104                self.vec_znx_copy(
105                    &mut sk_ij.as_vec_znx_mut(),
106                    j,
107                    &sk_tensor.at(i, j).as_vec_znx(),
108                    0,
109                );
110            }
111
112            let (seed_xa_tmp, _) = source_xa.branch();
113
114            res.at_mut(i).encrypt_sk(
115                self,
116                &sk_ij,
117                &sk_prepared,
118                seed_xa_tmp,
119                source_xe,
120                scratch_3,
121            );
122        }
123    }
124}