poulpy/rlwe/src/encryption.rs
rlwe/src/encryption.rsRaw
1use std::cmp::min;
2
3use base2k::{
4 AddNormal, Backend, FFT64, FillUniform, Module, ScalarZnxDft, ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnx,
5 VecZnxAlloc, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut,
6 VecZnxDftToRef, VecZnxToMut, VecZnxToRef,
7};
8
9use sampling::source::Source;
10
11use crate::{
12 elem::{Infos, RLWECt, RLWECtDft, RLWEPt},
13 keys::SecretKeyDft,
14};
15
16pub fn encrypt_rlwe_sk_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
17 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
18}
19
20pub fn encrypt_rlwe_sk<C, P, S>(
21 module: &Module<FFT64>,
22 ct: &mut RLWECt<C>,
23 pt: Option<&RLWEPt<P>>,
24 sk: &SecretKeyDft<S, FFT64>,
25 source_xa: &mut Source,
26 source_xe: &mut Source,
27 scratch: &mut Scratch,
28 sigma: f64,
29 bound: f64,
30) where
31 VecZnx<C>: VecZnxToMut + VecZnxToRef,
32 VecZnx<P>: VecZnxToRef,
33 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
34{
35 let log_base2k: usize = ct.log_base2k();
36 let log_k: usize = ct.log_k();
37 let size: usize = ct.size();
38
39 // c1 = a
40 ct.data.fill_uniform(log_base2k, 1, size, source_xa);
41
42 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
43
44 {
45 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
46 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
47
48 // c0_dft = DFT(a) * DFT(s)
49 module.svp_apply_inplace(&mut c0_dft, 0, sk, 0);
50
51 // c0_big = IDFT(c0_dft)
52 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
53 }
54
55 // c0_big = m - c0_big
56 if let Some(pt) = pt {
57 module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, pt, 0);
58 }
59 // c0_big += e
60 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
61
62 // c0 = norm(c0_big = -as + m + e)
63 module.vec_znx_big_normalize(log_base2k, ct, 0, &c0_big, 0, scratch_1);
64}
65
66pub fn decrypt_rlwe<P, C, S>(
67 module: &Module<FFT64>,
68 pt: &mut RLWEPt<P>,
69 ct: &RLWECt<C>,
70 sk: &SecretKeyDft<S, FFT64>,
71 scratch: &mut Scratch,
72) where
73 VecZnx<P>: VecZnxToMut + VecZnxToRef,
74 VecZnx<C>: VecZnxToRef,
75 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
76{
77 let size: usize = min(pt.size(), ct.size());
78
79 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
80
81 {
82 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
83 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
84
85 // c0_dft = DFT(a) * DFT(s)
86 module.svp_apply_inplace(&mut c0_dft, 0, sk, 0);
87
88 // c0_big = IDFT(c0_dft)
89 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
90 }
91
92 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
93 module.vec_znx_big_add_small_inplace(&mut c0_big, 0, ct, 0);
94
95 // pt = norm(BIG(m + e))
96 module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
97
98 pt.log_base2k = ct.log_base2k();
99 pt.log_k = min(pt.log_k(), ct.log_k());
100}
101
102pub fn decrypt_rlwe_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
103 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
104}
105
106impl<C> RLWECt<C> {
107 pub fn encrypt_sk<P, S>(
108 &mut self,
109 module: &Module<FFT64>,
110 pt: Option<&RLWEPt<P>>,
111 sk: &SecretKeyDft<S, FFT64>,
112 source_xa: &mut Source,
113 source_xe: &mut Source,
114 scratch: &mut Scratch,
115 sigma: f64,
116 bound: f64,
117 ) where
118 VecZnx<C>: VecZnxToMut + VecZnxToRef,
119 VecZnx<P>: VecZnxToRef,
120 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
121 {
122 encrypt_rlwe_sk(
123 module, self, pt, sk, source_xa, source_xe, scratch, sigma, bound,
124 )
125 }
126
127 pub fn decrypt<P, S>(&self, module: &Module<FFT64>, pt: &mut RLWEPt<P>, sk: &SecretKeyDft<S, FFT64>, scratch: &mut Scratch)
128 where
129 VecZnx<P>: VecZnxToMut + VecZnxToRef,
130 VecZnx<C>: VecZnxToRef,
131 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
132 {
133 decrypt_rlwe(module, pt, self, sk, scratch);
134 }
135}
136
137pub(crate) fn encrypt_rlwe_zero_dft_scratch_bytes<B: Backend>(module: &Module<FFT64>, size: usize) -> usize {
138 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
139}
140
141impl<C> RLWECtDft<C, FFT64> {
142 fn encrypt_zero<S>(
143 module: &Module<FFT64>,
144 ct: &mut RLWECtDft<C, FFT64>,
145 sk: &SecretKeyDft<S, FFT64>,
146 source_xa: &mut Source,
147 source_xe: &mut Source,
148 scratch: &mut Scratch,
149 sigma: f64,
150 bound: f64,
151 ) where
152 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
153 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
154 {
155 let log_base2k: usize = ct.log_base2k();
156 let log_k: usize = ct.log_k();
157 let size: usize = ct.size();
158
159 // ct[1] = DFT(a)
160 {
161 let (mut tmp_znx, _) = scratch.tmp_vec_znx(module, 1, size);
162 tmp_znx.fill_uniform(log_base2k, 1, size, source_xa);
163 module.vec_znx_dft(ct, 1, &tmp_znx, 0);
164 }
165
166 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
167
168 {
169 let (mut tmp_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
170 // c0_dft = DFT(a) * DFT(s)
171 module.svp_apply(&mut tmp_dft, 0, sk, 0, ct, 1);
172 // c0_big = IDFT(c0_dft)
173 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut tmp_dft, 0);
174 }
175
176 // c0_big += e
177 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
178
179 // c0 = norm(c0_big = -as + e)
180 let (mut tmp_znx, scratch_2) = scratch_1.tmp_vec_znx(module, 1, size);
181 module.vec_znx_big_normalize(log_base2k, &mut tmp_znx, 0, &c0_big, 0, scratch_2);
182 // ct[0] = DFT(-as + e)
183 module.vec_znx_dft(ct, 0, &tmp_znx, 0);
184 }
185
186 fn encrypt_zero_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
187 (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
188 + module.bytes_of_vec_znx_big(1, size)
189 + module.bytes_of_vec_znx(1, size)
190 + module.vec_znx_big_normalize_tmp_bytes()
191 }
192}
193
194#[cfg(test)]
195mod tests {
196 use base2k::{Encoding, FFT64, Module, ScratchOwned, ZnxZero};
197 use itertools::izip;
198 use sampling::source::Source;
199
200 use crate::{
201 elem::{Infos, RLWECt, RLWEPt},
202 keys::{SecretKey, SecretKeyDft},
203 };
204
205 use super::{decrypt_rlwe_scratch_bytes, encrypt_rlwe_sk_scratch_bytes};
206
207 #[test]
208 fn encrypt_sk_vec_znx_fft64() {
209 let module: Module<FFT64> = Module::<FFT64>::new(32);
210 let log_base2k: usize = 8;
211 let log_k_ct: usize = 54;
212 let log_k_pt: usize = 40;
213
214 let sigma: f64 = 3.2;
215 let bound: f64 = sigma * 6.0;
216
217 let mut ct: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_ct, 2);
218 let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_pt);
219
220 let mut source_xe: Source = Source::new([0u8; 32]);
221 let mut source_xa: Source = Source::new([0u8; 32]);
222
223 let mut scratch: ScratchOwned =
224 ScratchOwned::new(encrypt_rlwe_sk_scratch_bytes(&module, ct.size()) | decrypt_rlwe_scratch_bytes(&module, ct.size()));
225
226 let sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
227 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
228 sk_dft.dft(&module, &sk);
229
230 let mut data_want: Vec<i64> = vec![0i64; module.n()];
231
232 data_want
233 .iter_mut()
234 .for_each(|x| *x = source_xa.next_i64() & 0xFF);
235
236 pt.data
237 .encode_vec_i64(0, log_base2k, log_k_pt, &data_want, 10);
238
239 ct.encrypt_sk(
240 &module,
241 Some(&pt),
242 &sk_dft,
243 &mut source_xa,
244 &mut source_xe,
245 scratch.borrow(),
246 sigma,
247 bound,
248 );
249
250 pt.data.zero();
251
252 ct.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
253
254 let mut data_have: Vec<i64> = vec![0i64; module.n()];
255
256 pt.data
257 .decode_vec_i64(0, log_base2k, pt.size() * log_base2k, &mut data_have);
258
259 let scale: f64 = (1 << (pt.size() * log_base2k - log_k_pt)) as f64;
260 izip!(data_want.iter(), data_have.iter()).for_each(|(a, b)| {
261 let b_scaled = (*b as f64) / scale;
262 assert!(
263 (*a as f64 - b_scaled).abs() < 0.1,
264 "{} {}",
265 *a as f64,
266 b_scaled
267 )
268 });
269
270 module.free();
271 }
272}