poulpy_core/layouts/compressed/
lwe.rs1use std::fmt;
2
3use poulpy_hal::{
4 api::VecZnxFillUniformSourceBackend,
5 layouts::{
6 Backend, Data, DataView, DataViewMut, FillUniform, HostBytesBackend, HostDataMut, HostDataRef, Module, ReaderFrom,
7 VecZnx, VecZnxToBackendMut, VecZnxToBackendRef, WriterTo, ZnxView, ZnxViewMut, vec_znx_backend_mut_from_mut,
8 vec_znx_backend_ref_from_mut, vec_znx_backend_ref_from_ref,
9 },
10 source::Source,
11};
12
13use crate::layouts::{Base2K, Degree, LWE, LWEInfos, LWEToBackendMut, TorusPrecision};
14
15#[derive(PartialEq, Eq, Clone)]
21pub struct LWECompressed<D: Data> {
22 pub(crate) data: VecZnx<D>,
23 pub(crate) k: TorusPrecision,
24 pub(crate) base2k: Base2K,
25 pub(crate) seed: [u8; 32],
26}
27
28pub type LWECompressedBackendRef<'a, BE> = LWECompressed<<BE as Backend>::BufRef<'a>>;
29pub type LWECompressedBackendMut<'a, BE> = LWECompressed<<BE as Backend>::BufMut<'a>>;
30
31impl<D: Data> LWEInfos for LWECompressed<D> {
32 fn base2k(&self) -> Base2K {
33 self.base2k
34 }
35
36 fn n(&self) -> Degree {
37 Degree(self.data.n() as u32)
38 }
39
40 fn size(&self) -> usize {
41 self.data.size()
42 }
43}
44
45impl<D: HostDataRef> fmt::Debug for LWECompressed<D> {
46 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
47 write!(f, "{self}")
48 }
49}
50
51impl<D: HostDataRef> fmt::Display for LWECompressed<D> {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 write!(
54 f,
55 "LWECompressed: base2k={} k={} seed={:?}: {}",
56 self.base2k(),
57 self.max_k(),
58 self.seed,
59 self.data
60 )
61 }
62}
63
64impl<D: HostDataMut> FillUniform for LWECompressed<D> {
65 fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
66 self.data.fill_uniform(log_bound, source);
67 }
68}
69
70impl LWECompressed<Vec<u8>> {
71 pub(crate) fn alloc_from_infos<A>(infos: &A) -> Self
73 where
74 A: LWEInfos,
75 {
76 Self::alloc(infos.base2k(), infos.max_k())
77 }
78
79 pub(crate) fn alloc(base2k: Base2K, k: TorusPrecision) -> Self {
84 let size: usize = k.0.div_ceil(base2k.0) as usize;
85 LWECompressed {
86 data: VecZnx::from_data(
87 poulpy_hal::layouts::HostBytesBackend::alloc_bytes(VecZnx::<Vec<u8>>::bytes_of(1, 1, size)),
88 1,
89 1,
90 size,
91 ),
92 k,
93 base2k,
94 seed: [0u8; 32],
95 }
96 }
97
98 pub fn bytes_of_from_infos<A>(infos: &A) -> usize
99 where
100 A: LWEInfos,
101 {
102 Self::bytes_of(infos.base2k(), infos.max_k())
103 }
104
105 pub fn bytes_of(base2k: Base2K, k: TorusPrecision) -> usize {
106 VecZnx::bytes_of(1, 1, k.0.div_ceil(base2k.0) as usize)
107 }
108}
109
110use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
111
112impl<D: HostDataMut> ReaderFrom for LWECompressed<D> {
113 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
114 self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
115 self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
116 reader.read_exact(&mut self.seed)?;
117 self.data.read_from(reader)
118 }
119}
120
121impl<D: HostDataRef> WriterTo for LWECompressed<D> {
122 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
123 writer.write_u32::<LittleEndian>(self.k.into())?;
124 writer.write_u32::<LittleEndian>(self.base2k.into())?;
125 writer.write_all(&self.seed)?;
126 self.data.write_to(writer)
127 }
128}
129
130pub trait LWEDecompress
131where
132 Self: VecZnxFillUniformSourceBackend<Self::Backend>,
133{
134 type Backend: Backend;
135
136 fn decompress_lwe<R, O>(&self, res: &mut R, other: &O)
137 where
138 R: LWEToBackendMut<HostBytesBackend>,
139 O: LWECompressedToBackendRef<HostBytesBackend>,
140 {
141 let mut res_ref = res.to_backend_mut();
142 let res: &mut LWE<&mut [u8]> = &mut res_ref;
143 let other = other.to_backend_ref();
144
145 assert_eq!(res.lwe_layout(), other.lwe_layout());
146
147 let mut source: Source = Source::new(other.seed);
148 let mut res_backend = VecZnx::from_data(
149 <Self::Backend as Backend>::from_host_bytes(res.mask.data()),
150 res.mask.n(),
151 res.mask.cols(),
152 res.mask.size(),
153 );
154 {
155 let mut res_backend_mut =
156 <VecZnx<<Self::Backend as Backend>::OwnedBuf> as VecZnxToBackendMut<Self::Backend>>::to_backend_mut(
157 &mut res_backend,
158 );
159 self.vec_znx_fill_uniform_source_backend(other.base2k().into(), &mut res_backend_mut, 0, &mut source);
160 }
161 <Self::Backend as Backend>::copy_to_host(res_backend.data(), res.mask.data_mut());
162 for i in 0..res.size() {
163 res.body.at_mut(0, i)[0] = other.data.at(0, i)[0];
164 }
165 }
166}
167
168impl<B: Backend> LWEDecompress for Module<B>
169where
170 Self: VecZnxFillUniformSourceBackend<B>,
171{
172 type Backend = B;
173}
174
175pub trait LWECompressedToBackendRef<BE: Backend> {
178 fn to_backend_ref(&self) -> LWECompressedBackendRef<'_, BE>;
179}
180
181impl<BE: Backend> LWECompressedToBackendRef<BE> for LWECompressed<BE::OwnedBuf> {
182 fn to_backend_ref(&self) -> LWECompressedBackendRef<'_, BE> {
183 LWECompressed {
184 k: self.k,
185 base2k: self.base2k,
186 seed: self.seed,
187 data: <VecZnx<BE::OwnedBuf> as VecZnxToBackendRef<BE>>::to_backend_ref(&self.data),
188 }
189 }
190}
191
192impl<'b, BE: Backend + 'b> LWECompressedToBackendRef<BE> for &LWECompressed<BE::BufRef<'b>> {
193 fn to_backend_ref(&self) -> LWECompressedBackendRef<'_, BE> {
194 LWECompressed {
195 k: self.k,
196 base2k: self.base2k,
197 seed: self.seed,
198 data: vec_znx_backend_ref_from_ref::<BE>(&self.data),
199 }
200 }
201}
202
203impl<'b, BE: Backend + 'b> LWECompressedToBackendRef<BE> for &mut LWECompressed<BE::BufMut<'b>> {
204 fn to_backend_ref(&self) -> LWECompressedBackendRef<'_, BE> {
205 LWECompressed {
206 k: self.k,
207 base2k: self.base2k,
208 seed: self.seed,
209 data: vec_znx_backend_ref_from_mut::<BE>(&self.data),
210 }
211 }
212}
213
214pub trait LWECompressedToBackendMut<BE: Backend>: LWECompressedToBackendRef<BE> {
215 fn to_backend_mut(&mut self) -> LWECompressedBackendMut<'_, BE>;
216}
217
218impl<BE: Backend> LWECompressedToBackendMut<BE> for LWECompressed<BE::OwnedBuf> {
219 fn to_backend_mut(&mut self) -> LWECompressedBackendMut<'_, BE> {
220 LWECompressed {
221 k: self.k,
222 base2k: self.base2k,
223 seed: self.seed,
224 data: <VecZnx<BE::OwnedBuf> as VecZnxToBackendMut<BE>>::to_backend_mut(&mut self.data),
225 }
226 }
227}
228
229impl<'b, BE: Backend + 'b> LWECompressedToBackendMut<BE> for &mut LWECompressed<BE::BufMut<'b>> {
230 fn to_backend_mut(&mut self) -> LWECompressedBackendMut<'_, BE> {
231 LWECompressed {
232 k: self.k,
233 base2k: self.base2k,
234 seed: self.seed,
235 data: vec_znx_backend_mut_from_mut::<BE>(&mut self.data),
236 }
237 }
238}