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 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, VecZnxAddInplace,
4 VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
5 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
6 VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
7 VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, 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 + VecZnxDftFromVecZnx<B>
40 + SvpApplyInplace<B>
41 + VecZnxDftToVecZnxBigConsume<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 + VmpApplyTmpBytes
60 + VmpApply<B>
61 + VmpApplyAdd<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, n, basek, ksk.k(), rank_in, rank_out)
90 | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_in.k())
91 | GLWECiphertext::keyswitch_scratch_space(
92 module,
93 n,
94 basek,
95 ct_out.k(),
96 ct_in.k(),
97 ksk.k(),
98 digits,
99 rank_in,
100 rank_out,
101 ),
102 );
103
104 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
105 sk_in.fill_ternary_prob(0.5, &mut source_xs);
106 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
107
108 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
109 sk_out.fill_ternary_prob(0.5, &mut source_xs);
110 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
111
112 ksk.encrypt_sk(
113 module,
114 &sk_in,
115 &sk_out,
116 &mut source_xa,
117 &mut source_xe,
118 scratch.borrow(),
119 );
120
121 ct_in.encrypt_sk(
122 module,
123 &pt_want,
124 &sk_in_prepared,
125 &mut source_xa,
126 &mut source_xe,
127 scratch.borrow(),
128 );
129
130 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
131
132 ct_out.keyswitch(module, &ct_in, &ksk_prepared, scratch.borrow());
133
134 let max_noise: f64 = log2_std_noise_gglwe_product(
135 module.n() as f64,
136 basek * digits,
137 0.5,
138 0.5,
139 0f64,
140 SIGMA * SIGMA,
141 0f64,
142 rank_in as f64,
143 k_in,
144 k_ksk,
145 );
146
147 ct_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
148}
149
150pub fn test_glwe_keyswitch_inplace<B>(module: &Module<B>, basek: usize, k_ct: usize, k_ksk: usize, digits: usize, rank: usize)
151where
152 Module<B>: VecZnxDftAllocBytes
153 + VecZnxBigNormalize<B>
154 + VecZnxDftFromVecZnx<B>
155 + SvpApplyInplace<B>
156 + VecZnxDftToVecZnxBigConsume<B>
157 + VecZnxFillUniform
158 + VecZnxSubABInplace
159 + VecZnxAddInplace
160 + VecZnxNormalizeInplace<B>
161 + VecZnxAddNormal
162 + VecZnxNormalize<B>
163 + VecZnxSub
164 + SvpPrepare<B>
165 + SvpPPolAllocBytes
166 + SvpPPolAlloc<B>
167 + VecZnxBigAllocBytes
168 + VecZnxBigAddInplace<B>
169 + VecZnxBigAddSmallInplace<B>
170 + VecZnxNormalizeTmpBytes
171 + VecZnxAddScalarInplace
172 + VmpPMatAlloc<B>
173 + VmpPrepare<B>
174 + VmpApplyTmpBytes
175 + VmpApply<B>
176 + VmpApplyAdd<B>
177 + VecZnxBigNormalizeTmpBytes
178 + VecZnxSwithcDegree,
179 B: Backend
180 + TakeVecZnxDftImpl<B>
181 + TakeVecZnxBigImpl<B>
182 + TakeSvpPPolImpl<B>
183 + ScratchOwnedAllocImpl<B>
184 + ScratchOwnedBorrowImpl<B>
185 + ScratchAvailableImpl<B>
186 + TakeScalarZnxImpl<B>
187 + TakeVecZnxImpl<B>,
188{
189 let n: usize = module.n();
190 let rows: usize = k_ct.div_ceil(basek * digits);
191
192 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
193 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
194 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
195
196 let mut source_xs: Source = Source::new([0u8; 32]);
197 let mut source_xe: Source = Source::new([0u8; 32]);
198 let mut source_xa: Source = Source::new([0u8; 32]);
199
200 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
201
202 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
203 GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, ksk.k(), rank, rank)
204 | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_glwe.k())
205 | GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, ct_glwe.k(), ksk.k(), digits, rank),
206 );
207
208 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
209 sk_in.fill_ternary_prob(0.5, &mut source_xs);
210 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
211
212 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
213 sk_out.fill_ternary_prob(0.5, &mut source_xs);
214 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
215
216 ksk.encrypt_sk(
217 module,
218 &sk_in,
219 &sk_out,
220 &mut source_xa,
221 &mut source_xe,
222 scratch.borrow(),
223 );
224
225 ct_glwe.encrypt_sk(
226 module,
227 &pt_want,
228 &sk_in_prepared,
229 &mut source_xa,
230 &mut source_xe,
231 scratch.borrow(),
232 );
233
234 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
235
236 ct_glwe.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
237
238 let max_noise: f64 = log2_std_noise_gglwe_product(
239 module.n() as f64,
240 basek * digits,
241 0.5,
242 0.5,
243 0f64,
244 SIGMA * SIGMA,
245 0f64,
246 rank as f64,
247 k_ct,
248 k_ksk,
249 );
250
251 ct_glwe.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
252}