poulpy/core/src/gglwe.rs
core/src/gglwe.rsRaw
1use backend::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, Module, ScalarZnx, Scratch, VecZnxAlloc, VecZnxDftAlloc, VecZnxOps,
3 ZnxInfos, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{FourierGLWECiphertext, GLWECiphertext, GLWESecret, GetRow, Infos, ScratchCore, SetRow, div_ceil};
8
9pub struct GGLWECiphertext<C, B: Backend> {
10 pub(crate) data: MatZnxDft<C, B>,
11 pub(crate) basek: usize,
12 pub(crate) k: usize,
13 pub(crate) digits: usize,
14}
15
16impl<B: Backend> GGLWECiphertext<Vec<u8>, B> {
17 pub fn alloc(
18 module: &Module<B>,
19 basek: usize,
20 k: usize,
21 rows: usize,
22 digits: usize,
23 rank_in: usize,
24 rank_out: usize,
25 ) -> Self {
26 let size: usize = k.div_ceil(basek);
27 debug_assert!(
28 size > digits,
29 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
30 size,
31 digits
32 );
33
34 debug_assert!(
35 rows * digits <= size,
36 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
37 rows,
38 digits,
39 size
40 );
41
42 Self {
43 data: module.new_mat_znx_dft(rows, rank_in, rank_out + 1, size),
44 basek: basek,
45 k,
46 digits,
47 }
48 }
49
50 pub fn bytes_of(
51 module: &Module<FFT64>,
52 basek: usize,
53 k: usize,
54 rows: usize,
55 digits: usize,
56 rank_in: usize,
57 rank_out: usize,
58 ) -> usize {
59 let size: usize = k.div_ceil(basek);
60 debug_assert!(
61 size > digits,
62 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
63 size,
64 digits
65 );
66
67 assert!(
68 rows * digits <= size,
69 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
70 rows,
71 digits,
72 size
73 );
74
75 module.bytes_of_mat_znx_dft(rows, rank_in, rank_out + 1, size)
76 }
77}
78
79impl<T, B: Backend> Infos for GGLWECiphertext<T, B> {
80 type Inner = MatZnxDft<T, B>;
81
82 fn inner(&self) -> &Self::Inner {
83 &self.data
84 }
85
86 fn basek(&self) -> usize {
87 self.basek
88 }
89
90 fn k(&self) -> usize {
91 self.k
92 }
93}
94
95impl<T, B: Backend> GGLWECiphertext<T, B> {
96 pub fn rank(&self) -> usize {
97 self.data.cols_out() - 1
98 }
99
100 pub fn digits(&self) -> usize {
101 self.digits
102 }
103
104 pub fn rank_in(&self) -> usize {
105 self.data.cols_in()
106 }
107
108 pub fn rank_out(&self) -> usize {
109 self.data.cols_out() - 1
110 }
111}
112
113impl GGLWECiphertext<Vec<u8>, FFT64> {
114 pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
115 let size = k.div_ceil(basek);
116 GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
117 + module.bytes_of_vec_znx(rank + 1, size)
118 + module.bytes_of_vec_znx(1, size)
119 + module.bytes_of_vec_znx_dft(rank + 1, size)
120 }
121
122 pub fn generate_from_pk_scratch_space(_module: &Module<FFT64>, _basek: usize, _k: usize, _rank: usize) -> usize {
123 unimplemented!()
124 }
125}
126
127impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GGLWECiphertext<DataSelf, FFT64> {
128 pub fn encrypt_sk<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
129 &mut self,
130 module: &Module<FFT64>,
131 pt: &ScalarZnx<DataPt>,
132 sk: &GLWESecret<DataSk, FFT64>,
133 source_xa: &mut Source,
134 source_xe: &mut Source,
135 sigma: f64,
136 scratch: &mut Scratch,
137 ) {
138 #[cfg(debug_assertions)]
139 {
140 assert_eq!(self.rank_in(), pt.cols());
141 assert_eq!(self.rank_out(), sk.rank());
142 assert_eq!(self.n(), module.n());
143 assert_eq!(sk.n(), module.n());
144 assert_eq!(pt.n(), module.n());
145 assert!(
146 scratch.available()
147 >= GGLWECiphertext::generate_from_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
148 "scratch.available: {} < GGLWECiphertext::generate_from_sk_scratch_space(module, self.rank()={}, \
149 self.size()={}): {}",
150 scratch.available(),
151 self.rank(),
152 self.size(),
153 GGLWECiphertext::generate_from_sk_scratch_space(module, self.basek(), self.k(), self.rank())
154 );
155 assert!(
156 self.rows() * self.digits() * self.basek() <= self.k(),
157 "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
158 self.rows(),
159 self.digits(),
160 self.basek(),
161 self.rows() * self.digits() * self.basek(),
162 self.k()
163 );
164 }
165
166 let rows: usize = self.rows();
167 let digits: usize = self.digits();
168 let basek: usize = self.basek();
169 let k: usize = self.k();
170 let rank_in: usize = self.rank_in();
171 let rank_out: usize = self.rank_out();
172
173 let (mut tmp_pt, scrach_1) = scratch.tmp_glwe_pt(module, basek, k);
174 let (mut tmp_ct, scrach_2) = scrach_1.tmp_glwe_ct(module, basek, k, rank_out);
175 let (mut tmp_ct_dft, scratch_3) = scrach_2.tmp_glwe_fourier(module, basek, k, rank_out);
176
177 // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
178 //
179 // Example for ksk rank 2 to rank 3:
180 //
181 // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
182 // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
183 //
184 // Example ksk rank 2 to rank 1
185 //
186 // (-(a*s) + s0, a)
187 // (-(b*s) + s1, b)
188 (0..rank_in).for_each(|col_i| {
189 (0..rows).for_each(|row_i| {
190 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
191 tmp_pt.data.zero(); // zeroes for next iteration
192 module.vec_znx_add_scalar_inplace(
193 &mut tmp_pt.data,
194 0,
195 (digits - 1) + row_i * digits,
196 pt,
197 col_i,
198 );
199 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch_3);
200
201 // rlwe encrypt of vec_znx_pt into vec_znx_ct
202 tmp_ct.encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scratch_3);
203
204 // Switch vec_znx_ct into DFT domain
205 tmp_ct.dft(module, &mut tmp_ct_dft);
206
207 // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
208 self.set_row(module, row_i, col_i, &tmp_ct_dft);
209 });
210 });
211 }
212}
213
214impl<C: AsRef<[u8]>> GetRow<FFT64> for GGLWECiphertext<C, FFT64> {
215 fn get_row<R: AsMut<[u8]> + AsRef<[u8]>>(
216 &self,
217 module: &Module<FFT64>,
218 row_i: usize,
219 col_j: usize,
220 res: &mut FourierGLWECiphertext<R, FFT64>,
221 ) {
222 module.mat_znx_dft_get_row(&mut res.data, &self.data, row_i, col_j);
223 }
224}
225
226impl<C: AsMut<[u8]> + AsRef<[u8]>> SetRow<FFT64> for GGLWECiphertext<C, FFT64> {
227 fn set_row<R: AsRef<[u8]>>(
228 &mut self,
229 module: &Module<FFT64>,
230 row_i: usize,
231 col_j: usize,
232 a: &FourierGLWECiphertext<R, FFT64>,
233 ) {
234 module.mat_znx_dft_set_row(&mut self.data, row_i, col_j, &a.data);
235 }
236}