poulpy/core/src/ggsw_ciphertext.rs
core/src/ggsw_ciphertext.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
3 ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut,
4 VecZnxDftToRef, VecZnxOps, ZnxInfos, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::{GetRow, Infos, SetRow},
10 glwe_ciphertext::GLWECiphertext,
11 glwe_ciphertext_fourier::GLWECiphertextFourier,
12 glwe_plaintext::GLWEPlaintext,
13 keys::SecretKeyFourier,
14 utils::derive_size,
15};
16
17pub struct GGSWCiphertext<C, B: Backend> {
18 pub data: MatZnxDft<C, B>,
19 pub basek: usize,
20 pub k: usize,
21}
22
23impl<B: Backend> GGSWCiphertext<Vec<u8>, B> {
24 pub fn new(module: &Module<B>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
25 Self {
26 data: module.new_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k)),
27 basek: basek,
28 k: k,
29 }
30 }
31}
32
33impl<T, B: Backend> Infos for GGSWCiphertext<T, B> {
34 type Inner = MatZnxDft<T, B>;
35
36 fn inner(&self) -> &Self::Inner {
37 &self.data
38 }
39
40 fn basek(&self) -> usize {
41 self.basek
42 }
43
44 fn k(&self) -> usize {
45 self.k
46 }
47}
48
49impl<T, B: Backend> GGSWCiphertext<T, B> {
50 pub fn rank(&self) -> usize {
51 self.data.cols_out() - 1
52 }
53}
54
55impl<C, B: Backend> MatZnxDftToMut<B> for GGSWCiphertext<C, B>
56where
57 MatZnxDft<C, B>: MatZnxDftToMut<B>,
58{
59 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
60 self.data.to_mut()
61 }
62}
63
64impl<C, B: Backend> MatZnxDftToRef<B> for GGSWCiphertext<C, B>
65where
66 MatZnxDft<C, B>: MatZnxDftToRef<B>,
67{
68 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
69 self.data.to_ref()
70 }
71}
72
73impl GGSWCiphertext<Vec<u8>, FFT64> {
74 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
75 GLWECiphertext::encrypt_sk_scratch_space(module, size)
76 + module.bytes_of_vec_znx(rank + 1, size)
77 + module.bytes_of_vec_znx(1, size)
78 + module.bytes_of_vec_znx_dft(rank + 1, size)
79 }
80
81 pub fn external_product_scratch_space(
82 module: &Module<FFT64>,
83 out_size: usize,
84 in_size: usize,
85 ggsw_size: usize,
86 rank: usize,
87 ) -> usize {
88 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
89 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
90 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
91 tmp_in + tmp_out + ggsw
92 }
93
94 pub fn external_product_inplace_scratch_space(
95 module: &Module<FFT64>,
96 out_size: usize,
97 ggsw_size: usize,
98 rank: usize,
99 ) -> usize {
100 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
101 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
102 tmp + ggsw
103 }
104}
105
106impl<DataSelf> GGSWCiphertext<DataSelf, FFT64>
107where
108 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
109{
110 pub fn encrypt_sk<DataPt, DataSk>(
111 &mut self,
112 module: &Module<FFT64>,
113 pt: &ScalarZnx<DataPt>,
114 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
115 source_xa: &mut Source,
116 source_xe: &mut Source,
117 sigma: f64,
118 scratch: &mut Scratch,
119 ) where
120 ScalarZnx<DataPt>: ScalarZnxToRef,
121 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
122 {
123 #[cfg(debug_assertions)]
124 {
125 assert_eq!(self.rank(), sk_dft.rank());
126 assert_eq!(self.n(), module.n());
127 assert_eq!(pt.n(), module.n());
128 assert_eq!(sk_dft.n(), module.n());
129 }
130
131 let size: usize = self.size();
132 let basek: usize = self.basek();
133 let k: usize = self.k();
134 let cols: usize = self.rank() + 1;
135
136 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
137 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, cols, size);
138
139 let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
140 data: tmp_znx_pt,
141 basek: basek,
142 k: k,
143 };
144
145 let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
146 data: tmp_znx_ct,
147 basek: basek,
148 k,
149 };
150
151 (0..self.rows()).for_each(|row_j| {
152 vec_znx_pt.data.zero();
153
154 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
155 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
156 module.vec_znx_normalize_inplace(basek, &mut vec_znx_pt, 0, scrach_2);
157
158 (0..cols).for_each(|col_i| {
159 // rlwe encrypt of vec_znx_pt into vec_znx_ct
160
161 vec_znx_ct.encrypt_sk_private(
162 module,
163 Some((&vec_znx_pt, col_i)),
164 sk_dft,
165 source_xa,
166 source_xe,
167 sigma,
168 scrach_2,
169 );
170
171 // Switch vec_znx_ct into DFT domain
172 {
173 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, cols, size);
174
175 (0..cols).for_each(|i| {
176 module.vec_znx_dft(&mut vec_znx_dft_ct, i, &vec_znx_ct, i);
177 });
178
179 module.vmp_prepare_row(self, row_j, col_i, &vec_znx_dft_ct);
180 }
181 });
182 });
183 }
184
185 pub fn external_product<DataLhs, DataRhs>(
186 &mut self,
187 module: &Module<FFT64>,
188 lhs: &GGSWCiphertext<DataLhs, FFT64>,
189 rhs: &GGSWCiphertext<DataRhs, FFT64>,
190 scratch: &mut Scratch,
191 ) where
192 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
193 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
194 {
195 #[cfg(debug_assertions)]
196 {
197 assert_eq!(
198 self.rank(),
199 lhs.rank(),
200 "ggsw_out rank: {} != ggsw_in rank: {}",
201 self.rank(),
202 lhs.rank()
203 );
204 assert_eq!(
205 self.rank(),
206 rhs.rank(),
207 "ggsw_in rank: {} != ggsw_apply rank: {}",
208 self.rank(),
209 rhs.rank()
210 );
211 }
212
213 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank() + 1, lhs.size());
214
215 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
216 data: tmp_in_data,
217 basek: lhs.basek(),
218 k: lhs.k(),
219 };
220
221 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
222
223 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
224 data: tmp_out_data,
225 basek: self.basek(),
226 k: self.k(),
227 };
228
229 (0..self.rank() + 1).for_each(|col_i| {
230 (0..self.rows()).for_each(|row_j| {
231 lhs.get_row(module, row_j, col_i, &mut tmp_in);
232 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
233 self.set_row(module, row_j, col_i, &tmp_out);
234 });
235 });
236
237 tmp_out.data.zero();
238
239 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
240 (0..self.rank() + 1).for_each(|col_j| {
241 self.set_row(module, row_i, col_j, &tmp_out);
242 });
243 });
244 }
245
246 pub fn external_product_inplace<DataRhs>(
247 &mut self,
248 module: &Module<FFT64>,
249 rhs: &GGSWCiphertext<DataRhs, FFT64>,
250 scratch: &mut Scratch,
251 ) where
252 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
253 {
254 #[cfg(debug_assertions)]
255 {
256 assert_eq!(
257 self.rank(),
258 rhs.rank(),
259 "ggsw_out rank: {} != ggsw_apply: {}",
260 self.rank(),
261 rhs.rank()
262 );
263 }
264
265 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
266
267 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
268 data: tmp_data,
269 basek: self.basek(),
270 k: self.k(),
271 };
272
273 (0..self.rank() + 1).for_each(|col_i| {
274 (0..self.rows()).for_each(|row_j| {
275 self.get_row(module, row_j, col_i, &mut tmp);
276 tmp.external_product_inplace(module, rhs, scratch1);
277 self.set_row(module, row_j, col_i, &tmp);
278 });
279 });
280 }
281}
282
283impl<DataSelf> GetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
284where
285 MatZnxDft<DataSelf, FFT64>: MatZnxDftToRef<FFT64>,
286{
287 fn get_row<DataRes>(
288 &self,
289 module: &Module<FFT64>,
290 row_i: usize,
291 col_j: usize,
292 res: &mut GLWECiphertextFourier<DataRes, FFT64>,
293 ) where
294 VecZnxDft<DataRes, FFT64>: VecZnxDftToMut<FFT64>,
295 {
296 module.vmp_extract_row(res, self, row_i, col_j);
297 }
298}
299
300impl<DataSelf> SetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
301where
302 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
303{
304 fn set_row<DataRes>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<DataRes, FFT64>)
305 where
306 VecZnxDft<DataRes, FFT64>: VecZnxDftToRef<FFT64>,
307 {
308 module.vmp_prepare_row(self, row_i, col_j, a);
309 }
310}