poulpy/core/src/gglwe/layout.rs
core/src/gglwe/layout.rsRaw
1use backend::hal::{
2 api::{MatZnxAlloc, MatZnxAllocBytes, VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
3 layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
4};
5
6use crate::{GLWECiphertext, Infos};
7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
8
9pub trait GGLWEExecLayoutFamily<B: Backend> = VmpPMatAlloc<B> + VmpPMatAllocBytes + VmpPMatPrepare<B>;
10
11#[derive(PartialEq, Eq)]
12pub struct GGLWECiphertext<D: Data> {
13 pub(crate) data: MatZnx<D>,
14 pub(crate) basek: usize,
15 pub(crate) k: usize,
16 pub(crate) digits: usize,
17}
18
19impl<D: DataRef> GGLWECiphertext<D> {
20 pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
21 GLWECiphertext {
22 data: self.data.at(row, col),
23 basek: self.basek,
24 k: self.k,
25 }
26 }
27}
28
29impl<D: DataMut> GGLWECiphertext<D> {
30 pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
31 GLWECiphertext {
32 data: self.data.at_mut(row, col),
33 basek: self.basek,
34 k: self.k,
35 }
36 }
37}
38
39impl GGLWECiphertext<Vec<u8>> {
40 pub fn alloc<B: Backend>(
41 module: &Module<B>,
42 basek: usize,
43 k: usize,
44 rows: usize,
45 digits: usize,
46 rank_in: usize,
47 rank_out: usize,
48 ) -> Self
49 where
50 Module<B>: MatZnxAlloc,
51 {
52 let size: usize = k.div_ceil(basek);
53 debug_assert!(
54 size > digits,
55 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
56 size,
57 digits
58 );
59
60 assert!(
61 rows * digits <= size,
62 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
63 rows,
64 digits,
65 size
66 );
67
68 Self {
69 data: module.mat_znx_alloc(rows, rank_in, rank_out + 1, size),
70 basek: basek,
71 k,
72 digits,
73 }
74 }
75
76 pub fn bytes_of<B: Backend>(
77 module: &Module<B>,
78 basek: usize,
79 k: usize,
80 rows: usize,
81 digits: usize,
82 rank_in: usize,
83 rank_out: usize,
84 ) -> usize
85 where
86 Module<B>: MatZnxAllocBytes,
87 {
88 let size: usize = k.div_ceil(basek);
89 debug_assert!(
90 size > digits,
91 "invalid gglwe: ceil(k/basek): {} <= digits: {}",
92 size,
93 digits
94 );
95
96 assert!(
97 rows * digits <= size,
98 "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
99 rows,
100 digits,
101 size
102 );
103
104 module.mat_znx_alloc_bytes(rows, rank_in, rank_out + 1, rows)
105 }
106}
107
108impl<D: Data> Infos for GGLWECiphertext<D> {
109 type Inner = MatZnx<D>;
110
111 fn inner(&self) -> &Self::Inner {
112 &self.data
113 }
114
115 fn basek(&self) -> usize {
116 self.basek
117 }
118
119 fn k(&self) -> usize {
120 self.k
121 }
122}
123
124impl<D: Data> GGLWECiphertext<D> {
125 pub fn rank(&self) -> usize {
126 self.data.cols_out() - 1
127 }
128
129 pub fn digits(&self) -> usize {
130 self.digits
131 }
132
133 pub fn rank_in(&self) -> usize {
134 self.data.cols_in()
135 }
136
137 pub fn rank_out(&self) -> usize {
138 self.data.cols_out() - 1
139 }
140}
141
142impl<D: DataMut> ReaderFrom for GGLWECiphertext<D> {
143 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
144 self.k = reader.read_u64::<LittleEndian>()? as usize;
145 self.basek = reader.read_u64::<LittleEndian>()? as usize;
146 self.digits = reader.read_u64::<LittleEndian>()? as usize;
147 self.data.read_from(reader)
148 }
149}
150
151impl<D: DataRef> WriterTo for GGLWECiphertext<D> {
152 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
153 writer.write_u64::<LittleEndian>(self.k as u64)?;
154 writer.write_u64::<LittleEndian>(self.basek as u64)?;
155 writer.write_u64::<LittleEndian>(self.digits as u64)?;
156 self.data.write_to(writer)
157 }
158}
159
160#[derive(PartialEq, Eq)]
161pub struct GLWESwitchingKey<D: Data> {
162 pub(crate) key: GGLWECiphertext<D>,
163 pub(crate) sk_in_n: usize, // Degree of sk_in
164 pub(crate) sk_out_n: usize, // Degree of sk_out
165}
166
167impl GLWESwitchingKey<Vec<u8>> {
168 pub fn alloc<B: Backend>(
169 module: &Module<B>,
170 basek: usize,
171 k: usize,
172 rows: usize,
173 digits: usize,
174 rank_in: usize,
175 rank_out: usize,
176 ) -> Self
177 where
178 Module<B>: MatZnxAlloc,
179 {
180 GLWESwitchingKey {
181 key: GGLWECiphertext::alloc(module, basek, k, rows, digits, rank_in, rank_out),
182 sk_in_n: 0,
183 sk_out_n: 0,
184 }
185 }
186
187 pub fn bytes_of<B: Backend>(
188 module: &Module<B>,
189 basek: usize,
190 k: usize,
191 rows: usize,
192 digits: usize,
193 rank_in: usize,
194 rank_out: usize,
195 ) -> usize
196 where
197 Module<B>: MatZnxAllocBytes,
198 {
199 GGLWECiphertext::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, rank_in, rank_out)
200 }
201}
202
203impl<D: Data> Infos for GLWESwitchingKey<D> {
204 type Inner = MatZnx<D>;
205
206 fn inner(&self) -> &Self::Inner {
207 self.key.inner()
208 }
209
210 fn basek(&self) -> usize {
211 self.key.basek()
212 }
213
214 fn k(&self) -> usize {
215 self.key.k()
216 }
217}
218
219impl<D: Data> GLWESwitchingKey<D> {
220 pub fn rank(&self) -> usize {
221 self.key.data.cols_out() - 1
222 }
223
224 pub fn rank_in(&self) -> usize {
225 self.key.data.cols_in()
226 }
227
228 pub fn rank_out(&self) -> usize {
229 self.key.data.cols_out() - 1
230 }
231
232 pub fn digits(&self) -> usize {
233 self.key.digits()
234 }
235
236 pub fn sk_degree_in(&self) -> usize {
237 self.sk_in_n
238 }
239
240 pub fn sk_degree_out(&self) -> usize {
241 self.sk_out_n
242 }
243}
244
245impl<D: DataRef> GLWESwitchingKey<D> {
246 pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
247 self.key.at(row, col)
248 }
249}
250
251impl<D: DataMut> GLWESwitchingKey<D> {
252 pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
253 self.key.at_mut(row, col)
254 }
255}
256
257impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
258 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
259 self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
260 self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
261 self.key.read_from(reader)
262 }
263}
264
265impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
266 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
267 writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
268 writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
269 self.key.write_to(writer)
270 }
271}
272
273#[derive(PartialEq, Eq)]
274pub struct AutomorphismKey<D: Data> {
275 pub(crate) key: GLWESwitchingKey<D>,
276 pub(crate) p: i64,
277}
278
279impl AutomorphismKey<Vec<u8>> {
280 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
281 where
282 Module<B>: MatZnxAlloc,
283 {
284 AutomorphismKey {
285 key: GLWESwitchingKey::alloc(module, basek, k, rows, digits, rank, rank),
286 p: 0,
287 }
288 }
289
290 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
291 where
292 Module<B>: MatZnxAllocBytes,
293 {
294 GLWESwitchingKey::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, rank, rank)
295 }
296}
297
298impl<D: Data> Infos for AutomorphismKey<D> {
299 type Inner = MatZnx<D>;
300
301 fn inner(&self) -> &Self::Inner {
302 &self.key.inner()
303 }
304
305 fn basek(&self) -> usize {
306 self.key.basek()
307 }
308
309 fn k(&self) -> usize {
310 self.key.k()
311 }
312}
313
314impl<D: Data> AutomorphismKey<D> {
315 pub fn p(&self) -> i64 {
316 self.p
317 }
318
319 pub fn digits(&self) -> usize {
320 self.key.digits()
321 }
322
323 pub fn rank(&self) -> usize {
324 self.key.rank()
325 }
326
327 pub fn rank_in(&self) -> usize {
328 self.key.rank_in()
329 }
330
331 pub fn rank_out(&self) -> usize {
332 self.key.rank_out()
333 }
334}
335
336impl<D: DataRef> AutomorphismKey<D> {
337 pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
338 self.key.at(row, col)
339 }
340}
341
342impl<D: DataMut> AutomorphismKey<D> {
343 pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
344 self.key.at_mut(row, col)
345 }
346}
347
348impl<D: DataMut> ReaderFrom for AutomorphismKey<D> {
349 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
350 self.p = reader.read_u64::<LittleEndian>()? as i64;
351 self.key.read_from(reader)
352 }
353}
354
355impl<D: DataRef> WriterTo for AutomorphismKey<D> {
356 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
357 writer.write_u64::<LittleEndian>(self.p as u64)?;
358 self.key.write_to(writer)
359 }
360}
361
362#[derive(PartialEq, Eq)]
363pub struct GLWETensorKey<D: Data> {
364 pub(crate) keys: Vec<GLWESwitchingKey<D>>,
365}
366
367impl GLWETensorKey<Vec<u8>> {
368 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
369 where
370 Module<B>: MatZnxAlloc,
371 {
372 let mut keys: Vec<GLWESwitchingKey<Vec<u8>>> = Vec::new();
373 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
374 (0..pairs).for_each(|_| {
375 keys.push(GLWESwitchingKey::alloc(
376 module, basek, k, rows, digits, 1, rank,
377 ));
378 });
379 Self { keys: keys }
380 }
381
382 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
383 where
384 Module<B>: MatZnxAllocBytes,
385 {
386 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
387 pairs * GLWESwitchingKey::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, 1, rank)
388 }
389}
390
391impl<D: Data> Infos for GLWETensorKey<D> {
392 type Inner = MatZnx<D>;
393
394 fn inner(&self) -> &Self::Inner {
395 &self.keys[0].inner()
396 }
397
398 fn basek(&self) -> usize {
399 self.keys[0].basek()
400 }
401
402 fn k(&self) -> usize {
403 self.keys[0].k()
404 }
405}
406
407impl<D: Data> GLWETensorKey<D> {
408 pub fn rank(&self) -> usize {
409 self.keys[0].rank()
410 }
411
412 pub fn rank_in(&self) -> usize {
413 self.keys[0].rank_in()
414 }
415
416 pub fn rank_out(&self) -> usize {
417 self.keys[0].rank_out()
418 }
419
420 pub fn digits(&self) -> usize {
421 self.keys[0].digits()
422 }
423}
424
425impl<D: DataMut> GLWETensorKey<D> {
426 // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
427 pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<D> {
428 if i > j {
429 std::mem::swap(&mut i, &mut j);
430 };
431 let rank: usize = self.rank();
432 &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
433 }
434}
435
436impl<D: DataRef> GLWETensorKey<D> {
437 // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
438 pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<D> {
439 if i > j {
440 std::mem::swap(&mut i, &mut j);
441 };
442 let rank: usize = self.rank();
443 &self.keys[i * rank + j - (i * (i + 1) / 2)]
444 }
445}
446
447impl<D: DataMut> ReaderFrom for GLWETensorKey<D> {
448 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
449 let len: usize = reader.read_u64::<LittleEndian>()? as usize;
450 if self.keys.len() != len {
451 return Err(std::io::Error::new(
452 std::io::ErrorKind::InvalidData,
453 format!("self.keys.len()={} != read len={}", self.keys.len(), len),
454 ));
455 }
456 for key in &mut self.keys {
457 key.read_from(reader)?;
458 }
459 Ok(())
460 }
461}
462
463impl<D: DataRef> WriterTo for GLWETensorKey<D> {
464 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
465 writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
466 for key in &self.keys {
467 key.write_to(writer)?;
468 }
469 Ok(())
470 }
471}