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