1use backend::{Backend, FFT64, Module, ScratchOwned, VecZnxDft};
 2use sampling::source::Source;
 3
 4use crate::{FourierGLWECiphertext, GLWESecret, Infos, keys::SecretDistribution};
 5
 6pub struct GLWEPublicKey<D, B: Backend> {
 7    pub(crate) data: FourierGLWECiphertext<D, B>,
 8    pub(crate) dist: SecretDistribution,
 9}
10
11impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
12    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
13        Self {
14            data: FourierGLWECiphertext::alloc(module, basek, k, rank),
15            dist: SecretDistribution::NONE,
16        }
17    }
18
19    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
20        FourierGLWECiphertext::<Vec<u8>, B>::bytes_of(module, basek, k, rank)
21    }
22}
23
24impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
25    type Inner = VecZnxDft<T, B>;
26
27    fn inner(&self) -> &Self::Inner {
28        &self.data.data
29    }
30
31    fn basek(&self) -> usize {
32        self.data.basek
33    }
34
35    fn k(&self) -> usize {
36        self.data.k
37    }
38}
39
40impl<T, B: Backend> GLWEPublicKey<T, B> {
41    pub fn rank(&self) -> usize {
42        self.cols() - 1
43    }
44}
45
46impl<C: AsRef<[u8]> + AsMut<[u8]>> GLWEPublicKey<C, FFT64> {
47    pub fn generate_from_sk<S: AsRef<[u8]>>(
48        &mut self,
49        module: &Module<FFT64>,
50        sk: &GLWESecret<S, FFT64>,
51        source_xa: &mut Source,
52        source_xe: &mut Source,
53        sigma: f64,
54    ) {
55        #[cfg(debug_assertions)]
56        {
57            match sk.dist {
58                SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
59                _ => {}
60            }
61        }
62
63        // Its ok to allocate scratch space here since pk is usually generated only once.
64        let mut scratch: ScratchOwned = ScratchOwned::new(FourierGLWECiphertext::encrypt_sk_scratch_space(
65            module,
66            self.basek(),
67            self.k(),
68            self.rank(),
69        ));
70
71        self.data
72            .encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch.borrow());
73        self.dist = sk.dist;
74    }
75}