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