poulpy/core/src/glwe_ciphertext_fourier.rs
core/src/glwe_ciphertext_fourier.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftScratch, MatZnxDftToRef, Module, ScalarZnxDft, ScalarZnxDftOps,
3 ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft,
4 VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::Infos, ggsw_ciphertext::GGSWCiphertext, glwe_ciphertext::GLWECiphertext, glwe_plaintext::GLWEPlaintext,
10 keys::SecretKeyFourier, keyswitch_key::GLWESwitchingKey, utils::derive_size,
11};
12
13pub struct GLWECiphertextFourier<C, B: Backend> {
14 pub data: VecZnxDft<C, B>,
15 pub basek: usize,
16 pub k: usize,
17}
18
19impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
20 pub fn new(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
21 Self {
22 data: module.new_vec_znx_dft(rank + 1, derive_size(basek, k)),
23 basek: basek,
24 k: k,
25 }
26 }
27}
28
29impl<T, B: Backend> Infos for GLWECiphertextFourier<T, B> {
30 type Inner = VecZnxDft<T, B>;
31
32 fn inner(&self) -> &Self::Inner {
33 &self.data
34 }
35
36 fn basek(&self) -> usize {
37 self.basek
38 }
39
40 fn k(&self) -> usize {
41 self.k
42 }
43}
44
45impl<T, B: Backend> GLWECiphertextFourier<T, B> {
46 pub fn rank(&self) -> usize {
47 self.cols() - 1
48 }
49}
50
51impl<C, B: Backend> VecZnxDftToMut<B> for GLWECiphertextFourier<C, B>
52where
53 VecZnxDft<C, B>: VecZnxDftToMut<B>,
54{
55 fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
56 self.data.to_mut()
57 }
58}
59
60impl<C, B: Backend> VecZnxDftToRef<B> for GLWECiphertextFourier<C, B>
61where
62 VecZnxDft<C, B>: VecZnxDftToRef<B>,
63{
64 fn to_ref(&self) -> VecZnxDft<&[u8], B> {
65 self.data.to_ref()
66 }
67}
68
69impl GLWECiphertextFourier<Vec<u8>, FFT64> {
70 #[allow(dead_code)]
71 pub(crate) fn idft_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
72 module.bytes_of_vec_znx(1, size) + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
73 }
74
75 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, ct_size: usize) -> usize {
76 module.bytes_of_vec_znx(rank + 1, ct_size) + GLWECiphertext::encrypt_sk_scratch_space(module, ct_size)
77 }
78
79 pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
80 (module.vec_znx_big_normalize_tmp_bytes()
81 | module.bytes_of_vec_znx_dft(1, ct_size)
82 | (module.bytes_of_vec_znx_big(1, ct_size) + module.vec_znx_idft_tmp_bytes()))
83 + module.bytes_of_vec_znx_big(1, ct_size)
84 }
85
86 pub fn keyswitch_scratch_space(
87 module: &Module<FFT64>,
88 out_size: usize,
89 out_rank: usize,
90 in_size: usize,
91 in_rank: usize,
92 ksk_size: usize,
93 ) -> usize {
94 module.bytes_of_vec_znx(out_rank + 1, out_size)
95 + GLWECiphertext::keyswitch_from_fourier_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size)
96 }
97
98 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
99 Self::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
100 }
101
102 pub fn external_product_scratch_space(
103 module: &Module<FFT64>,
104 out_size: usize,
105 in_size: usize,
106 ggsw_size: usize,
107 rank: usize,
108 ) -> usize {
109 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
110 let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, rank + 1, rank + 1, ggsw_size);
111 let res_small: usize = module.bytes_of_vec_znx(rank + 1, out_size);
112 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
113
114 res_dft + (vmp | (res_small + normalize))
115 }
116
117 pub fn external_product_inplace_scratch_space(
118 module: &Module<FFT64>,
119 out_size: usize,
120 ggsw_size: usize,
121 rank: usize,
122 ) -> usize {
123 Self::external_product_scratch_space(module, out_size, out_size, ggsw_size, rank)
124 }
125}
126
127impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
128where
129 VecZnxDft<DataSelf, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
130{
131 pub fn encrypt_zero_sk<DataSk>(
132 &mut self,
133 module: &Module<FFT64>,
134 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
135 source_xa: &mut Source,
136 source_xe: &mut Source,
137 sigma: f64,
138 scratch: &mut Scratch,
139 ) where
140 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
141 {
142 let (vec_znx_tmp, scratch_1) = scratch.tmp_vec_znx(module, self.rank() + 1, self.size());
143 let mut ct_idft = GLWECiphertext {
144 data: vec_znx_tmp,
145 basek: self.basek,
146 k: self.k,
147 };
148 ct_idft.encrypt_zero_sk(module, sk_dft, source_xa, source_xe, sigma, scratch_1);
149
150 ct_idft.dft(module, self);
151 }
152
153 pub fn keyswitch<DataLhs, DataRhs>(
154 &mut self,
155 module: &Module<FFT64>,
156 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
157 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
158 scratch: &mut Scratch,
159 ) where
160 VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
161 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
162 {
163 let cols_out: usize = rhs.rank_out() + 1;
164
165 // Space fr normalized VMP result outside of DFT domain
166 let (res_idft_data, scratch1) = scratch.tmp_vec_znx(module, cols_out, lhs.size());
167
168 let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
169 data: res_idft_data,
170 basek: lhs.basek,
171 k: lhs.k,
172 };
173
174 res_idft.keyswitch_from_fourier(module, lhs, rhs, scratch1);
175
176 (0..cols_out).for_each(|i| {
177 module.vec_znx_dft(self, i, &res_idft, i);
178 });
179 }
180
181 pub fn keyswitch_inplace<DataRhs>(
182 &mut self,
183 module: &Module<FFT64>,
184 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
185 scratch: &mut Scratch,
186 ) where
187 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
188 {
189 unsafe {
190 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
191 self.keyswitch(&module, &*self_ptr, rhs, scratch);
192 }
193 }
194
195 pub fn external_product<DataLhs, DataRhs>(
196 &mut self,
197 module: &Module<FFT64>,
198 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
199 rhs: &GGSWCiphertext<DataRhs, FFT64>,
200 scratch: &mut Scratch,
201 ) where
202 VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
203 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
204 {
205 let basek: usize = self.basek();
206
207 #[cfg(debug_assertions)]
208 {
209 assert_eq!(rhs.rank(), lhs.rank());
210 assert_eq!(rhs.rank(), self.rank());
211 assert_eq!(self.basek(), basek);
212 assert_eq!(lhs.basek(), basek);
213 assert_eq!(rhs.n(), module.n());
214 assert_eq!(self.n(), module.n());
215 assert_eq!(lhs.n(), module.n());
216 }
217
218 let cols: usize = rhs.rank() + 1;
219
220 // Space for VMP result in DFT domain and high precision
221 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
222
223 {
224 module.vmp_apply(&mut res_dft, lhs, rhs, scratch1);
225 }
226
227 // VMP result in high precision
228 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
229
230 // Space for VMP result normalized
231 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
232 (0..cols).for_each(|i| {
233 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
234 module.vec_znx_dft(self, i, &res_small, i);
235 });
236 }
237
238 pub fn external_product_inplace<DataRhs>(
239 &mut self,
240 module: &Module<FFT64>,
241 rhs: &GGSWCiphertext<DataRhs, FFT64>,
242 scratch: &mut Scratch,
243 ) where
244 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
245 {
246 unsafe {
247 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
248 self.external_product(&module, &*self_ptr, rhs, scratch);
249 }
250 }
251}
252
253impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
254where
255 VecZnxDft<DataSelf, FFT64>: VecZnxDftToRef<FFT64>,
256{
257 pub fn decrypt<DataPt, DataSk>(
258 &self,
259 module: &Module<FFT64>,
260 pt: &mut GLWEPlaintext<DataPt>,
261 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
262 scratch: &mut Scratch,
263 ) where
264 VecZnx<DataPt>: VecZnxToMut + VecZnxToRef,
265 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
266 {
267 #[cfg(debug_assertions)]
268 {
269 assert_eq!(self.rank(), sk_dft.rank());
270 assert_eq!(self.n(), module.n());
271 assert_eq!(pt.n(), module.n());
272 assert_eq!(sk_dft.n(), module.n());
273 }
274
275 let cols = self.rank() + 1;
276
277 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
278 pt_big.zero();
279
280 {
281 (1..cols).for_each(|i| {
282 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
283 module.svp_apply(&mut ci_dft, 0, sk_dft, i - 1, self, i);
284 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
285 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
286 });
287 }
288
289 {
290 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
291 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
292 module.vec_znx_idft(&mut c0_big, 0, self, 0, scratch_2);
293 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
294 }
295
296 // pt = norm(BIG(m + e))
297 module.vec_znx_big_normalize(self.basek(), pt, 0, &mut pt_big, 0, scratch_1);
298
299 pt.basek = self.basek();
300 pt.k = pt.k().min(self.k());
301 }
302
303 #[allow(dead_code)]
304 pub(crate) fn idft<DataRes>(&self, module: &Module<FFT64>, res: &mut GLWECiphertext<DataRes>, scratch: &mut Scratch)
305 where
306 GLWECiphertext<DataRes>: VecZnxToMut,
307 {
308 #[cfg(debug_assertions)]
309 {
310 assert_eq!(self.rank(), res.rank());
311 assert_eq!(self.basek(), res.basek())
312 }
313
314 let min_size: usize = self.size().min(res.size());
315
316 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
317
318 (0..self.rank() + 1).for_each(|i| {
319 module.vec_znx_idft(&mut res_big, 0, self, i, scratch1);
320 module.vec_znx_big_normalize(self.basek(), res, i, &res_big, 0, scratch1);
321 });
322 }
323}