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