1use backend::hal::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
 3    layouts::{Backend, DataMut, DataRef, Module, ScratchOwned},
 4    oep::{ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxDftImpl, TakeVecZnxImpl},
 5};
 6use sampling::source::Source;
 7
 8use crate::{
 9    dist::Distribution,
10    layouts::{GLWECiphertext, GLWEPublicKey, Infos, prepared::GLWESecretExec},
11};
12
13use crate::trait_families::GLWEEncryptSkFamily;
14
15impl<D: DataMut> GLWEPublicKey<D> {
16    pub fn generate_from_sk<S: DataRef, B: Backend>(
17        &mut self,
18        module: &Module<B>,
19        sk: &GLWESecretExec<S, B>,
20        source_xa: &mut Source,
21        source_xe: &mut Source,
22        sigma: f64,
23    ) where
24        Module<B>: GLWEEncryptSkFamily<B>,
25        B: ScratchOwnedAllocImpl<B>
26            + ScratchOwnedBorrowImpl<B>
27            + TakeVecZnxDftImpl<B>
28            + ScratchAvailableImpl<B>
29            + TakeVecZnxImpl<B>,
30    {
31        #[cfg(debug_assertions)]
32        {
33            assert_eq!(self.n(), sk.n());
34
35            match sk.dist {
36                Distribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
37                _ => {}
38            }
39        }
40
41        // Its ok to allocate scratch space here since pk is usually generated only once.
42        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWECiphertext::encrypt_sk_scratch_space(
43            module,
44            self.n(),
45            self.basek(),
46            self.k(),
47        ));
48
49        let mut tmp: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(self.n(), self.basek(), self.k(), self.rank());
50        tmp.encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch.borrow());
51        self.dist = sk.dist;
52    }
53}