arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/test_fft64/ggsw.rs
core/src/test_fft64/ggsw.rsRaw
  1use base2k::{
  2    FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps,
  3    VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, VecZnxToMut, ZnxViewMut, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    elem::{GetRow, Infos},
  9    ggsw_ciphertext::GGSWCiphertext,
 10    glwe_ciphertext_fourier::GLWECiphertextFourier,
 11    glwe_plaintext::GLWEPlaintext,
 12    keys::{SecretKey, SecretKeyFourier},
 13    keyswitch_key::GLWESwitchingKey,
 14};
 15
 16#[test]
 17fn encrypt_sk() {
 18    (1..4).for_each(|rank| {
 19        println!("test encrypt_sk rank: {}", rank);
 20        test_encrypt_sk(11, 8, 54, 3.2, rank);
 21    });
 22}
 23
 24#[test]
 25fn external_product() {
 26    (1..4).for_each(|rank| {
 27        println!("test external_product rank: {}", rank);
 28        test_external_product(12, 12, 60, rank, 3.2);
 29    });
 30}
 31
 32#[test]
 33fn external_product_inplace() {
 34    (1..4).for_each(|rank| {
 35        println!("test external_product rank: {}", rank);
 36        test_external_product_inplace(12, 15, 60, rank, 3.2);
 37    });
 38}
 39
 40fn test_encrypt_sk(log_n: usize, basek: usize, k_ggsw: usize, sigma: f64, rank: usize) {
 41    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 42
 43    let rows: usize = (k_ggsw + basek - 1) / basek;
 44
 45    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
 46    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
 47    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
 48    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 49
 50    let mut source_xs: Source = Source::new([0u8; 32]);
 51    let mut source_xe: Source = Source::new([0u8; 32]);
 52    let mut source_xa: Source = Source::new([0u8; 32]);
 53
 54    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
 55
 56    let mut scratch: ScratchOwned = ScratchOwned::new(
 57        GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size())
 58            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size()),
 59    );
 60
 61    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
 62    sk.fill_ternary_prob(0.5, &mut source_xs);
 63
 64    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
 65    sk_dft.dft(&module, &sk);
 66
 67    ct.encrypt_sk(
 68        &module,
 69        &pt_scalar,
 70        &sk_dft,
 71        &mut source_xa,
 72        &mut source_xe,
 73        sigma,
 74        scratch.borrow(),
 75    );
 76
 77    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
 78    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
 79    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
 80
 81    (0..ct.rank() + 1).for_each(|col_j| {
 82        (0..ct.rows()).for_each(|row_i| {
 83            module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
 84
 85            // mul with sk[col_j-1]
 86            if col_j > 0 {
 87                module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
 88                module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
 89                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
 90                module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
 91            }
 92
 93            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
 94
 95            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
 96
 97            module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
 98
 99            let std_pt: f64 = pt_have.data.std(0, basek) * (k_ggsw as f64).exp2();
100            assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
101
102            pt_want.data.zero();
103        });
104    });
105}
106
107fn test_external_product(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
108    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
109
110    let rows: usize = (k_ggsw + basek - 1) / basek;
111
112    let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
113    let mut ct_ggsw_lhs_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
114    let mut ct_ggsw_lhs_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
115    let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
116    let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
117
118    let mut source_xs: Source = Source::new([0u8; 32]);
119    let mut source_xe: Source = Source::new([0u8; 32]);
120    let mut source_xa: Source = Source::new([0u8; 32]);
121
122    pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
123
124    let k: usize = 1;
125
126    pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
127
128    let mut scratch: ScratchOwned = ScratchOwned::new(
129        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_ggsw_rhs.size())
130            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_ggsw_lhs_out.size())
131            | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_ggsw_lhs_in.size())
132            | GGSWCiphertext::external_product_scratch_space(
133                &module,
134                ct_ggsw_lhs_out.size(),
135                ct_ggsw_lhs_in.size(),
136                ct_ggsw_rhs.size(),
137                rank,
138            ),
139    );
140
141    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
142    sk.fill_ternary_prob(0.5, &mut source_xs);
143
144    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
145    sk_dft.dft(&module, &sk);
146
147    ct_ggsw_rhs.encrypt_sk(
148        &module,
149        &pt_ggsw_rhs,
150        &sk_dft,
151        &mut source_xa,
152        &mut source_xe,
153        sigma,
154        scratch.borrow(),
155    );
156
157    ct_ggsw_lhs_in.encrypt_sk(
158        &module,
159        &pt_ggsw_lhs,
160        &sk_dft,
161        &mut source_xa,
162        &mut source_xe,
163        sigma,
164        scratch.borrow(),
165    );
166
167    ct_ggsw_lhs_out.external_product(&module, &ct_ggsw_lhs_in, &ct_ggsw_rhs, scratch.borrow());
168
169    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
170    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
171    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs_out.size());
172    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs_out.size());
173    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
174
175    module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
176
177    (0..ct_ggsw_lhs_out.rank() + 1).for_each(|col_j| {
178        (0..ct_ggsw_lhs_out.rows()).for_each(|row_i| {
179            module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_ggsw_lhs, 0);
180
181            if col_j > 0 {
182                module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
183                module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
184                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
185                module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
186            }
187
188            ct_ggsw_lhs_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
189            ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
190
191            module.vec_znx_sub_ab_inplace(&mut pt, 0, &pt_want, 0);
192
193            let noise_have: f64 = pt.data.std(0, basek).log2();
194
195            let var_gct_err_lhs: f64 = sigma * sigma;
196            let var_gct_err_rhs: f64 = 0f64;
197
198            let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
199            let var_a0_err: f64 = sigma * sigma;
200            let var_a1_err: f64 = 1f64 / 12f64;
201
202            let noise_want: f64 = noise_ggsw_product(
203                module.n() as f64,
204                basek,
205                0.5,
206                var_msg,
207                var_a0_err,
208                var_a1_err,
209                var_gct_err_lhs,
210                var_gct_err_rhs,
211                rank as f64,
212                k_ggsw,
213                k_ggsw,
214            );
215
216            assert!(
217                (noise_have - noise_want).abs() <= 0.1,
218                "have: {} want: {}",
219                noise_have,
220                noise_want
221            );
222
223            pt_want.data.zero();
224        });
225    });
226}
227
228fn test_external_product_inplace(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
229    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
230    let rows: usize = (k_ggsw + basek - 1) / basek;
231
232    let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
233    let mut ct_ggsw_lhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
234    let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
235    let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
236
237    let mut source_xs: Source = Source::new([0u8; 32]);
238    let mut source_xe: Source = Source::new([0u8; 32]);
239    let mut source_xa: Source = Source::new([0u8; 32]);
240
241    pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
242
243    let k: usize = 1;
244
245    pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
246
247    let mut scratch: ScratchOwned = ScratchOwned::new(
248        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_ggsw_rhs.size())
249            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_ggsw_lhs.size())
250            | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_ggsw_lhs.size())
251            | GGSWCiphertext::external_product_inplace_scratch_space(&module, ct_ggsw_lhs.size(), ct_ggsw_rhs.size(), rank),
252    );
253
254    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
255    sk.fill_ternary_prob(0.5, &mut source_xs);
256
257    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
258    sk_dft.dft(&module, &sk);
259
260    ct_ggsw_rhs.encrypt_sk(
261        &module,
262        &pt_ggsw_rhs,
263        &sk_dft,
264        &mut source_xa,
265        &mut source_xe,
266        sigma,
267        scratch.borrow(),
268    );
269
270    ct_ggsw_lhs.encrypt_sk(
271        &module,
272        &pt_ggsw_lhs,
273        &sk_dft,
274        &mut source_xa,
275        &mut source_xe,
276        sigma,
277        scratch.borrow(),
278    );
279
280    ct_ggsw_lhs.external_product_inplace(&module, &ct_ggsw_rhs, scratch.borrow());
281
282    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
283    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
284    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs.size());
285    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs.size());
286    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
287
288    module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
289
290    (0..ct_ggsw_lhs.rank() + 1).for_each(|col_j| {
291        (0..ct_ggsw_lhs.rows()).for_each(|row_i| {
292            module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_ggsw_lhs, 0);
293
294            if col_j > 0 {
295                module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
296                module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
297                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
298                module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
299            }
300
301            ct_ggsw_lhs.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
302            ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
303
304            module.vec_znx_sub_ab_inplace(&mut pt, 0, &pt_want, 0);
305
306            let noise_have: f64 = pt.data.std(0, basek).log2();
307
308            let var_gct_err_lhs: f64 = sigma * sigma;
309            let var_gct_err_rhs: f64 = 0f64;
310
311            let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
312            let var_a0_err: f64 = sigma * sigma;
313            let var_a1_err: f64 = 1f64 / 12f64;
314
315            let noise_want: f64 = noise_ggsw_product(
316                module.n() as f64,
317                basek,
318                0.5,
319                var_msg,
320                var_a0_err,
321                var_a1_err,
322                var_gct_err_lhs,
323                var_gct_err_rhs,
324                rank as f64,
325                k_ggsw,
326                k_ggsw,
327            );
328
329            assert!(
330                (noise_have - noise_want).abs() <= 0.1,
331                "have: {} want: {}",
332                noise_have,
333                noise_want
334            );
335
336            pt_want.data.zero();
337        });
338    });
339}
340pub(crate) fn noise_ggsw_product(
341    n: f64,
342    basek: usize,
343    var_xs: f64,
344    var_msg: f64,
345    var_a0_err: f64,
346    var_a1_err: f64,
347    var_gct_err_lhs: f64,
348    var_gct_err_rhs: f64,
349    rank: f64,
350    a_logq: usize,
351    b_logq: usize,
352) -> f64 {
353    let a_logq: usize = a_logq.min(b_logq);
354    let a_cols: usize = (a_logq + basek - 1) / basek;
355
356    let b_scale = 2.0f64.powi(b_logq as i32);
357    let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
358
359    let base: f64 = (1 << (basek)) as f64;
360    let var_base: f64 = base * base / 12f64;
361
362    // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
363    // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
364    let mut noise: f64 = (rank + 1.0) * (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
365    noise += var_msg * var_a0_err * a_scale * a_scale * n;
366    noise += var_msg * var_a1_err * a_scale * a_scale * n * var_xs * rank;
367    noise = noise.sqrt();
368    noise /= b_scale;
369    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
370}