poulpy/core/src/gglwe/encryption.rs
core/src/gglwe/encryption.rsRaw
1use backend::{
2 FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScalarZnxOps, Scratch, VecZnxAlloc, VecZnxDftAlloc, VecZnxOps,
3 ZnxInfos, ZnxView, ZnxViewMut, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{
8 FourierGLWESecret, GGLWECiphertext, GLWEAutomorphismKey, GLWECiphertext, GLWESecret, GLWESwitchingKey, GLWETensorKey, Infos,
9 ScratchCore, SetRow,
10};
11
12impl GGLWECiphertext<Vec<u8>, FFT64> {
13 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
14 let size = k.div_ceil(basek);
15 GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
16 + module.bytes_of_vec_znx(rank + 1, size)
17 + module.bytes_of_vec_znx(1, size)
18 + module.bytes_of_vec_znx_dft(rank + 1, size)
19 }
20
21 pub fn encrypt_pk_scratch_space(_module: &Module<FFT64>, _basek: usize, _k: usize, _rank: usize) -> usize {
22 unimplemented!()
23 }
24}
25
26impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GGLWECiphertext<DataSelf, FFT64> {
27 pub fn encrypt_sk<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
28 &mut self,
29 module: &Module<FFT64>,
30 pt: &ScalarZnx<DataPt>,
31 sk: &FourierGLWESecret<DataSk, FFT64>,
32 source_xa: &mut Source,
33 source_xe: &mut Source,
34 sigma: f64,
35 scratch: &mut Scratch,
36 ) {
37 #[cfg(debug_assertions)]
38 {
39 assert_eq!(
40 self.rank_in(),
41 pt.cols(),
42 "self.rank_in(): {} != pt.cols(): {}",
43 self.rank_in(),
44 pt.cols()
45 );
46 assert_eq!(
47 self.rank_out(),
48 sk.rank(),
49 "self.rank_out(): {} != sk.rank(): {}",
50 self.rank_out(),
51 sk.rank()
52 );
53 assert_eq!(self.n(), module.n());
54 assert_eq!(sk.n(), module.n());
55 assert_eq!(pt.n(), module.n());
56 assert!(
57 scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
58 "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
59 scratch.available(),
60 self.rank(),
61 self.size(),
62 GGLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
63 );
64 assert!(
65 self.rows() * self.digits() * self.basek() <= self.k(),
66 "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
67 self.rows(),
68 self.digits(),
69 self.basek(),
70 self.rows() * self.digits() * self.basek(),
71 self.k()
72 );
73 }
74
75 let rows: usize = self.rows();
76 let digits: usize = self.digits();
77 let basek: usize = self.basek();
78 let k: usize = self.k();
79 let rank_in: usize = self.rank_in();
80 let rank_out: usize = self.rank_out();
81
82 let (mut tmp_pt, scrach_1) = scratch.tmp_glwe_pt(module, basek, k);
83 let (mut tmp_ct, scrach_2) = scrach_1.tmp_glwe_ct(module, basek, k, rank_out);
84 let (mut tmp_ct_dft, scratch_3) = scrach_2.tmp_fourier_glwe_ct(module, basek, k, rank_out);
85
86 // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
87 //
88 // Example for ksk rank 2 to rank 3:
89 //
90 // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
91 // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
92 //
93 // Example ksk rank 2 to rank 1
94 //
95 // (-(a*s) + s0, a)
96 // (-(b*s) + s1, b)
97 (0..rank_in).for_each(|col_i| {
98 (0..rows).for_each(|row_i| {
99 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
100 tmp_pt.data.zero(); // zeroes for next iteration
101 module.vec_znx_add_scalar_inplace(
102 &mut tmp_pt.data,
103 0,
104 (digits - 1) + row_i * digits,
105 pt,
106 col_i,
107 );
108 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch_3);
109
110 // rlwe encrypt of vec_znx_pt into vec_znx_ct
111 tmp_ct.encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scratch_3);
112
113 // Switch vec_znx_ct into DFT domain
114 tmp_ct.dft(module, &mut tmp_ct_dft);
115
116 // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
117 self.set_row(module, row_i, col_i, &tmp_ct_dft);
118 });
119 });
120 }
121}
122
123impl GLWESwitchingKey<Vec<u8>, FFT64> {
124 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank_in: usize, rank_out: usize) -> usize {
125 GGLWECiphertext::encrypt_sk_scratch_space(module, basek, k, rank_out)
126 + module.bytes_of_scalar_znx(rank_in)
127 + FourierGLWESecret::bytes_of(module, rank_out)
128 }
129
130 pub fn encrypt_pk_scratch_space(
131 module: &Module<FFT64>,
132 _basek: usize,
133 _k: usize,
134 _rank_in: usize,
135 _rank_out: usize,
136 ) -> usize {
137 GGLWECiphertext::encrypt_pk_scratch_space(module, _basek, _k, _rank_out)
138 }
139}
140
141impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWESwitchingKey<DataSelf, FFT64> {
142 pub fn encrypt_sk<DataSkIn: AsRef<[u8]>, DataSkOut: AsRef<[u8]>>(
143 &mut self,
144 module: &Module<FFT64>,
145 sk_in: &GLWESecret<DataSkIn>,
146 sk_out: &FourierGLWESecret<DataSkOut, FFT64>,
147 source_xa: &mut Source,
148 source_xe: &mut Source,
149 sigma: f64,
150 scratch: &mut Scratch,
151 ) {
152 #[cfg(debug_assertions)]
153 {
154 assert!(sk_in.n() <= module.n());
155 assert!(sk_out.n() <= module.n());
156 }
157
158 let (mut sk_in_tmp, scratch1) = scratch.tmp_scalar_znx(module, sk_in.rank());
159 sk_in_tmp.zero();
160
161 (0..sk_in.rank()).for_each(|i| {
162 sk_in_tmp
163 .at_mut(i, 0)
164 .iter_mut()
165 .step_by(module.n() / sk_in.n())
166 .zip(sk_in.data.at(i, 0).iter())
167 .for_each(|(x, y)| *x = *y);
168 });
169
170 let (mut sk_out_tmp, scratch2) = scratch1.tmp_fourier_glwe_secret(module, sk_out.rank());
171 (0..sk_out.rank()).for_each(|i| {
172 sk_out_tmp
173 .data
174 .at_mut(i, 0)
175 .chunks_exact_mut(module.n() / sk_out.n())
176 .zip(sk_out.data.at(i, 0).iter())
177 .for_each(|(a_chunk, &b_elem)| {
178 a_chunk.fill(b_elem);
179 });
180 });
181
182 self.key.encrypt_sk(
183 module,
184 &sk_in_tmp,
185 &sk_out_tmp,
186 source_xa,
187 source_xe,
188 sigma,
189 scratch2,
190 );
191 self.sk_in_n = sk_in.n();
192 self.sk_out_n = sk_out.n();
193 }
194}
195
196impl GLWEAutomorphismKey<Vec<u8>, FFT64> {
197 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
198 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank) + GLWESecret::bytes_of(module, rank)
199 }
200
201 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, _basek: usize, _k: usize, _rank: usize) -> usize {
202 GLWESwitchingKey::encrypt_pk_scratch_space(module, _basek, _k, _rank, _rank)
203 }
204}
205
206impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWEAutomorphismKey<DataSelf, FFT64> {
207 pub fn encrypt_sk<DataSk: AsRef<[u8]>>(
208 &mut self,
209 module: &Module<FFT64>,
210 p: i64,
211 sk: &GLWESecret<DataSk>,
212 source_xa: &mut Source,
213 source_xe: &mut Source,
214 sigma: f64,
215 scratch: &mut Scratch,
216 ) {
217 #[cfg(debug_assertions)]
218 {
219 assert_eq!(self.n(), module.n());
220 assert_eq!(sk.n(), module.n());
221 assert_eq!(self.rank_out(), self.rank_in());
222 assert_eq!(sk.rank(), self.rank());
223 assert!(
224 scratch.available() >= GLWEAutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
225 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
226 scratch.available(),
227 self.rank(),
228 self.size(),
229 GLWEAutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
230 )
231 }
232
233 let (mut sk_out_dft, scratch_1) = scratch.tmp_fourier_glwe_secret(module, sk.rank());
234
235 {
236 let (mut sk_out, _) = scratch_1.tmp_glwe_secret(module, sk.rank());
237 (0..self.rank()).for_each(|i| {
238 module.scalar_znx_automorphism(
239 module.galois_element_inv(p),
240 &mut sk_out.data,
241 i,
242 &sk.data,
243 i,
244 );
245 });
246 sk_out_dft.set(module, &sk_out);
247 }
248
249 self.key.encrypt_sk(
250 module,
251 &sk,
252 &sk_out_dft,
253 source_xa,
254 source_xe,
255 sigma,
256 scratch_1,
257 );
258
259 self.p = p;
260 }
261}
262
263impl GLWETensorKey<Vec<u8>, FFT64> {
264 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
265 GLWESecret::bytes_of(module, 1)
266 + FourierGLWESecret::bytes_of(module, 1)
267 + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank)
268 }
269}
270
271impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWETensorKey<DataSelf, FFT64> {
272 pub fn encrypt_sk<DataSk: AsRef<[u8]>>(
273 &mut self,
274 module: &Module<FFT64>,
275 sk: &FourierGLWESecret<DataSk, FFT64>,
276 source_xa: &mut Source,
277 source_xe: &mut Source,
278 sigma: f64,
279 scratch: &mut Scratch,
280 ) {
281 #[cfg(debug_assertions)]
282 {
283 assert_eq!(self.rank(), sk.rank());
284 assert_eq!(self.n(), module.n());
285 assert_eq!(sk.n(), module.n());
286 }
287
288 let rank: usize = self.rank();
289
290 let (mut sk_ij, scratch1) = scratch.tmp_glwe_secret(module, 1);
291 let (mut sk_ij_dft, scratch2) = scratch1.tmp_fourier_glwe_secret(module, 1);
292
293 (0..rank).for_each(|i| {
294 (i..rank).for_each(|j| {
295 module.svp_apply(&mut sk_ij_dft.data, 0, &sk.data, i, &sk.data, j);
296 module.scalar_znx_idft(&mut sk_ij.data, 0, &sk_ij_dft.data, 0, scratch2);
297 self.at_mut(i, j)
298 .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch2);
299 });
300 })
301 }
302}