1use poulpy_hal::{
2 api::{
3 ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes,
5 VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace,
6 VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwitchRing,
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>, basek: usize, k: usize, rank: usize) -> usize
19 where
20 Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
21 {
22 GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank) + GLWESecret::bytes_of(module.n(), rank)
23 }
24
25 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, _basek: usize, _k: usize, _rank: usize) -> usize {
26 GGLWESwitchingKey::encrypt_pk_scratch_space(module, _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 scratch: &mut Scratch<B>,
40 ) where
41 Module<B>: VecZnxAddScalarInplace
42 + VecZnxDftAllocBytes
43 + VecZnxBigNormalize<B>
44 + VecZnxDftApply<B>
45 + SvpApplyDftToDftInplace<B>
46 + VecZnxIdftApplyConsume<B>
47 + VecZnxNormalizeTmpBytes
48 + VecZnxFillUniform
49 + VecZnxSubABInplace
50 + VecZnxAddInplace
51 + VecZnxNormalizeInplace<B>
52 + VecZnxAddNormal
53 + VecZnxNormalize<B>
54 + VecZnxSub
55 + SvpPrepare<B>
56 + VecZnxSwitchRing
57 + SvpPPolAllocBytes
58 + VecZnxAutomorphism,
59 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
60 {
61 #[cfg(debug_assertions)]
62 {
63 use crate::layouts::Infos;
64
65 assert_eq!(self.n(), sk.n());
66 assert_eq!(self.rank_out(), self.rank_in());
67 assert_eq!(sk.rank(), self.rank());
68 assert!(
69 scratch.available()
70 >= GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
71 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
72 scratch.available(),
73 self.rank(),
74 self.size(),
75 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
76 )
77 }
78
79 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
80
81 {
82 (0..self.rank()).for_each(|i| {
83 module.vec_znx_automorphism(
84 module.galois_element_inv(p),
85 &mut sk_out.data.as_vec_znx_mut(),
86 i,
87 &sk.data.as_vec_znx(),
88 i,
89 );
90 });
91 }
92
93 self.key
94 .encrypt_sk(module, sk, &sk_out, source_xa, source_xe, scratch_1);
95
96 self.p = p;
97 }
98}