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