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::{
11 GetRow, Infos, MatRLWEProduct, MatRLWEProductScratchSpace, ProdInplace, ProdInplaceScratchSpace, ProdScratchSpace,
12 Product, SetRow,
13 },
14 grlwe::GRLWECt,
15 keys::SecretKeyDft,
16 rlwe::{RLWECt, RLWECtDft, RLWEPt, encrypt_rlwe_sk},
17 utils::derive_size,
18};
19
20pub struct RGSWCt<C, B: Backend> {
21 pub data: MatZnxDft<C, B>,
22 pub log_base2k: usize,
23 pub log_k: usize,
24}
25
26impl<B: Backend> RGSWCt<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, 2, 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 RGSWCt<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 RGSWCt<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 RGSWCt<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 RGSWCt<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_rgsw_sk<C, P, S>(
80 module: &Module<FFT64>,
81 ct: &mut RGSWCt<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 size: usize = ct.size();
95 let log_base2k: usize = ct.log_base2k();
96
97 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
98 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, 2, size);
99
100 let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
101 data: tmp_znx_pt,
102 log_base2k: log_base2k,
103 log_k: ct.log_k(),
104 };
105
106 let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
107 data: tmp_znx_ct,
108 log_base2k: log_base2k,
109 log_k: ct.log_k(),
110 };
111
112 (0..ct.rows()).for_each(|row_j| {
113 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
114 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
115 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scrach_2);
116
117 (0..ct.cols()).for_each(|col_i| {
118 // rlwe encrypt of vec_znx_pt into vec_znx_ct
119 encrypt_rlwe_sk(
120 module,
121 &mut vec_znx_ct,
122 Some((&vec_znx_pt, col_i)),
123 sk_dft,
124 source_xa,
125 source_xe,
126 sigma,
127 bound,
128 scrach_2,
129 );
130
131 // Switch vec_znx_ct into DFT domain
132 {
133 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, 2, size);
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 module.vmp_prepare_row(ct, row_j, col_i, &vec_znx_dft_ct);
137 }
138 });
139
140 vec_znx_pt.data.zero(); // zeroes for next iteration
141 });
142}
143
144impl<C> RGSWCt<C, FFT64> {
145 pub fn encrypt_sk<P, S>(
146 &mut self,
147 module: &Module<FFT64>,
148 pt: &ScalarZnx<P>,
149 sk_dft: &SecretKeyDft<S, FFT64>,
150 source_xa: &mut Source,
151 source_xe: &mut Source,
152 sigma: f64,
153 bound: f64,
154 scratch: &mut Scratch,
155 ) where
156 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
157 ScalarZnx<P>: ScalarZnxToRef,
158 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
159 {
160 encrypt_rgsw_sk(
161 module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
162 )
163 }
164}
165
166impl<C> GetRow<FFT64> for RGSWCt<C, FFT64>
167where
168 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
169{
170 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, FFT64>)
171 where
172 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
173 {
174 module.vmp_extract_row(res, self, row_i, col_j);
175 }
176}
177
178impl<C> SetRow<FFT64> for RGSWCt<C, FFT64>
179where
180 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
181{
182 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &RLWECtDft<R, FFT64>)
183 where
184 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
185 {
186 module.vmp_prepare_row(self, row_i, col_j, a);
187 }
188}
189
190impl MatRLWEProductScratchSpace for RGSWCt<Vec<u8>, FFT64> {
191 fn prod_with_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
198impl<C> MatRLWEProduct for RGSWCt<C, FFT64>
199where
200 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
201{
202 fn prod_with_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
203 where
204 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
205 VecZnx<R>: VecZnxToMut,
206 VecZnx<A>: VecZnxToRef,
207 {
208 let log_base2k: usize = self.log_base2k();
209
210 #[cfg(debug_assertions)]
211 {
212 assert_eq!(res.log_base2k(), log_base2k);
213 assert_eq!(a.log_base2k(), log_base2k);
214 assert_eq!(self.n(), module.n());
215 assert_eq!(res.n(), module.n());
216 assert_eq!(a.n(), module.n());
217 }
218
219 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
220
221 {
222 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, a.size());
223 module.vec_znx_dft(&mut a_dft, 0, a, 0);
224 module.vec_znx_dft(&mut a_dft, 1, a, 1);
225 module.vmp_apply(&mut res_dft, &a_dft, self, scratch2);
226 }
227
228 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
229
230 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
231 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
232 }
233}
234
235impl ProdInplaceScratchSpace for RGSWCt<Vec<u8>, FFT64> {
236 fn prod_by_grlwe_inplace_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
237 <GRLWECt<Vec<u8>, FFT64> as MatRLWEProductScratchSpace>::prod_with_mat_rlwe_inplace_scratch_space(module, lhs, rhs)
238 }
239
240 fn prod_by_rgsw_inplace_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
241 <RGSWCt<Vec<u8>, FFT64> as MatRLWEProductScratchSpace>::prod_with_mat_rlwe_inplace_scratch_space(module, lhs, rhs)
242 }
243}
244
245impl ProdScratchSpace for RGSWCt<Vec<u8>, FFT64> {
246 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
247 <GRLWECt<Vec<u8>, FFT64> as MatRLWEProductScratchSpace>::prod_with_mat_rlwe_scratch_space(module, res_size, lhs, rhs)
248 }
249
250 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
251 <RGSWCt<Vec<u8>, FFT64> as MatRLWEProductScratchSpace>::prod_with_mat_rlwe_scratch_space(module, res_size, lhs, rhs)
252 }
253}
254
255impl<MUT, REF> ProdInplace<MUT, REF> for RGSWCt<MUT, FFT64>
256where
257 RGSWCt<MUT, FFT64>: GetRow<FFT64> + SetRow<FFT64> + Infos,
258 MatZnxDft<REF, FFT64>: MatZnxDftToRef<FFT64>,
259{
260 fn prod_by_grlwe_inplace(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<REF, FFT64>, scratch: &mut Scratch) {
261 rhs.prod_with_mat_rlwe_inplace(module, self, scratch);
262 }
263
264 fn prod_by_rgsw_inplace(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<REF, FFT64>, scratch: &mut Scratch) {
265 rhs.prod_with_mat_rlwe_inplace(module, self, scratch);
266 }
267}
268
269impl<MUT, REF> Product<MUT, REF> for RGSWCt<MUT, FFT64>
270where
271 MatZnxDft<MUT, FFT64>: MatZnxDftToRef<FFT64> + MatZnxDftToMut<FFT64>,
272 MatZnxDft<REF, FFT64>: MatZnxDftToRef<FFT64>,
273{
274 type Lhs = RGSWCt<REF, FFT64>;
275
276 fn prod_by_grlwe(&mut self, module: &Module<FFT64>, lhs: &Self::Lhs, rhs: &GRLWECt<REF, FFT64>, scratch: &mut Scratch) {
277 rhs.prod_with_mat_rlwe(module, self, lhs, scratch);
278 }
279
280 fn prod_by_rgsw(&mut self, module: &Module<FFT64>, lhs: &Self::Lhs, rhs: &RGSWCt<REF, FFT64>, scratch: &mut Scratch) {
281 rhs.prod_with_mat_rlwe(module, self, lhs, scratch);
282 }
283}