1use poulpy_hal::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
 3    layouts::{Backend, DataMut, Module, Scratch, ScratchOwned},
 4    source::Source,
 5};
 6
 7use crate::{
 8    Distribution, GetDistribution, GetDistributionMut, ScratchTakeCore,
 9    encryption::glwe_ct::GLWEEncryptSk,
10    layouts::{
11        GLWE, GLWEPublicKey, GLWEPublicKeyToMut, LWEInfos,
12        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
13    },
14};
15
16impl<D: DataMut> GLWEPublicKey<D> {
17    pub fn generate<S, M, BE: Backend>(&mut self, module: &M, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
18    where
19        S: GLWESecretPreparedToRef<BE> + GetDistribution,
20        M: GLWEPublicKeyGenerate<BE>,
21    {
22        module.glwe_public_key_generate(self, sk, source_xa, source_xe);
23    }
24}
25
26pub trait GLWEPublicKeyGenerate<BE: Backend> {
27    fn glwe_public_key_generate<R, S>(&self, res: &mut R, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
28    where
29        R: GLWEPublicKeyToMut + GetDistributionMut,
30        S: GLWESecretPreparedToRef<BE> + GetDistribution;
31}
32
33impl<BE: Backend> GLWEPublicKeyGenerate<BE> for Module<BE>
34where
35    Self: GLWEEncryptSk<BE>,
36    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
37    Scratch<BE>: ScratchTakeCore<BE>,
38{
39    fn glwe_public_key_generate<R, S>(&self, res: &mut R, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
40    where
41        R: GLWEPublicKeyToMut + GetDistributionMut,
42        S: GLWESecretPreparedToRef<BE> + GetDistribution,
43    {
44        {
45            let res: &mut GLWEPublicKey<&mut [u8]> = &mut res.to_mut();
46            let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
47
48            assert_eq!(res.n(), self.n() as u32);
49            assert_eq!(sk.n(), self.n() as u32);
50
51            if sk.dist == Distribution::NONE {
52                panic!("invalid sk: SecretDistribution::NONE")
53            }
54
55            // Its ok to allocate scratch space here since pk is usually generated only once.
56            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(self.glwe_encrypt_sk_tmp_bytes(res));
57
58            let mut tmp: GLWE<Vec<u8>> = GLWE::alloc_from_infos(res);
59
60            tmp.encrypt_zero_sk(self, sk, source_xa, source_xe, scratch.borrow());
61        }
62        *res.dist_mut() = *sk.dist();
63    }
64}