1use poulpy_hal::{
2 api::{
3 ScratchAvailable, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxBigNormalize,
5 VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform, VecZnxNormalize,
6 VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree, ZnxView, ZnxViewMut,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, Scratch},
9 source::Source,
10};
11
12use crate::{
13 TakeGLWESecret, TakeGLWESecretPrepared,
14 layouts::{GGLWESwitchingKey, GLWESecret, Infos, LWESecret, LWESwitchingKey, prepared::GLWESecretPrepared},
15};
16
17impl LWESwitchingKey<Vec<u8>> {
18 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
19 where
20 Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
21 {
22 GLWESecret::bytes_of(n, 1)
23 + GLWESecretPrepared::bytes_of(module, n, 1)
24 + GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, 1, 1)
25 }
26}
27
28impl<D: DataMut> LWESwitchingKey<D> {
29 #[allow(clippy::too_many_arguments)]
30 pub fn encrypt_sk<DIn, DOut, B: Backend>(
31 &mut self,
32 module: &Module<B>,
33 sk_lwe_in: &LWESecret<DIn>,
34 sk_lwe_out: &LWESecret<DOut>,
35 source_xa: &mut Source,
36 source_xe: &mut Source,
37 sigma: f64,
38 scratch: &mut Scratch<B>,
39 ) where
40 DIn: DataRef,
41 DOut: DataRef,
42 Module<B>: VecZnxAutomorphismInplace
43 + VecZnxAddScalarInplace
44 + VecZnxDftAllocBytes
45 + VecZnxBigNormalize<B>
46 + VecZnxDftFromVecZnx<B>
47 + SvpApplyInplace<B>
48 + VecZnxDftToVecZnxBigConsume<B>
49 + VecZnxNormalizeTmpBytes
50 + VecZnxFillUniform
51 + VecZnxSubABInplace
52 + VecZnxAddInplace
53 + VecZnxNormalizeInplace<B>
54 + VecZnxAddNormal
55 + VecZnxNormalize<B>
56 + VecZnxSub
57 + SvpPrepare<B>
58 + VecZnxSwithcDegree
59 + SvpPPolAllocBytes,
60 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
61 {
62 #[cfg(debug_assertions)]
63 {
64 assert!(sk_lwe_in.n() <= self.n());
65 assert!(sk_lwe_out.n() <= self.n());
66 assert!(self.n() <= module.n());
67 }
68
69 let (mut sk_in_glwe, scratch1) = scratch.take_glwe_secret(self.n(), 1);
70 let (mut sk_out_glwe, scratch2) = scratch1.take_glwe_secret(self.n(), 1);
71
72 sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
73 sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
74 module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0);
75
76 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
77 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
78 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0);
79
80 self.0.encrypt_sk(
81 module,
82 &sk_in_glwe,
83 &sk_out_glwe,
84 source_xa,
85 source_xe,
86 sigma,
87 scratch2,
88 );
89 }
90}