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