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