poulpy/core/src/ggsw_ciphertext.rs
core/src/ggsw_ciphertext.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
3 ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut,
4 VecZnxDftToRef, VecZnxOps, ZnxInfos, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 automorphism::AutomorphismKey,
10 elem::{GetRow, Infos, SetRow},
11 glwe_ciphertext::GLWECiphertext,
12 glwe_ciphertext_fourier::GLWECiphertextFourier,
13 glwe_plaintext::GLWEPlaintext,
14 keys::SecretKeyFourier,
15 utils::derive_size,
16};
17
18pub struct GGSWCiphertext<C, B: Backend> {
19 pub data: MatZnxDft<C, B>,
20 pub basek: usize,
21 pub k: usize,
22}
23
24impl<B: Backend> GGSWCiphertext<Vec<u8>, B> {
25 pub fn new(module: &Module<B>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
26 Self {
27 data: module.new_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k)),
28 basek: basek,
29 k: k,
30 }
31 }
32}
33
34impl<T, B: Backend> Infos for GGSWCiphertext<T, B> {
35 type Inner = MatZnxDft<T, B>;
36
37 fn inner(&self) -> &Self::Inner {
38 &self.data
39 }
40
41 fn basek(&self) -> usize {
42 self.basek
43 }
44
45 fn k(&self) -> usize {
46 self.k
47 }
48}
49
50impl<T, B: Backend> GGSWCiphertext<T, B> {
51 pub fn rank(&self) -> usize {
52 self.data.cols_out() - 1
53 }
54}
55
56impl<C, B: Backend> MatZnxDftToMut<B> for GGSWCiphertext<C, B>
57where
58 MatZnxDft<C, B>: MatZnxDftToMut<B>,
59{
60 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
61 self.data.to_mut()
62 }
63}
64
65impl<C, B: Backend> MatZnxDftToRef<B> for GGSWCiphertext<C, B>
66where
67 MatZnxDft<C, B>: MatZnxDftToRef<B>,
68{
69 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
70 self.data.to_ref()
71 }
72}
73
74impl GGSWCiphertext<Vec<u8>, FFT64> {
75 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
76 GLWECiphertext::encrypt_sk_scratch_space(module, size)
77 + module.bytes_of_vec_znx(rank + 1, size)
78 + module.bytes_of_vec_znx(1, size)
79 + module.bytes_of_vec_znx_dft(rank + 1, size)
80 }
81
82 pub fn automorphism_scratch_space(
83 module: &Module<FFT64>,
84 out_size: usize,
85 in_size: usize,
86 auto_key_size: usize,
87 rank: usize,
88 ) -> usize {
89 let size: usize = in_size.min(out_size);
90 let tmp_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, size);
91 let tmp_idft: usize = module.bytes_of_vec_znx(rank + 1, size);
92 let vmp: usize = GLWECiphertext::keyswitch_from_fourier_scratch_space(module, size, rank, size, rank, auto_key_size);
93 tmp_dft + tmp_idft + vmp
94 }
95
96 pub fn external_product_scratch_space(
97 module: &Module<FFT64>,
98 out_size: usize,
99 in_size: usize,
100 ggsw_size: usize,
101 rank: usize,
102 ) -> usize {
103 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
104 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
105 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
106 tmp_in + tmp_out + ggsw
107 }
108
109 pub fn external_product_inplace_scratch_space(
110 module: &Module<FFT64>,
111 out_size: usize,
112 ggsw_size: usize,
113 rank: usize,
114 ) -> usize {
115 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
116 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
117 tmp + ggsw
118 }
119}
120
121impl<DataSelf> GGSWCiphertext<DataSelf, FFT64>
122where
123 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
124{
125 pub fn encrypt_sk<DataPt, DataSk>(
126 &mut self,
127 module: &Module<FFT64>,
128 pt: &ScalarZnx<DataPt>,
129 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
130 source_xa: &mut Source,
131 source_xe: &mut Source,
132 sigma: f64,
133 scratch: &mut Scratch,
134 ) where
135 ScalarZnx<DataPt>: ScalarZnxToRef,
136 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
137 {
138 #[cfg(debug_assertions)]
139 {
140 assert_eq!(self.rank(), sk_dft.rank());
141 assert_eq!(self.n(), module.n());
142 assert_eq!(pt.n(), module.n());
143 assert_eq!(sk_dft.n(), module.n());
144 }
145
146 let size: usize = self.size();
147 let basek: usize = self.basek();
148 let k: usize = self.k();
149 let cols: usize = self.rank() + 1;
150
151 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
152 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, cols, size);
153
154 let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
155 data: tmp_znx_pt,
156 basek: basek,
157 k: k,
158 };
159
160 let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
161 data: tmp_znx_ct,
162 basek: basek,
163 k,
164 };
165
166 (0..self.rows()).for_each(|row_j| {
167 vec_znx_pt.data.zero();
168
169 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
170 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
171 module.vec_znx_normalize_inplace(basek, &mut vec_znx_pt, 0, scrach_2);
172
173 (0..cols).for_each(|col_i| {
174 // rlwe encrypt of vec_znx_pt into vec_znx_ct
175
176 vec_znx_ct.encrypt_sk_private(
177 module,
178 Some((&vec_znx_pt, col_i)),
179 sk_dft,
180 source_xa,
181 source_xe,
182 sigma,
183 scrach_2,
184 );
185
186 // Switch vec_znx_ct into DFT domain
187 {
188 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, cols, size);
189
190 (0..cols).for_each(|i| {
191 module.vec_znx_dft(&mut vec_znx_dft_ct, i, &vec_znx_ct, i);
192 });
193
194 module.vmp_prepare_row(self, row_j, col_i, &vec_znx_dft_ct);
195 }
196 });
197 });
198 }
199
200 pub fn automorphism<DataLhs, DataRhs>(
201 &mut self,
202 module: &Module<FFT64>,
203 lhs: &GGSWCiphertext<DataLhs, FFT64>,
204 rhs: &AutomorphismKey<DataRhs, FFT64>,
205 scratch: &mut Scratch,
206 ) where
207 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
208 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
209 {
210 #[cfg(debug_assertions)]
211 {
212 assert_eq!(
213 self.rank(),
214 lhs.rank(),
215 "ggsw_out rank: {} != ggsw_in rank: {}",
216 self.rank(),
217 lhs.rank()
218 );
219 assert_eq!(
220 self.rank(),
221 rhs.rank(),
222 "ggsw_in rank: {} != auto_key rank: {}",
223 self.rank(),
224 rhs.rank()
225 );
226 }
227
228 let size: usize = self.size().min(lhs.size());
229 let cols: usize = self.rank() + 1;
230
231 let (tmp_dft_data, scratch1) = scratch.tmp_vec_znx_dft(module, cols, size);
232
233 let mut tmp_dft: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
234 data: tmp_dft_data,
235 basek: lhs.basek(),
236 k: lhs.k(),
237 };
238
239 let (tmp_idft_data, scratch2) = scratch1.tmp_vec_znx(module, cols, size);
240
241 let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
242 data: tmp_idft_data,
243 basek: self.basek(),
244 k: self.k(),
245 };
246
247 (0..cols).for_each(|col_i| {
248 (0..self.rows()).for_each(|row_j| {
249 lhs.get_row(module, row_j, col_i, &mut tmp_dft);
250 tmp_idft.keyswitch_from_fourier(module, &tmp_dft, &rhs.key, scratch2);
251 (0..cols).for_each(|i| {
252 module.vec_znx_automorphism_inplace(rhs.p(), &mut tmp_idft, i);
253 });
254 self.set_row(module, row_j, col_i, &tmp_dft);
255 });
256 });
257
258 tmp_dft.data.zero();
259
260 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
261 (0..self.rank() + 1).for_each(|col_j| {
262 self.set_row(module, row_i, col_j, &tmp_dft);
263 });
264 });
265 }
266
267 pub fn external_product<DataLhs, DataRhs>(
268 &mut self,
269 module: &Module<FFT64>,
270 lhs: &GGSWCiphertext<DataLhs, FFT64>,
271 rhs: &GGSWCiphertext<DataRhs, FFT64>,
272 scratch: &mut Scratch,
273 ) where
274 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
275 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
276 {
277 #[cfg(debug_assertions)]
278 {
279 assert_eq!(
280 self.rank(),
281 lhs.rank(),
282 "ggsw_out rank: {} != ggsw_in rank: {}",
283 self.rank(),
284 lhs.rank()
285 );
286 assert_eq!(
287 self.rank(),
288 rhs.rank(),
289 "ggsw_in rank: {} != ggsw_apply rank: {}",
290 self.rank(),
291 rhs.rank()
292 );
293 }
294
295 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank() + 1, lhs.size());
296
297 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
298 data: tmp_in_data,
299 basek: lhs.basek(),
300 k: lhs.k(),
301 };
302
303 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
304
305 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
306 data: tmp_out_data,
307 basek: self.basek(),
308 k: self.k(),
309 };
310
311 (0..self.rank() + 1).for_each(|col_i| {
312 (0..self.rows()).for_each(|row_j| {
313 lhs.get_row(module, row_j, col_i, &mut tmp_in);
314 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
315 self.set_row(module, row_j, col_i, &tmp_out);
316 });
317 });
318
319 tmp_out.data.zero();
320
321 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
322 (0..self.rank() + 1).for_each(|col_j| {
323 self.set_row(module, row_i, col_j, &tmp_out);
324 });
325 });
326 }
327
328 pub fn external_product_inplace<DataRhs>(
329 &mut self,
330 module: &Module<FFT64>,
331 rhs: &GGSWCiphertext<DataRhs, FFT64>,
332 scratch: &mut Scratch,
333 ) where
334 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
335 {
336 #[cfg(debug_assertions)]
337 {
338 assert_eq!(
339 self.rank(),
340 rhs.rank(),
341 "ggsw_out rank: {} != ggsw_apply: {}",
342 self.rank(),
343 rhs.rank()
344 );
345 }
346
347 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
348
349 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
350 data: tmp_data,
351 basek: self.basek(),
352 k: self.k(),
353 };
354
355 (0..self.rank() + 1).for_each(|col_i| {
356 (0..self.rows()).for_each(|row_j| {
357 self.get_row(module, row_j, col_i, &mut tmp);
358 tmp.external_product_inplace(module, rhs, scratch1);
359 self.set_row(module, row_j, col_i, &tmp);
360 });
361 });
362 }
363}
364
365impl<DataSelf> GetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
366where
367 MatZnxDft<DataSelf, FFT64>: MatZnxDftToRef<FFT64>,
368{
369 fn get_row<DataRes>(
370 &self,
371 module: &Module<FFT64>,
372 row_i: usize,
373 col_j: usize,
374 res: &mut GLWECiphertextFourier<DataRes, FFT64>,
375 ) where
376 VecZnxDft<DataRes, FFT64>: VecZnxDftToMut<FFT64>,
377 {
378 module.vmp_extract_row(res, self, row_i, col_j);
379 }
380}
381
382impl<DataSelf> SetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
383where
384 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
385{
386 fn set_row<DataRes>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<DataRes, FFT64>)
387 where
388 VecZnxDft<DataRes, FFT64>: VecZnxDftToRef<FFT64>,
389 {
390 module.vmp_prepare_row(self, row_i, col_j, a);
391 }
392}