poulpy/poulpy-core/src/tests/generics/keyswitch/glwe_ct.rs
poulpy-core/src/tests/generics/keyswitch/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 DFT, IDFTConsume, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes, VecZnxFillUniform,
6 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree,
7 VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
8 },
9 layouts::{Backend, Module, ScratchOwned},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
13 },
14 source::Source,
15};
16
17use crate::{
18 encryption::SIGMA,
19 layouts::{
20 GGLWESwitchingKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
21 prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
22 },
23 noise::log2_std_noise_gglwe_product,
24};
25
26#[allow(clippy::too_many_arguments)]
27pub fn test_glwe_keyswitch<B>(
28 module: &Module<B>,
29 basek: usize,
30 k_out: usize,
31 k_in: usize,
32 k_ksk: usize,
33 digits: usize,
34 rank_in: usize,
35 rank_out: usize,
36) where
37 Module<B>: VecZnxDftAllocBytes
38 + VecZnxBigNormalize<B>
39 + DFT<B>
40 + SvpApplyInplace<B>
41 + IDFTConsume<B>
42 + VecZnxFillUniform
43 + VecZnxSubABInplace
44 + VecZnxAddInplace
45 + VecZnxNormalizeInplace<B>
46 + VecZnxAddNormal
47 + VecZnxNormalize<B>
48 + VecZnxSub
49 + SvpPrepare<B>
50 + SvpPPolAllocBytes
51 + SvpPPolAlloc<B>
52 + VecZnxBigAllocBytes
53 + VecZnxBigAddInplace<B>
54 + VecZnxBigAddSmallInplace<B>
55 + VecZnxNormalizeTmpBytes
56 + VecZnxAddScalarInplace
57 + VmpPMatAlloc<B>
58 + VmpPrepare<B>
59 + VmpApplyDftToDftTmpBytes
60 + VmpApplyDftToDft<B>
61 + VmpApplyDftToDftAdd<B>
62 + VecZnxBigNormalizeTmpBytes
63 + VecZnxSwithcDegree,
64 B: Backend
65 + TakeVecZnxDftImpl<B>
66 + TakeVecZnxBigImpl<B>
67 + TakeSvpPPolImpl<B>
68 + ScratchOwnedAllocImpl<B>
69 + ScratchOwnedBorrowImpl<B>
70 + ScratchAvailableImpl<B>
71 + TakeScalarZnxImpl<B>
72 + TakeVecZnxImpl<B>,
73{
74 let n = module.n();
75 let rows: usize = k_in.div_ceil(basek * digits);
76
77 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
78 let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_in, rank_in);
79 let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_out, rank_out);
80 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
81
82 let mut source_xs: Source = Source::new([0u8; 32]);
83 let mut source_xe: Source = Source::new([0u8; 32]);
84 let mut source_xa: Source = Source::new([0u8; 32]);
85
86 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
87
88 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
89 GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank_in, rank_out)
90 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_in.k())
91 | GLWECiphertext::keyswitch_scratch_space(
92 module,
93 basek,
94 ct_out.k(),
95 ct_in.k(),
96 ksk.k(),
97 digits,
98 rank_in,
99 rank_out,
100 ),
101 );
102
103 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
104 sk_in.fill_ternary_prob(0.5, &mut source_xs);
105 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
106
107 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
108 sk_out.fill_ternary_prob(0.5, &mut source_xs);
109 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
110
111 ksk.encrypt_sk(
112 module,
113 &sk_in,
114 &sk_out,
115 &mut source_xa,
116 &mut source_xe,
117 scratch.borrow(),
118 );
119
120 ct_in.encrypt_sk(
121 module,
122 &pt_want,
123 &sk_in_prepared,
124 &mut source_xa,
125 &mut source_xe,
126 scratch.borrow(),
127 );
128
129 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
130
131 ct_out.keyswitch(module, &ct_in, &ksk_prepared, scratch.borrow());
132
133 let max_noise: f64 = log2_std_noise_gglwe_product(
134 module.n() as f64,
135 basek * digits,
136 0.5,
137 0.5,
138 0f64,
139 SIGMA * SIGMA,
140 0f64,
141 rank_in as f64,
142 k_in,
143 k_ksk,
144 );
145
146 ct_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
147}
148
149pub fn test_glwe_keyswitch_inplace<B>(module: &Module<B>, basek: usize, k_ct: usize, k_ksk: usize, digits: usize, rank: usize)
150where
151 Module<B>: VecZnxDftAllocBytes
152 + VecZnxBigNormalize<B>
153 + DFT<B>
154 + SvpApplyInplace<B>
155 + IDFTConsume<B>
156 + VecZnxFillUniform
157 + VecZnxSubABInplace
158 + VecZnxAddInplace
159 + VecZnxNormalizeInplace<B>
160 + VecZnxAddNormal
161 + VecZnxNormalize<B>
162 + VecZnxSub
163 + SvpPrepare<B>
164 + SvpPPolAllocBytes
165 + SvpPPolAlloc<B>
166 + VecZnxBigAllocBytes
167 + VecZnxBigAddInplace<B>
168 + VecZnxBigAddSmallInplace<B>
169 + VecZnxNormalizeTmpBytes
170 + VecZnxAddScalarInplace
171 + VmpPMatAlloc<B>
172 + VmpPrepare<B>
173 + VmpApplyDftToDftTmpBytes
174 + VmpApplyDftToDft<B>
175 + VmpApplyDftToDftAdd<B>
176 + VecZnxBigNormalizeTmpBytes
177 + VecZnxSwithcDegree,
178 B: Backend
179 + TakeVecZnxDftImpl<B>
180 + TakeVecZnxBigImpl<B>
181 + TakeSvpPPolImpl<B>
182 + ScratchOwnedAllocImpl<B>
183 + ScratchOwnedBorrowImpl<B>
184 + ScratchAvailableImpl<B>
185 + TakeScalarZnxImpl<B>
186 + TakeVecZnxImpl<B>,
187{
188 let n: usize = module.n();
189 let rows: usize = k_ct.div_ceil(basek * digits);
190
191 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
192 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
193 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
194
195 let mut source_xs: Source = Source::new([0u8; 32]);
196 let mut source_xe: Source = Source::new([0u8; 32]);
197 let mut source_xa: Source = Source::new([0u8; 32]);
198
199 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
200
201 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
202 GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank, rank)
203 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_glwe.k())
204 | GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, ct_glwe.k(), ksk.k(), digits, rank),
205 );
206
207 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
208 sk_in.fill_ternary_prob(0.5, &mut source_xs);
209 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
210
211 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
212 sk_out.fill_ternary_prob(0.5, &mut source_xs);
213 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
214
215 ksk.encrypt_sk(
216 module,
217 &sk_in,
218 &sk_out,
219 &mut source_xa,
220 &mut source_xe,
221 scratch.borrow(),
222 );
223
224 ct_glwe.encrypt_sk(
225 module,
226 &pt_want,
227 &sk_in_prepared,
228 &mut source_xa,
229 &mut source_xe,
230 scratch.borrow(),
231 );
232
233 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
234
235 ct_glwe.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
236
237 let max_noise: f64 = log2_std_noise_gglwe_product(
238 module.n() as f64,
239 basek * digits,
240 0.5,
241 0.5,
242 0f64,
243 SIGMA * SIGMA,
244 0f64,
245 rank as f64,
246 k_ct,
247 k_ksk,
248 );
249
250 ct_glwe.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
251}