poulpy/core/src/grlwe.rs
core/src/grlwe.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::{FromProdBy, FromProdByScratchSpace, GetRow, Infos, MatZnxDftProducts, ProdBy, ProdByScratchSpace, SetRow},
11 keys::SecretKeyDft,
12 rgsw::RGSWCt,
13 rlwe::{RLWECt, RLWECtDft, RLWEPt},
14 utils::derive_size,
15};
16
17pub struct GRLWECt<C, B: Backend> {
18 pub data: MatZnxDft<C, B>,
19 pub log_base2k: usize,
20 pub log_k: usize,
21}
22
23impl<B: Backend> GRLWECt<Vec<u8>, B> {
24 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
25 Self {
26 data: module.new_mat_znx_dft(rows, 1, 2, derive_size(log_base2k, log_k)),
27 log_base2k: log_base2k,
28 log_k: log_k,
29 }
30 }
31}
32
33impl<C> GRLWECt<C, FFT64>
34where
35 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
36{
37 pub fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, res: &mut RLWECtDft<R, FFT64>)
38 where
39 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
40 {
41 module.vmp_extract_row(res, self, row_i, 0);
42 }
43}
44
45impl<T, B: Backend> Infos for GRLWECt<T, B> {
46 type Inner = MatZnxDft<T, B>;
47
48 fn inner(&self) -> &Self::Inner {
49 &self.data
50 }
51
52 fn log_base2k(&self) -> usize {
53 self.log_base2k
54 }
55
56 fn log_k(&self) -> usize {
57 self.log_k
58 }
59}
60
61impl<C, B: Backend> MatZnxDftToMut<B> for GRLWECt<C, B>
62where
63 MatZnxDft<C, B>: MatZnxDftToMut<B>,
64{
65 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
66 self.data.to_mut()
67 }
68}
69
70impl<C, B: Backend> MatZnxDftToRef<B> for GRLWECt<C, B>
71where
72 MatZnxDft<C, B>: MatZnxDftToRef<B>,
73{
74 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
75 self.data.to_ref()
76 }
77}
78
79impl GRLWECt<Vec<u8>, FFT64> {
80 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
81 RLWECt::encrypt_sk_scratch_space(module, size)
82 + module.bytes_of_vec_znx(2, size)
83 + module.bytes_of_vec_znx(1, size)
84 + module.bytes_of_vec_znx_dft(2, size)
85 }
86}
87
88pub fn encrypt_grlwe_sk<C, P, S>(
89 module: &Module<FFT64>,
90 ct: &mut GRLWECt<C, FFT64>,
91 pt: &ScalarZnx<P>,
92 sk_dft: &SecretKeyDft<S, FFT64>,
93 source_xa: &mut Source,
94 source_xe: &mut Source,
95 sigma: f64,
96 bound: f64,
97 scratch: &mut Scratch,
98) where
99 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
100 ScalarZnx<P>: ScalarZnxToRef,
101 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
102{
103 let rows: usize = ct.rows();
104 let size: usize = ct.size();
105 let log_base2k: usize = ct.log_base2k();
106
107 let (tmp_znx_pt, scrach_1) = scratch.tmp_vec_znx(module, 1, size);
108 let (tmp_znx_ct, scrach_2) = scrach_1.tmp_vec_znx(module, 2, size);
109 let (mut vec_znx_dft_ct, scratch_3) = scrach_2.tmp_vec_znx_dft(module, 2, size);
110
111 let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
112 data: tmp_znx_pt,
113 log_base2k: log_base2k,
114 log_k: ct.log_k(),
115 };
116
117 let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
118 data: tmp_znx_ct,
119 log_base2k: log_base2k,
120 log_k: ct.log_k(),
121 };
122
123 (0..rows).for_each(|row_i| {
124 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
125 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_i, pt, 0);
126 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scratch_3);
127
128 // rlwe encrypt of vec_znx_pt into vec_znx_ct
129 vec_znx_ct.encrypt_sk(
130 module,
131 Some(&vec_znx_pt),
132 sk_dft,
133 source_xa,
134 source_xe,
135 sigma,
136 bound,
137 scratch_3,
138 );
139
140 vec_znx_pt.data.zero(); // zeroes for next iteration
141
142 // Switch vec_znx_ct into DFT domain
143 module.vec_znx_dft(&mut vec_znx_dft_ct, 0, &vec_znx_ct, 0);
144 module.vec_znx_dft(&mut vec_znx_dft_ct, 1, &vec_znx_ct, 1);
145
146 // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
147 module.vmp_prepare_row(ct, row_i, 0, &vec_znx_dft_ct);
148 });
149}
150
151impl<C> GRLWECt<C, FFT64> {
152 pub fn encrypt_sk<P, S>(
153 &mut self,
154 module: &Module<FFT64>,
155 pt: &ScalarZnx<P>,
156 sk_dft: &SecretKeyDft<S, FFT64>,
157 source_xa: &mut Source,
158 source_xe: &mut Source,
159 sigma: f64,
160 bound: f64,
161 scratch: &mut Scratch,
162 ) where
163 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
164 ScalarZnx<P>: ScalarZnxToRef,
165 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
166 {
167 encrypt_grlwe_sk(
168 module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
169 )
170 }
171}
172
173impl<C> GetRow<FFT64> for GRLWECt<C, FFT64>
174where
175 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
176{
177 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, FFT64>)
178 where
179 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
180 {
181 #[cfg(debug_assertions)]
182 {
183 assert_eq!(col_j, 0);
184 }
185 module.vmp_extract_row(res, self, row_i, col_j);
186 }
187}
188
189impl<C> SetRow<FFT64> for GRLWECt<C, FFT64>
190where
191 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
192{
193 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &RLWECtDft<R, FFT64>)
194 where
195 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
196 {
197 #[cfg(debug_assertions)]
198 {
199 assert_eq!(col_j, 0);
200 }
201 module.vmp_prepare_row(self, row_i, col_j, a);
202 }
203}
204
205impl<C> MatZnxDftProducts<GRLWECt<C, FFT64>, C> for GRLWECt<C, FFT64>
206where
207 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
208{
209 fn mul_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
210 module.bytes_of_vec_znx_dft(2, grlwe_size)
211 + (module.vec_znx_big_normalize_tmp_bytes()
212 | (module.vmp_apply_tmp_bytes(res_size, a_size, a_size, 1, 2, grlwe_size)
213 + module.bytes_of_vec_znx_dft(1, a_size)))
214 }
215
216 fn mul_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
217 where
218 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
219 VecZnx<R>: VecZnxToMut,
220 VecZnx<A>: VecZnxToRef,
221 {
222 let log_base2k: usize = self.log_base2k();
223
224 #[cfg(debug_assertions)]
225 {
226 assert_eq!(res.log_base2k(), log_base2k);
227 assert_eq!(a.log_base2k(), log_base2k);
228 assert_eq!(self.n(), module.n());
229 assert_eq!(res.n(), module.n());
230 assert_eq!(a.n(), module.n());
231 }
232
233 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
234
235 {
236 let (mut a1_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 1, a.size());
237 module.vec_znx_dft(&mut a1_dft, 0, a, 1);
238 module.vmp_apply(&mut res_dft, &a1_dft, self, scratch2);
239 }
240
241 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
242
243 module.vec_znx_big_add_small_inplace(&mut res_big, 0, a, 0);
244
245 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
246 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
247 }
248}
249
250impl ProdByScratchSpace for GRLWECt<Vec<u8>, FFT64> {
251 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
252 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_mat_rlwe_inplace_scratch_space(
253 module, lhs, rhs,
254 )
255 }
256
257 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
258 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_mat_rlwe_inplace_scratch_space(
259 module, lhs, rhs,
260 )
261 }
262}
263
264impl FromProdByScratchSpace for GRLWECt<Vec<u8>, FFT64> {
265 fn from_prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
266 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_mat_rlwe_scratch_space(
267 module, res_size, lhs, rhs,
268 )
269 }
270
271 fn from_prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
272 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_mat_rlwe_scratch_space(
273 module, res_size, lhs, rhs,
274 )
275 }
276}
277
278impl<MUT> ProdBy<GRLWECt<MUT, FFT64>> for GRLWECt<MUT, FFT64>
279where
280 GRLWECt<MUT, FFT64>: GetRow<FFT64> + SetRow<FFT64> + Infos,
281{
282 fn prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
283 where
284 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
285 {
286 rhs.mul_mat_rlwe_inplace(module, self, scratch);
287 }
288
289 fn prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
290 where
291 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
292 {
293 rhs.mul_mat_rlwe_inplace(module, self, scratch);
294 }
295}
296
297impl<MUT, REF> FromProdBy<GRLWECt<MUT, FFT64>, GRLWECt<REF, FFT64>> for GRLWECt<MUT, FFT64>
298where
299 GRLWECt<MUT, FFT64>: GetRow<FFT64> + SetRow<FFT64> + Infos,
300 GRLWECt<REF, FFT64>: GetRow<FFT64> + Infos,
301{
302 fn from_prod_by_grlwe<R>(
303 &mut self,
304 module: &Module<FFT64>,
305 lhs: &GRLWECt<REF, FFT64>,
306 rhs: &GRLWECt<R, FFT64>,
307 scratch: &mut Scratch,
308 ) where
309 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
310 {
311 rhs.mul_mat_rlwe(module, self, lhs, scratch);
312 }
313
314 fn from_prod_by_rgsw<R>(
315 &mut self,
316 module: &Module<FFT64>,
317 lhs: &GRLWECt<REF, FFT64>,
318 rhs: &RGSWCt<R, FFT64>,
319 scratch: &mut Scratch,
320 ) where
321 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
322 {
323 rhs.mul_mat_rlwe(module, self, lhs, scratch);
324 }
325}