1use poulpy_hal::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
 3    layouts::{Backend, DataMut, Module, Scratch, ScratchOwned},
 4    source::Source,
 5};
 6
 7use crate::{
 8    Distribution, 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>,
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,
30        S: GLWESecretPreparedToRef<BE>;
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,
42        S: GLWESecretPreparedToRef<BE>,
43    {
44        let res: &mut GLWEPublicKey<&mut [u8]> = &mut res.to_mut();
45        let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
46
47        assert_eq!(res.n(), self.n() as u32);
48        assert_eq!(sk.n(), self.n() as u32);
49
50        if sk.dist == Distribution::NONE {
51            panic!("invalid sk: SecretDistribution::NONE")
52        }
53
54        // Its ok to allocate scratch space here since pk is usually generated only once.
55        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(self.glwe_encrypt_sk_tmp_bytes(res));
56
57        let mut tmp: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self, res);
58
59        tmp.encrypt_zero_sk(self, sk, source_xa, source_xe, scratch.borrow());
60        res.dist = sk.dist;
61    }
62}