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, rank, 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(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
86 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
87 module, res_size, lhs, rhs,
88 )
89 }
90
91 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
92 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_inplace_scratch_space(
93 module, res_size, rhs,
94 )
95 }
96
97 pub fn external_product_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
98 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
99 module, res_size, lhs, rhs,
100 )
101 }
102
103 pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
104 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
105 module, res_size, rhs,
106 )
107 }
108}
109
110impl<DataSelf> GGSWCiphertext<DataSelf, FFT64>
111where
112 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
113{
114 pub fn encrypt_sk<DataPt, DataSk>(
115 &mut self,
116 module: &Module<FFT64>,
117 pt: &ScalarZnx<DataPt>,
118 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
119 source_xa: &mut Source,
120 source_xe: &mut Source,
121 sigma: f64,
122 bound: f64,
123 scratch: &mut Scratch,
124 ) where
125 ScalarZnx<DataPt>: ScalarZnxToRef,
126 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
127 {
128 #[cfg(debug_assertions)]
129 {
130 assert_eq!(self.rank(), sk_dft.rank());
131 assert_eq!(self.n(), module.n());
132 assert_eq!(pt.n(), module.n());
133 assert_eq!(sk_dft.n(), module.n());
134 }
135
136 let size: usize = self.size();
137 let log_base2k: usize = self.basek();
138 let k: usize = self.k();
139 let cols: usize = self.rank() + 1;
140
141 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
142 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, cols, size);
143
144 let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
145 data: tmp_znx_pt,
146 basek: log_base2k,
147 k: k,
148 };
149
150 let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
151 data: tmp_znx_ct,
152 basek: log_base2k,
153 k,
154 };
155
156 (0..self.rows()).for_each(|row_j| {
157 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
158 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
159 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scrach_2);
160
161 (0..cols).for_each(|col_i| {
162 // rlwe encrypt of vec_znx_pt into vec_znx_ct
163
164 vec_znx_ct.encrypt_sk_private(
165 module,
166 Some((&vec_znx_pt, col_i)),
167 sk_dft,
168 source_xa,
169 source_xe,
170 sigma,
171 bound,
172 scrach_2,
173 );
174
175 // Switch vec_znx_ct into DFT domain
176 {
177 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, cols, size);
178
179 (0..cols).for_each(|i| {
180 module.vec_znx_dft(&mut vec_znx_dft_ct, i, &vec_znx_ct, i);
181 });
182
183 module.vmp_prepare_row(self, row_j, col_i, &vec_znx_dft_ct);
184 }
185 });
186
187 vec_znx_pt.data.zero(); // zeroes for next iteration
188 });
189 }
190
191 pub fn keyswitch<DataLhs, DataRhs>(
192 &mut self,
193 module: &Module<FFT64>,
194 lhs: &GGSWCiphertext<DataLhs, FFT64>,
195 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
196 scratch: &mut Scratch,
197 ) where
198 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
199 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
200 {
201 rhs.0.prod_with_vec_glwe(module, self, lhs, scratch);
202 }
203
204 pub fn keyswitch_inplace<DataRhs>(
205 &mut self,
206 module: &Module<FFT64>,
207 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
208 scratch: &mut Scratch,
209 ) where
210 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
211 {
212 rhs.0.prod_with_vec_glwe_inplace(module, self, scratch);
213 }
214
215 pub fn external_product<DataLhs, DataRhs>(
216 &mut self,
217 module: &Module<FFT64>,
218 lhs: &GGSWCiphertext<DataLhs, FFT64>,
219 rhs: &GGSWCiphertext<DataRhs, FFT64>,
220 scratch: &mut Scratch,
221 ) where
222 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
223 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
224 {
225 rhs.prod_with_vec_glwe(module, self, lhs, scratch);
226 }
227
228 pub fn external_product_inplace<DataRhs>(
229 &mut self,
230 module: &Module<FFT64>,
231 rhs: &GGSWCiphertext<DataRhs, FFT64>,
232 scratch: &mut Scratch,
233 ) where
234 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
235 {
236 rhs.prod_with_vec_glwe_inplace(module, self, scratch);
237 }
238}
239
240impl<DataSelf> GetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
241where
242 MatZnxDft<DataSelf, FFT64>: MatZnxDftToRef<FFT64>,
243{
244 fn get_row<DataRes>(
245 &self,
246 module: &Module<FFT64>,
247 row_i: usize,
248 col_j: usize,
249 res: &mut GLWECiphertextFourier<DataRes, FFT64>,
250 ) where
251 VecZnxDft<DataRes, FFT64>: VecZnxDftToMut<FFT64>,
252 {
253 module.vmp_extract_row(res, self, row_i, col_j);
254 }
255}
256
257impl<DataSelf> SetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
258where
259 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
260{
261 fn set_row<DataRes>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<DataRes, FFT64>)
262 where
263 VecZnxDft<DataRes, FFT64>: VecZnxDftToRef<FFT64>,
264 {
265 module.vmp_prepare_row(self, row_i, col_j, a);
266 }
267}
268
269impl VecGLWEProductScratchSpace for GGSWCiphertext<Vec<u8>, FFT64> {
270 fn prod_with_glwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, rgsw_size: usize) -> usize {
271 module.bytes_of_vec_znx_dft(2, rgsw_size)
272 + ((module.bytes_of_vec_znx_dft(2, a_size) + module.vmp_apply_tmp_bytes(res_size, a_size, a_size, 2, 2, rgsw_size))
273 | module.vec_znx_big_normalize_tmp_bytes())
274 }
275}
276
277impl<C> VecGLWEProduct for GGSWCiphertext<C, FFT64>
278where
279 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
280{
281 fn prod_with_glwe<R, A>(
282 &self,
283 module: &Module<FFT64>,
284 res: &mut GLWECiphertext<R>,
285 a: &GLWECiphertext<A>,
286 scratch: &mut Scratch,
287 ) where
288 VecZnx<R>: VecZnxToMut,
289 VecZnx<A>: VecZnxToRef,
290 {
291 let log_base2k: usize = self.basek();
292
293 #[cfg(debug_assertions)]
294 {
295 assert_eq!(res.basek(), log_base2k);
296 assert_eq!(a.basek(), log_base2k);
297 assert_eq!(self.n(), module.n());
298 assert_eq!(res.n(), module.n());
299 assert_eq!(a.n(), module.n());
300 }
301
302 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
303
304 {
305 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, a.size());
306 module.vec_znx_dft(&mut a_dft, 0, a, 0);
307 module.vec_znx_dft(&mut a_dft, 1, a, 1);
308 module.vmp_apply(&mut res_dft, &a_dft, self, scratch2);
309 }
310
311 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
312
313 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
314 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
315 }
316}