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