poulpy/core/src/test_fft64/rlwe.rs
core/src/test_fft64/rlwe.rsRaw
1#[cfg(test)]
2mod tests {
3 use base2k::{Decoding, Encoding, FFT64, Module, ScratchOwned, Stats, VecZnxOps, ZnxZero};
4 use itertools::izip;
5 use sampling::source::Source;
6
7 use crate::{
8 elem::Infos,
9 keys::{PublicKey, SecretKey, SecretKeyDft},
10 rlwe::{RLWECt, RLWECtDft, RLWEPt},
11 };
12
13 #[test]
14 fn encrypt_sk() {
15 let module: Module<FFT64> = Module::<FFT64>::new(32);
16 let log_base2k: usize = 8;
17 let log_k_ct: usize = 54;
18 let log_k_pt: usize = 30;
19
20 let sigma: f64 = 3.2;
21 let bound: f64 = sigma * 6.0;
22
23 let mut ct: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_ct);
24 let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_pt);
25
26 let mut source_xs: Source = Source::new([0u8; 32]);
27 let mut source_xe: Source = Source::new([0u8; 32]);
28 let mut source_xa: Source = Source::new([0u8; 32]);
29
30 let mut scratch: ScratchOwned = ScratchOwned::new(
31 RLWECt::encrypt_sk_scratch_space(&module, ct.size()) | RLWECt::decrypt_scratch_space(&module, ct.size()),
32 );
33
34 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
35 sk.fill_ternary_prob(0.5, &mut source_xs);
36
37 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
38 sk_dft.dft(&module, &sk);
39
40 let mut data_want: Vec<i64> = vec![0i64; module.n()];
41
42 data_want
43 .iter_mut()
44 .for_each(|x| *x = source_xa.next_i64() & 0xFF);
45
46 pt.data
47 .encode_vec_i64(0, log_base2k, log_k_pt, &data_want, 10);
48
49 ct.encrypt_sk(
50 &module,
51 Some(&pt),
52 &sk_dft,
53 &mut source_xa,
54 &mut source_xe,
55 sigma,
56 bound,
57 scratch.borrow(),
58 );
59
60 pt.data.zero();
61
62 ct.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
63
64 let mut data_have: Vec<i64> = vec![0i64; module.n()];
65
66 pt.data
67 .decode_vec_i64(0, log_base2k, pt.size() * log_base2k, &mut data_have);
68
69 // TODO: properly assert the decryption noise through std(dec(ct) - pt)
70 let scale: f64 = (1 << (pt.size() * log_base2k - log_k_pt)) as f64;
71 izip!(data_want.iter(), data_have.iter()).for_each(|(a, b)| {
72 let b_scaled = (*b as f64) / scale;
73 assert!(
74 (*a as f64 - b_scaled).abs() < 0.1,
75 "{} {}",
76 *a as f64,
77 b_scaled
78 )
79 });
80
81 module.free();
82 }
83
84 #[test]
85 fn encrypt_zero_sk() {
86 let module: Module<FFT64> = Module::<FFT64>::new(1024);
87 let log_base2k: usize = 8;
88 let log_k_ct: usize = 55;
89
90 let sigma: f64 = 3.2;
91 let bound: f64 = sigma * 6.0;
92
93 let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
94
95 let mut source_xs: Source = Source::new([0u8; 32]);
96 let mut source_xe: Source = Source::new([1u8; 32]);
97 let mut source_xa: Source = Source::new([0u8; 32]);
98
99 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
100 sk.fill_ternary_prob(0.5, &mut source_xs);
101 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
102 sk_dft.dft(&module, &sk);
103
104 let mut ct_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct);
105
106 let mut scratch: ScratchOwned = ScratchOwned::new(
107 RLWECtDft::decrypt_scratch_space(&module, ct_dft.size())
108 | RLWECtDft::encrypt_zero_sk_scratch_space(&module, ct_dft.size()),
109 );
110
111 ct_dft.encrypt_zero_sk(
112 &module,
113 &sk_dft,
114 &mut source_xa,
115 &mut source_xe,
116 sigma,
117 bound,
118 scratch.borrow(),
119 );
120 ct_dft.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
121
122 assert!((sigma - pt.data.std(0, log_base2k) * (log_k_ct as f64).exp2()) <= 0.2);
123 module.free();
124 }
125
126 #[test]
127 fn encrypt_pk() {
128 let module: Module<FFT64> = Module::<FFT64>::new(32);
129 let log_base2k: usize = 8;
130 let log_k_ct: usize = 54;
131 let log_k_pk: usize = 64;
132
133 let sigma: f64 = 3.2;
134 let bound: f64 = sigma * 6.0;
135
136 let mut ct: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_ct);
137 let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
138
139 let mut source_xs: Source = Source::new([0u8; 32]);
140 let mut source_xe: Source = Source::new([0u8; 32]);
141 let mut source_xa: Source = Source::new([0u8; 32]);
142 let mut source_xu: Source = Source::new([0u8; 32]);
143
144 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
145 sk.fill_ternary_prob(0.5, &mut source_xs);
146 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
147 sk_dft.dft(&module, &sk);
148
149 let mut pk: PublicKey<Vec<u8>, FFT64> = PublicKey::new(&module, log_base2k, log_k_pk);
150 pk.generate(
151 &module,
152 &sk_dft,
153 &mut source_xa,
154 &mut source_xe,
155 sigma,
156 bound,
157 );
158
159 let mut scratch: ScratchOwned = ScratchOwned::new(
160 RLWECt::encrypt_sk_scratch_space(&module, ct.size())
161 | RLWECt::decrypt_scratch_space(&module, ct.size())
162 | RLWECt::encrypt_pk_scratch_space(&module, pk.size()),
163 );
164
165 let mut data_want: Vec<i64> = vec![0i64; module.n()];
166
167 data_want
168 .iter_mut()
169 .for_each(|x| *x = source_xa.next_i64() & 0);
170
171 pt_want
172 .data
173 .encode_vec_i64(0, log_base2k, log_k_ct, &data_want, 10);
174
175 ct.encrypt_pk(
176 &module,
177 Some(&pt_want),
178 &pk,
179 &mut source_xu,
180 &mut source_xe,
181 sigma,
182 bound,
183 scratch.borrow(),
184 );
185
186 let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
187
188 ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
189
190 module.vec_znx_sub_ab_inplace(&mut pt_want, 0, &pt_have, 0);
191
192 assert!(((1.0f64 / 12.0).sqrt() - pt_want.data.std(0, log_base2k) * (log_k_ct as f64).exp2()).abs() < 0.2);
193
194 module.free();
195 }
196}