poulpy/poulpy-core/src/tests/test_suite/encryption/glwe_tsk.rs
poulpy-core/src/tests/test_suite/encryption/glwe_tsk.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf,
4 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigAlloc, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf,
6 VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace,
7 VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSubScalarInplace, VecZnxSwitchRing,
8 },
9 layouts::{Backend, Module, ScratchOwned, VecZnxDft},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
13 },
14 source::Source,
15};
16
17use crate::{
18 encryption::SIGMA,
19 layouts::{
20 Dsize, GLWEPlaintext, GLWESecret, TensorKey, TensorKeyLayout,
21 compressed::{Decompress, TensorKeyCompressed},
22 prepared::{GLWESecretPrepared, PrepareAlloc},
23 },
24};
25
26pub fn test_gglwe_tensor_key_encrypt_sk<B>(module: &Module<B>)
27where
28 Module<B>: VecZnxDftBytesOf
29 + VecZnxBigNormalize<B>
30 + VecZnxDftApply<B>
31 + SvpApplyDftToDftInplace<B>
32 + VecZnxIdftApplyConsume<B>
33 + VecZnxNormalizeTmpBytes
34 + VecZnxFillUniform
35 + VecZnxSubInplace
36 + VecZnxAddInplace
37 + VecZnxNormalizeInplace<B>
38 + VecZnxAddNormal
39 + VecZnxNormalize<B>
40 + VecZnxSub
41 + SvpPrepare<B>
42 + SvpPPolBytesOf
43 + SvpPPolAlloc<B>
44 + VecZnxBigAddSmallInplace<B>
45 + VecZnxBigBytesOf
46 + VecZnxBigAddInplace<B>
47 + VecZnxCopy
48 + VecZnxDftAlloc<B>
49 + SvpApplyDftToDft<B>
50 + VecZnxBigAlloc<B>
51 + VecZnxIdftApplyTmpA<B>
52 + VecZnxAddScalarInplace
53 + VecZnxSwitchRing
54 + VecZnxSubScalarInplace,
55 B: Backend
56 + TakeVecZnxDftImpl<B>
57 + TakeVecZnxBigImpl<B>
58 + ScratchOwnedAllocImpl<B>
59 + ScratchOwnedBorrowImpl<B>
60 + ScratchAvailableImpl<B>
61 + TakeScalarZnxImpl<B>
62 + TakeVecZnxImpl<B>
63 + VecZnxDftAllocBytesImpl<B>
64 + VecZnxBigAllocBytesImpl<B>
65 + TakeSvpPPolImpl<B>,
66{
67 let base2k: usize = 8;
68 let k: usize = 54;
69
70 for rank in 1_usize..3 {
71 let n: usize = module.n();
72 let dnum: usize = k / base2k;
73
74 let tensor_key_infos = TensorKeyLayout {
75 n: n.into(),
76 base2k: base2k.into(),
77 k: k.into(),
78 dnum: dnum.into(),
79 dsize: Dsize(1),
80 rank: rank.into(),
81 };
82
83 let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&tensor_key_infos);
84
85 let mut source_xs: Source = Source::new([0u8; 32]);
86 let mut source_xe: Source = Source::new([0u8; 32]);
87 let mut source_xa: Source = Source::new([0u8; 32]);
88
89 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(TensorKey::encrypt_sk_scratch_space(
90 module,
91 &tensor_key_infos,
92 ));
93
94 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
95 sk.fill_ternary_prob(0.5, &mut source_xs);
96 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
97
98 tensor_key.encrypt_sk(
99 module,
100 &sk,
101 &mut source_xa,
102 &mut source_xe,
103 scratch.borrow(),
104 );
105
106 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
107
108 let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
109 let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
110 let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
111 let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(rank, 1);
112
113 for i in 0..rank {
114 module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
115 }
116
117 for i in 0..rank {
118 for j in 0..rank {
119 module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
120 module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
121 module.vec_znx_big_normalize(
122 base2k,
123 &mut sk_ij.data.as_vec_znx_mut(),
124 0,
125 base2k,
126 &sk_ij_big,
127 0,
128 scratch.borrow(),
129 );
130 for row_i in 0..dnum {
131 tensor_key
132 .at(i, j)
133 .at(row_i, 0)
134 .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
135
136 module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
137
138 let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
139 assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
140 }
141 }
142 }
143 }
144}
145
146pub fn test_gglwe_tensor_key_compressed_encrypt_sk<B>(module: &Module<B>)
147where
148 Module<B>: VecZnxDftBytesOf
149 + VecZnxBigNormalize<B>
150 + VecZnxDftApply<B>
151 + SvpApplyDftToDftInplace<B>
152 + VecZnxIdftApplyConsume<B>
153 + VecZnxNormalizeTmpBytes
154 + VecZnxFillUniform
155 + VecZnxSubInplace
156 + VecZnxAddInplace
157 + VecZnxNormalizeInplace<B>
158 + VecZnxAddNormal
159 + VecZnxNormalize<B>
160 + VecZnxSub
161 + SvpPrepare<B>
162 + SvpPPolBytesOf
163 + SvpPPolAlloc<B>
164 + VecZnxBigAddSmallInplace<B>
165 + VecZnxBigBytesOf
166 + VecZnxBigAddInplace<B>
167 + VecZnxCopy
168 + VecZnxDftAlloc<B>
169 + SvpApplyDftToDft<B>
170 + VecZnxBigAlloc<B>
171 + VecZnxIdftApplyTmpA<B>
172 + VecZnxAddScalarInplace
173 + VecZnxSwitchRing
174 + VecZnxSubScalarInplace,
175 B: Backend
176 + TakeVecZnxDftImpl<B>
177 + TakeVecZnxBigImpl<B>
178 + ScratchOwnedAllocImpl<B>
179 + ScratchOwnedBorrowImpl<B>
180 + ScratchAvailableImpl<B>
181 + TakeScalarZnxImpl<B>
182 + TakeVecZnxImpl<B>
183 + VecZnxDftAllocBytesImpl<B>
184 + VecZnxBigAllocBytesImpl<B>
185 + TakeSvpPPolImpl<B>,
186{
187 let base2k = 8;
188 let k = 54;
189 for rank in 1_usize..3 {
190 let n: usize = module.n();
191 let dnum: usize = k / base2k;
192
193 let tensor_key_infos: TensorKeyLayout = TensorKeyLayout {
194 n: n.into(),
195 base2k: base2k.into(),
196 k: k.into(),
197 dnum: dnum.into(),
198 dsize: Dsize(1),
199 rank: rank.into(),
200 };
201
202 let mut tensor_key_compressed: TensorKeyCompressed<Vec<u8>> = TensorKeyCompressed::alloc_from_infos(&tensor_key_infos);
203
204 let mut source_xs: Source = Source::new([0u8; 32]);
205 let mut source_xe: Source = Source::new([0u8; 32]);
206
207 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(TensorKeyCompressed::encrypt_sk_scratch_space(
208 module,
209 &tensor_key_infos,
210 ));
211
212 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&tensor_key_infos);
213 sk.fill_ternary_prob(0.5, &mut source_xs);
214 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
215
216 let seed_xa: [u8; 32] = [1u8; 32];
217
218 tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
219
220 let mut tensor_key: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&tensor_key_infos);
221 tensor_key.decompress(module, &tensor_key_compressed);
222
223 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&tensor_key_infos);
224
225 let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
226 let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
227 let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), 1_u32.into());
228 let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(rank, 1);
229
230 for i in 0..rank {
231 module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
232 }
233
234 for i in 0..rank {
235 for j in 0..rank {
236 module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
237 module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
238 module.vec_znx_big_normalize(
239 base2k,
240 &mut sk_ij.data.as_vec_znx_mut(),
241 0,
242 base2k,
243 &sk_ij_big,
244 0,
245 scratch.borrow(),
246 );
247 for row_i in 0..dnum {
248 tensor_key
249 .at(i, j)
250 .at(row_i, 0)
251 .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
252
253 module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, 0);
254
255 let std_pt: f64 = pt.data.std(base2k, 0) * (k as f64).exp2();
256 assert!((SIGMA - std_pt).abs() <= 0.5, "{SIGMA} {std_pt}");
257 }
258 }
259 }
260 }
261}