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