poulpy/poulpy-core/src/encryption/gglwe_tsk.rs
poulpy-core/src/encryption/gglwe_tsk.rsRaw
1use poulpy_hal::{
2 api::{
3 ModuleN, ScratchAvailable, ScratchTakeBasic, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
4 VecZnxAddScalarInplace, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
5 VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
6 VecZnxSubInplace, VecZnxSwitchRing,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, Scratch},
9 oep::{VecZnxAddScalarInplaceImpl, VecZnxBigAllocBytesImpl, VecZnxDftApplyImpl, SvpApplyDftToDftImpl, VecZnxIdftApplyTmpAImpl, VecZnxBigNormalizeImpl},
10 source::Source,
11};
12
13use crate::{
14 ScratchTakeCore,
15 layouts::{
16 GetDist, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, LWEInfos, Rank, TensorKey, TensorKeyToMut,
17 prepared::GLWESecretPrepared,
18 },
19 encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
20};
21
22impl TensorKey<Vec<u8>> {
23 pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
24 where
25 A: GGLWEInfos,
26 M: GGLWETensorKeyEncryptSk<BE>
27 {
28 module.gglwe_tensor_key_encrypt_sk_tmp_bytes(infos)
29 }
30 // pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
31 // where
32 // A: GGLWEInfos,
33 // Module<B>: ModuleN + SvpPPolBytesOf + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigBytesOf,
34 // {
35 // GLWESecretPrepared::bytes_of(module, infos.rank_out())
36 // + module.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
37 // + module.bytes_of_vec_znx_big(1, 1)
38 // + module.bytes_of_vec_znx_dft(1, 1)
39 // + GLWESecret::bytes_of(module, Rank(1))
40 // + GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos)
41 // }
42}
43
44impl<DataSelf: DataMut> TensorKey<DataSelf> {
45 pub fn encrypt_sk<M, DataSk: DataRef, BE: Backend>(
46 &mut self,
47 module: &M,
48 sk: &GLWESecret<DataSk>,
49 source_xa: &mut Source,
50 source_xe: &mut Source,
51 scratch: &mut Scratch<BE>,
52 ) where
53 M: GGLWETensorKeyEncryptSk<BE>,
54 GLWESecret<DataSk>: GetDist,
55 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
56 {
57 module.gglwe_tensor_key_encrypt_sk(self, sk, source_xa, source_xe, scratch);
58 }
59
60 // pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
61 // &mut self,
62 // module: &Module<B>,
63 // sk: &GLWESecret<DataSk>,
64 // source_xa: &mut Source,
65 // source_xe: &mut Source,
66 // scratch: &mut Scratch<B>,
67 // ) where
68 // GLWESecret<DataSk>: GetDist,
69 // Module<B>: ModuleN
70 // + SvpApplyDftToDft<B>
71 // + VecZnxIdftApplyTmpA<B>
72 // + VecZnxAddScalarInplace
73 // + VecZnxDftBytesOf
74 // + VecZnxBigNormalize<B>
75 // + VecZnxDftApply<B>
76 // + SvpApplyDftToDftInplace<B>
77 // + VecZnxIdftApplyConsume<B>
78 // + VecZnxNormalizeTmpBytes
79 // + VecZnxFillUniform
80 // + VecZnxSubInplace
81 // + VecZnxAddInplace
82 // + VecZnxNormalizeInplace<B>
83 // + VecZnxAddNormal
84 // + VecZnxNormalize<B>
85 // + VecZnxSub
86 // + SvpPrepare<B>
87 // + VecZnxSwitchRing
88 // + SvpPPolBytesOf
89 // + VecZnxBigAllocBytesImpl<B>
90 // + VecZnxBigBytesOf
91 // + SvpPPolAlloc<B>,
92 // Scratch<B>: ScratchTakeBasic + ScratchTakeCore<B>,
93 // {
94 // #[cfg(debug_assertions)]
95 // {
96 // assert_eq!(self.rank_out(), sk.rank());
97 // assert_eq!(self.n(), sk.n());
98 // }
99
100 // // let n: RingDegree = sk.n();
101 // let rank: Rank = self.rank_out();
102
103 // let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(module, rank);
104 // sk_dft_prep.prepare(module, sk);
105
106 // let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(module, rank.into(), 1);
107
108 // (0..rank.into()).for_each(|i| {
109 // module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
110 // });
111
112 // let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(module, 1, 1);
113 // let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(module, Rank(1));
114 // let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(module, 1, 1);
115
116 // (0..rank.into()).for_each(|i| {
117 // (i..rank.into()).for_each(|j| {
118 // module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
119
120 // module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
121 // module.vec_znx_big_normalize(
122 // self.base2k().into(),
123 // &mut sk_ij.data.as_vec_znx_mut(),
124 // 0,
125 // self.base2k().into(),
126 // &sk_ij_big,
127 // 0,
128 // scratch_5,
129 // );
130
131 // self.at_mut(i, j)
132 // .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, scratch_5);
133 // });
134 // })
135 // }
136}
137
138pub trait GGLWETensorKeyEncryptSk<BE: Backend>
139where
140 Self: Sized
141 + ModuleN
142 + SvpPPolBytesOf
143 + SvpPPolAlloc<BE>
144 + VecZnxNormalizeTmpBytes
145 + VecZnxDftBytesOf
146 + VecZnxNormalizeTmpBytes
147 + VecZnxBigBytesOf,
148{
149 fn gglwe_tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
150 where
151 A: GGLWEInfos;
152
153 fn gglwe_tensor_key_encrypt_sk<R, S>(
154 &self,
155 res: &mut R,
156 sk: &S,
157 source_xa: &mut Source,
158 source_xe: &mut Source,
159 scratch: &mut Scratch<BE>,
160 ) where
161 R: TensorKeyToMut,
162 S: GLWESecretToRef + GetDist;
163}
164
165impl<BE: Backend> GGLWETensorKeyEncryptSk<BE> for Module<BE> where
166 Module<BE>: ModuleN
167 + SvpPPolBytesOf
168 + SvpPPolAlloc<BE>
169 + VecZnxNormalizeTmpBytes
170 + VecZnxDftBytesOf
171 + VecZnxNormalizeTmpBytes
172 + VecZnxBigBytesOf
173 + VecZnxAddScalarInplaceImpl<BE>
174 + VecZnxDftApply<BE>
175 + VecZnxDftApplyImpl<BE>
176 + SvpApplyDftToDftImpl<BE>
177 + GLWESwitchingKeyEncryptSk<BE>
178 + SvpApplyDftToDft<BE>
179 + VecZnxIdftApplyTmpAImpl<BE>
180 + VecZnxBigNormalizeImpl<BE>
181 + VecZnxIdftApplyTmpA<BE>
182 + VecZnxBigNormalize<BE>
183 + VecZnxAddScalarInplaceImpl<BE>
184 + SvpPrepare<BE>,
185 Scratch<BE>: ScratchTakeBasic + ScratchTakeCore<BE>,
186{
187 fn gglwe_tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
188 where
189 A: GGLWEInfos,
190 {
191 GLWESecretPrepared::bytes_of(self, infos.rank_out())
192 + self.bytes_of_vec_znx_dft(infos.rank_out().into(), 1)
193 + self.bytes_of_vec_znx_big(1, 1)
194 + self.bytes_of_vec_znx_dft(1, 1)
195 + GLWESecret::bytes_of(self, Rank(1))
196 + GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos)
197 }
198
199 fn gglwe_tensor_key_encrypt_sk<R, S>(
200 &self,
201 res: &mut R,
202 sk: &S,
203 source_xa: &mut Source,
204 source_xe: &mut Source,
205 scratch: &mut Scratch<BE>,
206 ) where
207 R: TensorKeyToMut,
208 S: GLWESecretToRef + GetDist,
209 {
210 let res: &mut TensorKey<&mut [u8]> = &mut res.to_mut();
211
212 // let n: RingDegree = sk.n();
213 let rank: Rank = res.rank_out();
214
215 let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, rank);
216 sk_dft_prep.prepare(self, sk);
217
218 let sk: &GLWESecret<&[u8]> = &sk.to_ref();
219
220 #[cfg(debug_assertions)]
221 {
222 assert_eq!(res.rank_out(), sk.rank());
223 assert_eq!(res.n(), sk.n());
224 }
225
226 let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, rank.into(), 1);
227
228 (0..rank.into()).for_each(|i| {
229 self.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
230 });
231
232 let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(self, 1, 1);
233 let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(self, Rank(1));
234 let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(self, 1, 1);
235
236 (0..rank.into()).for_each(|i| {
237 (i..rank.into()).for_each(|j| {
238 self.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
239
240 self.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
241 self.vec_znx_big_normalize(
242 res.base2k().into(),
243 &mut sk_ij.data.as_vec_znx_mut(),
244 0,
245 res.base2k().into(),
246 &sk_ij_big,
247 0,
248 scratch_5,
249 );
250
251 res.at_mut(i, j)
252 .encrypt_sk(self, &sk_ij, sk, source_xa, source_xe, scratch_5);
253 });
254 })
255 }
256}