1use poulpy_hal::{
2 layouts::{Data, DataMut, DataRef, FillUniform, MatZnx, MatZnxToMut, MatZnxToRef, ReaderFrom, WriterTo, ZnxInfos},
3 source::Source,
4};
5
6use crate::layouts::{Base2K, Degree, Dnum, Dsize, GLWE, GLWEInfos, LWEInfos, Rank, TorusPrecision};
7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
8
9use std::fmt;
10
11pub trait GGLWEInfos
12where
13 Self: GLWEInfos,
14{
15 fn dnum(&self) -> Dnum;
16 fn dsize(&self) -> Dsize;
17 fn rank_in(&self) -> Rank;
18 fn rank_out(&self) -> Rank;
19 fn gglwe_layout(&self) -> GGLWELayout {
20 GGLWELayout {
21 n: self.n(),
22 base2k: self.base2k(),
23 k: self.k(),
24 rank_in: self.rank_in(),
25 rank_out: self.rank_out(),
26 dsize: self.dsize(),
27 dnum: self.dnum(),
28 }
29 }
30}
31
32pub trait SetGGLWEInfos {
33 fn set_dsize(&mut self, dsize: usize);
34}
35
36#[derive(PartialEq, Eq, Copy, Clone, Debug)]
37pub struct GGLWELayout {
38 pub n: Degree,
39 pub base2k: Base2K,
40 pub k: TorusPrecision,
41 pub rank_in: Rank,
42 pub rank_out: Rank,
43 pub dnum: Dnum,
44 pub dsize: Dsize,
45}
46
47impl LWEInfos for GGLWELayout {
48 fn base2k(&self) -> Base2K {
49 self.base2k
50 }
51
52 fn k(&self) -> TorusPrecision {
53 self.k
54 }
55
56 fn n(&self) -> Degree {
57 self.n
58 }
59}
60
61impl GLWEInfos for GGLWELayout {
62 fn rank(&self) -> Rank {
63 self.rank_out
64 }
65}
66
67impl GGLWEInfos for GGLWELayout {
68 fn rank_in(&self) -> Rank {
69 self.rank_in
70 }
71
72 fn dsize(&self) -> Dsize {
73 self.dsize
74 }
75
76 fn rank_out(&self) -> Rank {
77 self.rank_out
78 }
79
80 fn dnum(&self) -> Dnum {
81 self.dnum
82 }
83}
84
85#[derive(PartialEq, Eq, Clone)]
86pub struct GGLWE<D: Data> {
87 pub(crate) data: MatZnx<D>,
88 pub(crate) k: TorusPrecision,
89 pub(crate) base2k: Base2K,
90 pub(crate) dsize: Dsize,
91}
92
93impl<D: Data> LWEInfos for GGLWE<D> {
94 fn base2k(&self) -> Base2K {
95 self.base2k
96 }
97
98 fn k(&self) -> TorusPrecision {
99 self.k
100 }
101
102 fn n(&self) -> Degree {
103 Degree(self.data.n() as u32)
104 }
105
106 fn size(&self) -> usize {
107 self.data.size()
108 }
109}
110
111impl<D: Data> GLWEInfos for GGLWE<D> {
112 fn rank(&self) -> Rank {
113 self.rank_out()
114 }
115}
116
117impl<D: Data> GGLWEInfos for GGLWE<D> {
118 fn rank_in(&self) -> Rank {
119 Rank(self.data.cols_in() as u32)
120 }
121
122 fn rank_out(&self) -> Rank {
123 Rank(self.data.cols_out() as u32 - 1)
124 }
125
126 fn dsize(&self) -> Dsize {
127 self.dsize
128 }
129
130 fn dnum(&self) -> Dnum {
131 Dnum(self.data.rows() as u32)
132 }
133}
134
135impl<D: DataRef> GGLWE<D> {
136 pub fn data(&self) -> &MatZnx<D> {
137 &self.data
138 }
139}
140
141impl<D: DataMut> GGLWE<D> {
142 pub fn data_mut(&mut self) -> &mut MatZnx<D> {
143 &mut self.data
144 }
145}
146
147impl<D: DataRef> fmt::Debug for GGLWE<D> {
148 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
149 write!(f, "{self}")
150 }
151}
152
153impl<D: DataMut> FillUniform for GGLWE<D> {
154 fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
155 self.data.fill_uniform(log_bound, source);
156 }
157}
158
159impl<D: DataRef> fmt::Display for GGLWE<D> {
160 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161 write!(
162 f,
163 "(GGLWE: k={} base2k={} dsize={}) {}",
164 self.k().0,
165 self.base2k().0,
166 self.dsize().0,
167 self.data
168 )
169 }
170}
171
172impl<D: DataRef> GGLWE<D> {
173 pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
174 GLWE {
175 k: self.k,
176 base2k: self.base2k,
177 data: self.data.at(row, col),
178 }
179 }
180}
181
182impl<D: DataMut> GGLWE<D> {
183 pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
184 GLWE {
185 k: self.k,
186 base2k: self.base2k,
187 data: self.data.at_mut(row, col),
188 }
189 }
190}
191
192impl GGLWE<Vec<u8>> {
193 pub fn alloc_from_infos<A>(infos: &A) -> Self
194 where
195 A: GGLWEInfos,
196 {
197 Self::alloc(
198 infos.n(),
199 infos.base2k(),
200 infos.k(),
201 infos.rank_in(),
202 infos.rank_out(),
203 infos.dnum(),
204 infos.dsize(),
205 )
206 }
207
208 pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank_in: Rank, rank_out: Rank, dnum: Dnum, dsize: Dsize) -> Self {
209 let size: usize = k.0.div_ceil(base2k.0) as usize;
210 debug_assert!(
211 size as u32 > dsize.0,
212 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
213 dsize.0
214 );
215
216 assert!(
217 dnum.0 * dsize.0 <= size as u32,
218 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
219 dnum.0,
220 dsize.0,
221 );
222
223 GGLWE {
224 data: MatZnx::alloc(
225 n.into(),
226 dnum.into(),
227 rank_in.into(),
228 (rank_out + 1).into(),
229 k.0.div_ceil(base2k.0) as usize,
230 ),
231 k,
232 base2k,
233 dsize,
234 }
235 }
236
237 pub fn bytes_of_from_infos<A>(infos: &A) -> usize
238 where
239 A: GGLWEInfos,
240 {
241 Self::bytes_of(
242 infos.n(),
243 infos.base2k(),
244 infos.k(),
245 infos.rank_in(),
246 infos.rank_out(),
247 infos.dnum(),
248 infos.dsize(),
249 )
250 }
251
252 pub fn bytes_of(
253 n: Degree,
254 base2k: Base2K,
255 k: TorusPrecision,
256 rank_in: Rank,
257 rank_out: Rank,
258 dnum: Dnum,
259 dsize: Dsize,
260 ) -> usize {
261 let size: usize = k.0.div_ceil(base2k.0) as usize;
262 debug_assert!(
263 size as u32 > dsize.0,
264 "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
265 dsize.0
266 );
267
268 assert!(
269 dnum.0 * dsize.0 <= size as u32,
270 "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
271 dnum.0,
272 dsize.0,
273 );
274
275 MatZnx::bytes_of(
276 n.into(),
277 dnum.into(),
278 rank_in.into(),
279 (rank_out + 1).into(),
280 k.0.div_ceil(base2k.0) as usize,
281 )
282 }
283}
284
285pub trait GGLWEToMut {
286 fn to_mut(&mut self) -> GGLWE<&mut [u8]>;
287}
288
289impl<D: DataMut> GGLWEToMut for GGLWE<D> {
290 fn to_mut(&mut self) -> GGLWE<&mut [u8]> {
291 GGLWE {
292 k: self.k(),
293 base2k: self.base2k(),
294 dsize: self.dsize(),
295 data: self.data.to_mut(),
296 }
297 }
298}
299
300pub trait GGLWEToRef {
301 fn to_ref(&self) -> GGLWE<&[u8]>;
302}
303
304impl<D: DataRef> GGLWEToRef for GGLWE<D> {
305 fn to_ref(&self) -> GGLWE<&[u8]> {
306 GGLWE {
307 k: self.k(),
308 base2k: self.base2k(),
309 dsize: self.dsize(),
310 data: self.data.to_ref(),
311 }
312 }
313}
314
315impl<D: DataMut> ReaderFrom for GGLWE<D> {
316 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
317 self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
318 self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
319 self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
320 self.data.read_from(reader)
321 }
322}
323
324impl<D: DataRef> WriterTo for GGLWE<D> {
325 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
326 writer.write_u32::<LittleEndian>(self.k.0)?;
327 writer.write_u32::<LittleEndian>(self.base2k.0)?;
328 writer.write_u32::<LittleEndian>(self.dsize.0)?;
329 self.data.write_to(writer)
330 }
331}