1use poulpy_hal::{
2 api::{
3 ScratchAvailable, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes,
5 VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace,
6 VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, Scratch},
9 source::Source,
10};
11
12use crate::{
13 TakeGLWESecret, TakeGLWESecretPrepared,
14 layouts::{GGLWEAutomorphismKey, GGLWESwitchingKey, GLWESecret},
15};
16
17impl GGLWEAutomorphismKey<Vec<u8>> {
18 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
19 where
20 Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
21 {
22 GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank, rank) + GLWESecret::bytes_of(n, rank)
23 }
24
25 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, _n: usize, _basek: usize, _k: usize, _rank: usize) -> usize {
26 GGLWESwitchingKey::encrypt_pk_scratch_space(module, _n, _basek, _k, _rank, _rank)
27 }
28}
29
30impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
31 #[allow(clippy::too_many_arguments)]
32 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
33 &mut self,
34 module: &Module<B>,
35 p: i64,
36 sk: &GLWESecret<DataSk>,
37 source_xa: &mut Source,
38 source_xe: &mut Source,
39 sigma: f64,
40 scratch: &mut Scratch<B>,
41 ) where
42 Module<B>: VecZnxAddScalarInplace
43 + VecZnxDftAllocBytes
44 + VecZnxBigNormalize<B>
45 + VecZnxDftFromVecZnx<B>
46 + SvpApplyInplace<B>
47 + VecZnxDftToVecZnxBigConsume<B>
48 + VecZnxNormalizeTmpBytes
49 + VecZnxFillUniform
50 + VecZnxSubABInplace
51 + VecZnxAddInplace
52 + VecZnxNormalizeInplace<B>
53 + VecZnxAddNormal
54 + VecZnxNormalize<B>
55 + VecZnxSub
56 + SvpPrepare<B>
57 + VecZnxSwithcDegree
58 + SvpPPolAllocBytes
59 + VecZnxAutomorphism,
60 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
61 {
62 #[cfg(debug_assertions)]
63 {
64 use crate::layouts::Infos;
65
66 assert_eq!(self.n(), sk.n());
67 assert_eq!(self.rank_out(), self.rank_in());
68 assert_eq!(sk.rank(), self.rank());
69 assert!(
70 scratch.available()
71 >= GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank()),
72 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
73 scratch.available(),
74 self.rank(),
75 self.size(),
76 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
77 )
78 }
79
80 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
81
82 {
83 (0..self.rank()).for_each(|i| {
84 module.vec_znx_automorphism(
85 module.galois_element_inv(p),
86 &mut sk_out.data.as_vec_znx_mut(),
87 i,
88 &sk.data.as_vec_znx(),
89 i,
90 );
91 });
92 }
93
94 self.key
95 .encrypt_sk(module, sk, &sk_out, source_xa, source_xe, sigma, scratch_1);
96
97 self.p = p;
98 }
99}