Skip to main content

poulpy_core/layouts/compressed/
glwe.rs

1use poulpy_hal::{
2    api::{VecZnxCopyBackend, VecZnxFillUniformSourceBackend},
3    layouts::{
4        Backend, Data, FillUniform, HostDataMut, HostDataRef, Module, ReaderFrom, VecZnx, VecZnxToBackendMut, VecZnxToBackendRef,
5        WriterTo,
6    },
7    source::Source,
8};
9
10use crate::layouts::{Base2K, Degree, GLWEInfos, GLWEToBackendMut, GetDegree, LWEInfos, Rank, SetLWEInfos, TorusPrecision};
11use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
12use std::fmt;
13use std::ops::{Deref, DerefMut};
14
15/// Seed-compressed GLWE ciphertext layout.
16///
17/// Stores only the body component of a [`GLWE`] ciphertext; the mask
18/// polynomials are regenerated deterministically from a 32-byte PRNG
19/// seed during decompression. This reduces the serialized size by a
20/// factor proportional to the rank.
21#[derive(PartialEq, Eq, Clone)]
22pub struct GLWECompressed<D: Data> {
23    pub(crate) data: VecZnx<D>,
24    pub(crate) base2k: Base2K,
25    pub(crate) rank: Rank,
26    pub(crate) seed: [u8; 32],
27}
28
29pub type GLWECompressedBackendRef<'a, BE> = GLWECompressed<<BE as Backend>::BufRef<'a>>;
30pub type GLWECompressedBackendMut<'a, BE> = GLWECompressed<<BE as Backend>::BufMut<'a>>;
31
32pub struct GLWECompressedViewRef<'a, BE: Backend + 'a> {
33    inner: GLWECompressedBackendRef<'a, BE>,
34}
35
36pub struct GLWECompressedViewMut<'a, BE: Backend + 'a> {
37    inner: GLWECompressedBackendMut<'a, BE>,
38}
39
40impl<'a, BE: Backend + 'a> GLWECompressedViewRef<'a, BE> {
41    pub fn from_inner(inner: GLWECompressedBackendRef<'a, BE>) -> Self {
42        Self { inner }
43    }
44
45    pub fn into_inner(self) -> GLWECompressedBackendRef<'a, BE> {
46        self.inner
47    }
48}
49
50impl<'a, BE: Backend + 'a> GLWECompressedViewMut<'a, BE> {
51    pub fn from_inner(inner: GLWECompressedBackendMut<'a, BE>) -> Self {
52        Self { inner }
53    }
54
55    pub fn into_inner(self) -> GLWECompressedBackendMut<'a, BE> {
56        self.inner
57    }
58}
59
60impl<'a, BE: Backend + 'a> Deref for GLWECompressedViewRef<'a, BE> {
61    type Target = GLWECompressedBackendRef<'a, BE>;
62
63    fn deref(&self) -> &Self::Target {
64        &self.inner
65    }
66}
67
68impl<'a, BE: Backend + 'a> Deref for GLWECompressedViewMut<'a, BE> {
69    type Target = GLWECompressedBackendMut<'a, BE>;
70
71    fn deref(&self) -> &Self::Target {
72        &self.inner
73    }
74}
75
76impl<'a, BE: Backend + 'a> DerefMut for GLWECompressedViewMut<'a, BE> {
77    fn deref_mut(&mut self) -> &mut Self::Target {
78        &mut self.inner
79    }
80}
81
82impl<'a, BE: Backend + 'a> LWEInfos for GLWECompressedViewRef<'a, BE> {
83    fn base2k(&self) -> Base2K {
84        self.inner.base2k()
85    }
86
87    fn size(&self) -> usize {
88        self.inner.size()
89    }
90
91    fn n(&self) -> Degree {
92        self.inner.n()
93    }
94}
95
96impl<'a, BE: Backend + 'a> LWEInfos for GLWECompressedViewMut<'a, BE> {
97    fn base2k(&self) -> Base2K {
98        self.inner.base2k()
99    }
100
101    fn size(&self) -> usize {
102        self.inner.size()
103    }
104
105    fn n(&self) -> Degree {
106        self.inner.n()
107    }
108}
109
110impl<'a, BE: Backend + 'a> GLWEInfos for GLWECompressedViewRef<'a, BE> {
111    fn rank(&self) -> Rank {
112        self.inner.rank()
113    }
114}
115
116impl<'a, BE: Backend + 'a> GLWEInfos for GLWECompressedViewMut<'a, BE> {
117    fn rank(&self) -> Rank {
118        self.inner.rank()
119    }
120}
121
122/// Provides mutable access to the PRNG seed of a compressed GLWE.
123pub trait GLWECompressedSeedMut {
124    /// Returns a mutable reference to the 32-byte PRNG seed.
125    fn seed_mut(&mut self) -> &mut [u8; 32];
126}
127
128impl<D: Data> GLWECompressedSeedMut for GLWECompressed<D> {
129    fn seed_mut(&mut self) -> &mut [u8; 32] {
130        &mut self.seed
131    }
132}
133
134/// Provides read access to the PRNG seed of a compressed GLWE.
135pub trait GLWECompressedSeed {
136    /// Returns a reference to the 32-byte PRNG seed.
137    fn seed(&self) -> &[u8; 32];
138}
139
140impl<D: HostDataRef> GLWECompressedSeed for GLWECompressed<D> {
141    fn seed(&self) -> &[u8; 32] {
142        &self.seed
143    }
144}
145
146impl<D: Data> LWEInfos for GLWECompressed<D> {
147    fn base2k(&self) -> Base2K {
148        self.base2k
149    }
150
151    fn size(&self) -> usize {
152        self.data.size()
153    }
154
155    fn n(&self) -> Degree {
156        Degree(self.data.n() as u32)
157    }
158}
159impl<D: Data> GLWEInfos for GLWECompressed<D> {
160    fn rank(&self) -> Rank {
161        self.rank
162    }
163}
164
165impl<D: Data> LWEInfos for &GLWECompressed<D> {
166    fn n(&self) -> Degree {
167        (*self).n()
168    }
169
170    fn base2k(&self) -> Base2K {
171        (*self).base2k()
172    }
173
174    fn size(&self) -> usize {
175        (*self).size()
176    }
177}
178
179impl<D: Data> GLWEInfos for &GLWECompressed<D> {
180    fn rank(&self) -> Rank {
181        (*self).rank()
182    }
183}
184
185impl<D: HostDataRef> fmt::Debug for GLWECompressed<D> {
186    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
187        write!(f, "{self}")
188    }
189}
190
191impl<D: HostDataRef> fmt::Display for GLWECompressed<D> {
192    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193        write!(
194            f,
195            "GLWECompressed: base2k={} k={} rank={} seed={:?}: {}",
196            self.base2k(),
197            self.max_k(),
198            self.rank(),
199            self.seed,
200            self.data
201        )
202    }
203}
204
205impl<D: HostDataMut> FillUniform for GLWECompressed<D> {
206    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
207        self.data.fill_uniform(log_bound, source);
208    }
209}
210
211impl GLWECompressed<Vec<u8>> {
212    /// Allocates a new compressed GLWE by copying parameters from an existing info provider.
213    pub(crate) fn alloc_from_infos<A>(infos: &A) -> Self
214    where
215        A: GLWEInfos,
216    {
217        Self::alloc(infos.n(), infos.base2k(), infos.max_k(), infos.rank())
218    }
219
220    /// Allocates a new compressed GLWE with the given parameters.
221    ///
222    /// The underlying `VecZnx` is sized to hold one column of
223    /// `ceil(k / base2k)` limbs at ring degree `n`.
224    pub(crate) fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
225        let size: usize = k.0.div_ceil(base2k.0) as usize;
226        GLWECompressed {
227            data: VecZnx::from_data(
228                poulpy_hal::layouts::HostBytesBackend::alloc_bytes(VecZnx::<Vec<u8>>::bytes_of(n.into(), 1, size)),
229                n.into(),
230                1,
231                size,
232            ),
233            base2k,
234            rank,
235            seed: [0u8; 32],
236        }
237    }
238
239    /// Returns the serialized byte size by copying parameters from an existing info provider.
240    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
241    where
242        A: GLWEInfos,
243    {
244        Self::bytes_of(infos.n(), infos.base2k(), infos.max_k())
245    }
246
247    /// Returns the serialized byte size for a compressed GLWE with the given parameters.
248    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
249        VecZnx::bytes_of(n.into(), 1, k.0.div_ceil(base2k.0) as usize)
250    }
251}
252
253/// Deserializes the metadata (k, base2k, rank, seed) followed by the body data.
254impl<D: HostDataMut> ReaderFrom for GLWECompressed<D> {
255    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
256        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
257        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
258        reader.read_exact(&mut self.seed)?;
259        self.data.read_from(reader)
260    }
261}
262
263/// Serializes the metadata (k, base2k, rank, seed) followed by the body data.
264impl<D: HostDataRef> WriterTo for GLWECompressed<D> {
265    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
266        writer.write_u32::<LittleEndian>(self.base2k.into())?;
267        writer.write_u32::<LittleEndian>(self.rank.into())?;
268        writer.write_all(&self.seed)?;
269        self.data.write_to(writer)
270    }
271}
272
273/// Trait for decompressing a [`GLWECompressed`] into a standard [`GLWE`].
274///
275/// Copies the body from the compressed ciphertext and regenerates
276/// the mask polynomials from the stored PRNG seed.
277pub trait GLWEDecompress
278where
279    Self: GetDegree + VecZnxFillUniformSourceBackend<Self::Backend> + VecZnxCopyBackend<Self::Backend>,
280{
281    type Backend: Backend;
282
283    /// Decompresses `other` into `res` by copying the body and regenerating the mask.
284    fn decompress_glwe<R, O>(&self, res: &mut R, other: &O)
285    where
286        R: GLWEToBackendMut<Self::Backend> + SetLWEInfos,
287        O: GLWECompressedToBackendRef<Self::Backend> + GLWEInfos,
288    {
289        let other = other.to_backend_ref();
290        {
291            let res = &mut res.to_backend_mut();
292            assert_eq!(
293                res.n(),
294                self.ring_degree(),
295                "invalid receiver: res.n()={} != other.n()={}",
296                res.n(),
297                self.ring_degree()
298            );
299
300            assert_eq!(res.glwe_layout(), other.glwe_layout());
301
302            let mut source: Source = Source::new(other.seed);
303
304            self.vec_znx_copy_backend(&mut res.data, 0, &other.data, 0);
305            (1..(other.rank() + 1).into()).for_each(|i| {
306                self.vec_znx_fill_uniform_source_backend(other.base2k.into(), &mut res.data, i, &mut source);
307            });
308        }
309
310        res.set_base2k(other.base2k());
311    }
312}
313
314impl<B: Backend> GLWEDecompress for Module<B>
315where
316    Self: GetDegree + VecZnxFillUniformSourceBackend<B> + VecZnxCopyBackend<B>,
317{
318    type Backend = B;
319}
320
321// module-only API: decompression is provided by `GLWEDecompress` on `Module`.
322
323pub trait GLWECompressedToBackendRef<BE: Backend> {
324    fn to_backend_ref(&self) -> GLWECompressedBackendRef<'_, BE>;
325}
326
327impl<BE: Backend> GLWECompressedToBackendRef<BE> for GLWECompressed<BE::OwnedBuf> {
328    fn to_backend_ref(&self) -> GLWECompressedBackendRef<'_, BE> {
329        GLWECompressed {
330            seed: self.seed,
331            base2k: self.base2k,
332            rank: self.rank,
333            data: <VecZnx<BE::OwnedBuf> as VecZnxToBackendRef<BE>>::to_backend_ref(&self.data),
334        }
335    }
336}
337
338impl<'a, BE: Backend + 'a> GLWECompressedToBackendRef<BE> for GLWECompressedViewRef<'a, BE> {
339    fn to_backend_ref(&self) -> GLWECompressedBackendRef<'_, BE> {
340        GLWECompressed {
341            seed: self.inner.seed,
342            base2k: self.inner.base2k,
343            rank: self.inner.rank,
344            data: poulpy_hal::layouts::vec_znx_backend_ref_from_ref::<BE>(&self.inner.data),
345        }
346    }
347}
348
349impl<'a, BE: Backend + 'a> GLWECompressedToBackendRef<BE> for GLWECompressedViewMut<'a, BE> {
350    fn to_backend_ref(&self) -> GLWECompressedBackendRef<'_, BE> {
351        GLWECompressed {
352            seed: self.inner.seed,
353            base2k: self.inner.base2k,
354            rank: self.inner.rank,
355            data: poulpy_hal::layouts::vec_znx_backend_ref_from_mut::<BE>(&self.inner.data),
356        }
357    }
358}
359
360pub trait GLWECompressedToBackendMut<BE: Backend>: GLWECompressedToBackendRef<BE> {
361    fn to_backend_mut(&mut self) -> GLWECompressedBackendMut<'_, BE>;
362}
363
364impl<BE: Backend> GLWECompressedToBackendMut<BE> for GLWECompressed<BE::OwnedBuf> {
365    fn to_backend_mut(&mut self) -> GLWECompressedBackendMut<'_, BE> {
366        GLWECompressed {
367            seed: self.seed,
368            base2k: self.base2k,
369            rank: self.rank,
370            data: <VecZnx<BE::OwnedBuf> as VecZnxToBackendMut<BE>>::to_backend_mut(&mut self.data),
371        }
372    }
373}