poulpy_core/layouts/compressed/
lwe.rs1use std::fmt;
2
3use poulpy_hal::{
4 api::ZnFillUniform,
5 layouts::{
6 Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, WriterTo, Zn, ZnToMut, ZnToRef, ZnxInfos, ZnxView,
7 ZnxViewMut,
8 },
9 source::Source,
10};
11
12use crate::layouts::{Base2K, Degree, LWE, LWEInfos, LWEToMut, TorusPrecision};
13
14#[derive(PartialEq, Eq, Clone)]
15pub struct LWECompressed<D: Data> {
16 pub(crate) data: Zn<D>,
17 pub(crate) k: TorusPrecision,
18 pub(crate) base2k: Base2K,
19 pub(crate) seed: [u8; 32],
20}
21
22impl<D: Data> LWEInfos for LWECompressed<D> {
23 fn base2k(&self) -> Base2K {
24 self.base2k
25 }
26
27 fn k(&self) -> TorusPrecision {
28 self.k
29 }
30
31 fn n(&self) -> Degree {
32 Degree(self.data.n() as u32)
33 }
34
35 fn size(&self) -> usize {
36 self.data.size()
37 }
38}
39
40impl<D: DataRef> fmt::Debug for LWECompressed<D> {
41 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
42 write!(f, "{self}")
43 }
44}
45
46impl<D: DataRef> fmt::Display for LWECompressed<D> {
47 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48 write!(
49 f,
50 "LWECompressed: base2k={} k={} seed={:?}: {}",
51 self.base2k(),
52 self.k(),
53 self.seed,
54 self.data
55 )
56 }
57}
58
59impl<D: DataMut> FillUniform for LWECompressed<D> {
60 fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
61 self.data.fill_uniform(log_bound, source);
62 }
63}
64
65impl LWECompressed<Vec<u8>> {
66 pub fn alloc_from_infos<A>(infos: &A) -> Self
67 where
68 A: LWEInfos,
69 {
70 Self::alloc(infos.base2k(), infos.k())
71 }
72
73 pub fn alloc(base2k: Base2K, k: TorusPrecision) -> Self {
74 LWECompressed {
75 data: Zn::alloc(1, 1, k.0.div_ceil(base2k.0) as usize),
76 k,
77 base2k,
78 seed: [0u8; 32],
79 }
80 }
81
82 pub fn bytes_of_from_infos<A>(infos: &A) -> usize
83 where
84 A: LWEInfos,
85 {
86 Self::bytes_of(infos.base2k(), infos.k())
87 }
88
89 pub fn bytes_of(base2k: Base2K, k: TorusPrecision) -> usize {
90 Zn::bytes_of(1, 1, k.0.div_ceil(base2k.0) as usize)
91 }
92}
93
94use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
95
96impl<D: DataMut> ReaderFrom for LWECompressed<D> {
97 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
98 self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
99 self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
100 reader.read_exact(&mut self.seed)?;
101 self.data.read_from(reader)
102 }
103}
104
105impl<D: DataRef> WriterTo for LWECompressed<D> {
106 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
107 writer.write_u32::<LittleEndian>(self.k.into())?;
108 writer.write_u32::<LittleEndian>(self.base2k.into())?;
109 writer.write_all(&self.seed)?;
110 self.data.write_to(writer)
111 }
112}
113
114pub trait LWEDecompress
115where
116 Self: ZnFillUniform,
117{
118 fn decompress_lwe<R, O>(&self, res: &mut R, other: &O)
119 where
120 R: LWEToMut,
121 O: LWECompressedToRef,
122 {
123 let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
124 let other: &LWECompressed<&[u8]> = &other.to_ref();
125
126 assert_eq!(res.lwe_layout(), other.lwe_layout());
127
128 let mut source: Source = Source::new(other.seed);
129 self.zn_fill_uniform(
130 res.n().into(),
131 other.base2k().into(),
132 &mut res.data,
133 0,
134 &mut source,
135 );
136 for i in 0..res.size() {
137 res.data.at_mut(0, i)[0] = other.data.at(0, i)[0];
138 }
139 }
140}
141
142impl<B: Backend> LWEDecompress for Module<B> where Self: ZnFillUniform {}
143
144impl<D: DataMut> LWE<D> {
145 pub fn decompress<O, M>(&mut self, module: &M, other: &O)
146 where
147 O: LWECompressedToRef,
148 M: LWEDecompress,
149 {
150 module.decompress_lwe(self, other);
151 }
152}
153
154pub trait LWECompressedToRef {
155 fn to_ref(&self) -> LWECompressed<&[u8]>;
156}
157
158impl<D: DataRef> LWECompressedToRef for LWECompressed<D> {
159 fn to_ref(&self) -> LWECompressed<&[u8]> {
160 LWECompressed {
161 k: self.k,
162 base2k: self.base2k,
163 seed: self.seed,
164 data: self.data.to_ref(),
165 }
166 }
167}
168
169pub trait LWECompressedToMut {
170 fn to_mut(&mut self) -> LWECompressed<&mut [u8]>;
171}
172
173impl<D: DataMut> LWECompressedToMut for LWECompressed<D> {
174 fn to_mut(&mut self) -> LWECompressed<&mut [u8]> {
175 LWECompressed {
176 k: self.k,
177 base2k: self.base2k,
178 seed: self.seed,
179 data: self.data.to_mut(),
180 }
181 }
182}