poulpy/rlwe/src/elem_rgsw.rs
rlwe/src/elem_rgsw.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
3 ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxOps,
4 ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::Infos,
10 elem_rlwe::{RLWECt, RLWECtDft, RLWEPt, encrypt_rlwe_sk},
11 keys::SecretKeyDft,
12 utils::derive_size,
13};
14
15pub struct RGSWCt<C, B: Backend> {
16 pub data: MatZnxDft<C, B>,
17 pub log_base2k: usize,
18 pub log_k: usize,
19}
20
21impl<B: Backend> RGSWCt<Vec<u8>, B> {
22 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
23 Self {
24 data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
25 log_base2k: log_base2k,
26 log_k: log_k,
27 }
28 }
29}
30
31impl<T, B: Backend> Infos for RGSWCt<T, B> {
32 type Inner = MatZnxDft<T, B>;
33
34 fn inner(&self) -> &Self::Inner {
35 &self.data
36 }
37
38 fn log_base2k(&self) -> usize {
39 self.log_base2k
40 }
41
42 fn log_k(&self) -> usize {
43 self.log_k
44 }
45}
46
47impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
48where
49 MatZnxDft<C, B>: MatZnxDftToMut<B>,
50{
51 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
52 self.data.to_mut()
53 }
54}
55
56impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
57where
58 MatZnxDft<C, B>: MatZnxDftToRef<B>,
59{
60 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
61 self.data.to_ref()
62 }
63}
64
65impl RGSWCt<Vec<u8>, FFT64> {
66 pub fn encrypt_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
67 RLWECt::encrypt_sk_scratch_bytes(module, size)
68 + module.bytes_of_vec_znx(2, size)
69 + module.bytes_of_vec_znx(1, size)
70 + module.bytes_of_vec_znx_dft(2, size)
71 }
72}
73
74impl<C> RGSWCt<C, FFT64>
75where
76 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
77{
78 pub fn get_row(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut RLWECtDft<C, FFT64>)
79 where
80 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64>,
81 {
82 module.vmp_extract_row(res, self, row_i, col_j);
83 }
84}
85
86pub fn encrypt_rgsw_sk<C, P, S>(
87 module: &Module<FFT64>,
88 ct: &mut RGSWCt<C, FFT64>,
89 pt: &ScalarZnx<P>,
90 sk_dft: &SecretKeyDft<S, FFT64>,
91 source_xa: &mut Source,
92 source_xe: &mut Source,
93 sigma: f64,
94 bound: f64,
95 scratch: &mut Scratch,
96) where
97 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
98 ScalarZnx<P>: ScalarZnxToRef,
99 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
100{
101 let size: usize = ct.size();
102 let log_base2k: usize = ct.log_base2k();
103
104 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
105 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, 2, size);
106
107 let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
108 data: tmp_znx_pt,
109 log_base2k: log_base2k,
110 log_k: ct.log_k(),
111 };
112
113 let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
114 data: tmp_znx_ct,
115 log_base2k: log_base2k,
116 log_k: ct.log_k(),
117 };
118
119 (0..ct.rows()).for_each(|row_j| {
120 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
121 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
122 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scrach_2);
123
124 (0..ct.cols()).for_each(|col_i| {
125 // rlwe encrypt of vec_znx_pt into vec_znx_ct
126 encrypt_rlwe_sk(
127 module,
128 &mut vec_znx_ct,
129 Some((&vec_znx_pt, col_i)),
130 sk_dft,
131 source_xa,
132 source_xe,
133 sigma,
134 bound,
135 scrach_2,
136 );
137
138 // Switch vec_znx_ct into DFT domain
139 {
140 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, 2, size);
141 module.vec_znx_dft(&mut vec_znx_dft_ct, 0, &vec_znx_ct, 0);
142 module.vec_znx_dft(&mut vec_znx_dft_ct, 1, &vec_znx_ct, 1);
143 module.vmp_prepare_row(ct, row_j, col_i, &vec_znx_dft_ct);
144 }
145 });
146
147 vec_znx_pt.data.zero(); // zeroes for next iteration
148 });
149}
150
151impl<C> RGSWCt<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_rgsw_sk(
168 module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
169 )
170 }
171}
172
173#[cfg(test)]
174mod tests {
175 use base2k::{
176 FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps,
177 VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, ZnxZero,
178 };
179 use sampling::source::Source;
180
181 use crate::{
182 elem::Infos,
183 elem_rlwe::{RLWECtDft, RLWEPt},
184 keys::{SecretKey, SecretKeyDft},
185 };
186
187 use super::RGSWCt;
188
189 #[test]
190 fn encrypt_rgsw_sk_fft64() {
191 let module: Module<FFT64> = Module::<FFT64>::new(2048);
192 let log_base2k: usize = 8;
193 let log_k_ct: usize = 54;
194 let rows: usize = 4;
195
196 let sigma: f64 = 3.2;
197 let bound: f64 = sigma * 6.0;
198
199 let mut ct: RGSWCt<Vec<u8>, FFT64> = RGSWCt::new(&module, log_base2k, log_k_ct, rows);
200 let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
201 let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
202 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
203
204 let mut source_xs: Source = Source::new([0u8; 32]);
205 let mut source_xe: Source = Source::new([0u8; 32]);
206 let mut source_xa: Source = Source::new([0u8; 32]);
207
208 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
209
210 let mut scratch: ScratchOwned = ScratchOwned::new(
211 RGSWCt::encrypt_sk_scratch_bytes(&module, ct.size()) | RLWECtDft::decrypt_scratch_bytes(&module, ct.size()),
212 );
213
214 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
215 sk.fill_ternary_prob(0.5, &mut source_xs);
216
217 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
218 sk_dft.dft(&module, &sk);
219
220 ct.encrypt_sk(
221 &module,
222 &pt_scalar,
223 &sk_dft,
224 &mut source_xa,
225 &mut source_xe,
226 sigma,
227 bound,
228 scratch.borrow(),
229 );
230
231 let mut ct_rlwe_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct, 2);
232 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
233 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
234
235 (0..ct.cols()).for_each(|col_j| {
236 (0..ct.rows()).for_each(|row_i| {
237 module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
238
239 if col_j == 1 {
240 module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
241 module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, 0);
242 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
243 module.vec_znx_big_normalize(log_base2k, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
244 }
245
246 ct.get_row(&module, row_i, col_j, &mut ct_rlwe_dft);
247
248 ct_rlwe_dft.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
249
250 module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
251
252 let std_pt: f64 = pt_have.data.std(0, log_base2k) * (log_k_ct as f64).exp2();
253 assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
254
255 pt_want.data.zero();
256 });
257 });
258
259 module.free();
260 }
261}