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