poulpy/core/src/gglwe_ciphertext.rs
core/src/gglwe_ciphertext.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::{GetRow, Infos, SetRow},
11 glwe_ciphertext::GLWECiphertext,
12 glwe_ciphertext_fourier::GLWECiphertextFourier,
13 glwe_plaintext::GLWEPlaintext,
14 keys::SecretKeyFourier,
15 utils::derive_size,
16 vec_glwe_product::{VecGLWEProduct, VecGLWEProductScratchSpace},
17};
18
19pub struct GGLWECiphertext<C, B: Backend> {
20 pub(crate) data: MatZnxDft<C, B>,
21 pub(crate) basek: usize,
22 pub(crate) k: usize,
23}
24
25impl<B: Backend> GGLWECiphertext<Vec<u8>, B> {
26 pub fn new(module: &Module<B>, base2k: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> Self {
27 Self {
28 data: module.new_mat_znx_dft(rows, rank_in, rank_out + 1, derive_size(base2k, k)),
29 basek: base2k,
30 k,
31 }
32 }
33}
34
35impl<T, B: Backend> Infos for GGLWECiphertext<T, B> {
36 type Inner = MatZnxDft<T, B>;
37
38 fn inner(&self) -> &Self::Inner {
39 &self.data
40 }
41
42 fn basek(&self) -> usize {
43 self.basek
44 }
45
46 fn k(&self) -> usize {
47 self.k
48 }
49}
50
51impl<T, B: Backend> GGLWECiphertext<T, B> {
52 pub fn rank(&self) -> usize {
53 self.data.cols_out() - 1
54 }
55}
56
57impl<C, B: Backend> MatZnxDftToMut<B> for GGLWECiphertext<C, B>
58where
59 MatZnxDft<C, B>: MatZnxDftToMut<B>,
60{
61 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
62 self.data.to_mut()
63 }
64}
65
66impl<C, B: Backend> MatZnxDftToRef<B> for GGLWECiphertext<C, B>
67where
68 MatZnxDft<C, B>: MatZnxDftToRef<B>,
69{
70 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
71 self.data.to_ref()
72 }
73}
74
75impl GGLWECiphertext<Vec<u8>, FFT64> {
76 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
77 GLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
78 + module.bytes_of_vec_znx(rank + 1, size)
79 + module.bytes_of_vec_znx(1, size)
80 + module.bytes_of_vec_znx_dft(rank + 1, size)
81 }
82
83 pub fn encrypt_pk_scratch_space(_module: &Module<FFT64>, _rank: usize, _pk_size: usize) -> usize {
84 unimplemented!()
85 }
86}
87
88impl<DataSelf> GGLWECiphertext<DataSelf, FFT64>
89where
90 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + ZnxInfos,
91{
92 pub fn encrypt_sk<DataPt, DataSk>(
93 &mut self,
94 module: &Module<FFT64>,
95 pt: &ScalarZnx<DataPt>,
96 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
97 source_xa: &mut Source,
98 source_xe: &mut Source,
99 sigma: f64,
100 bound: f64,
101 scratch: &mut Scratch,
102 ) where
103 ScalarZnx<DataPt>: ScalarZnxToRef,
104 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
105 {
106 #[cfg(debug_assertions)]
107 {
108 assert_eq!(self.rank(), sk_dft.rank());
109 assert_eq!(self.n(), module.n());
110 assert_eq!(sk_dft.n(), module.n());
111 assert_eq!(pt.n(), module.n());
112 }
113
114 let rows: usize = self.rows();
115 let size: usize = self.size();
116 let basek: usize = self.basek();
117 let k: usize = self.k();
118
119 let cols: usize = self.rank() + 1;
120
121 let (tmp_znx_pt, scrach_1) = scratch.tmp_vec_znx(module, 1, size);
122 let (tmp_znx_ct, scrach_2) = scrach_1.tmp_vec_znx(module, cols, size);
123 let (tmp_znx_dft_ct, scratch_3) = scrach_2.tmp_vec_znx_dft(module, cols, size);
124
125 let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
126 data: tmp_znx_pt,
127 basek,
128 k,
129 };
130
131 let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
132 data: tmp_znx_ct,
133 basek,
134 k,
135 };
136
137 let mut vec_znx_ct_dft: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier {
138 data: tmp_znx_dft_ct,
139 basek,
140 k,
141 };
142
143 (0..rows).for_each(|row_i| {
144 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
145 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_i, pt, 0);
146 module.vec_znx_normalize_inplace(basek, &mut vec_znx_pt, 0, scratch_3);
147
148 // rlwe encrypt of vec_znx_pt into vec_znx_ct
149 vec_znx_ct.encrypt_sk(
150 module,
151 &vec_znx_pt,
152 sk_dft,
153 source_xa,
154 source_xe,
155 sigma,
156 bound,
157 scratch_3,
158 );
159
160 vec_znx_pt.data.zero(); // zeroes for next iteration
161
162 // Switch vec_znx_ct into DFT domain
163 vec_znx_ct.dft(module, &mut vec_znx_ct_dft);
164
165 // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
166 module.vmp_prepare_row(self, row_i, 0, &vec_znx_ct_dft);
167 });
168 }
169}
170
171impl<C> GetRow<FFT64> for GGLWECiphertext<C, FFT64>
172where
173 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
174{
175 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
176 where
177 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
178 {
179 #[cfg(debug_assertions)]
180 {
181 assert_eq!(col_j, 0);
182 }
183 module.vmp_extract_row(res, self, row_i, col_j);
184 }
185}
186
187impl<C> SetRow<FFT64> for GGLWECiphertext<C, FFT64>
188where
189 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
190{
191 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
192 where
193 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
194 {
195 #[cfg(debug_assertions)]
196 {
197 assert_eq!(col_j, 0);
198 }
199 module.vmp_prepare_row(self, row_i, col_j, a);
200 }
201}
202
203impl VecGLWEProductScratchSpace for GGLWECiphertext<Vec<u8>, FFT64> {
204 fn prod_with_glwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
205 module.bytes_of_vec_znx_dft(2, grlwe_size)
206 + (module.vec_znx_big_normalize_tmp_bytes()
207 | (module.vmp_apply_tmp_bytes(res_size, a_size, a_size, 1, 2, grlwe_size)
208 + module.bytes_of_vec_znx_dft(1, a_size)))
209 }
210}
211
212impl<C> VecGLWEProduct for GGLWECiphertext<C, FFT64>
213where
214 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
215{
216 fn prod_with_glwe<R, A>(
217 &self,
218 module: &Module<FFT64>,
219 res: &mut GLWECiphertext<R>,
220 a: &GLWECiphertext<A>,
221 scratch: &mut Scratch,
222 ) where
223 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
224 VecZnx<R>: VecZnxToMut,
225 VecZnx<A>: VecZnxToRef,
226 {
227 let log_base2k: usize = self.basek();
228
229 #[cfg(debug_assertions)]
230 {
231 assert_eq!(res.basek(), log_base2k);
232 assert_eq!(a.basek(), log_base2k);
233 assert_eq!(self.n(), module.n());
234 assert_eq!(res.n(), module.n());
235 assert_eq!(a.n(), module.n());
236 }
237
238 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
239
240 {
241 let (mut a1_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 1, a.size());
242 module.vec_znx_dft(&mut a1_dft, 0, a, 1);
243 module.vmp_apply(&mut res_dft, &a1_dft, self, scratch2);
244 }
245
246 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
247
248 module.vec_znx_big_add_small_inplace(&mut res_big, 0, a, 0);
249
250 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
251 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
252 }
253}