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