poulpy_core/layouts/
lwe.rs

1use poulpy_hal::layouts::ZnxWord;
2use std::fmt;
3
4use poulpy_hal::{
5    layouts::{
6        Backend, Data, FillUniform, HostDataMut, HostDataRef, ReaderFrom, VecZnx, VecZnxToBackendMut, VecZnxToBackendRef,
7        WriterTo,
8    },
9    source::Source,
10};
11
12use crate::layouts::{Base2K, Degree, TorusPrecision};
13use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
14
15/// Trait providing the parameter accessors for an LWE ciphertext.
16///
17/// An LWE ciphertext is a scalar (non-polynomial) ciphertext consisting of
18/// a body `b` and a mask `(a_1, ..., a_n)`.
19pub trait LWEInfos {
20    /// Returns the LWE dimension, i.e. the number of mask elements (= GLWE ring degree N).
21    fn n(&self) -> Degree;
22    /// Returns `log2(n)`.
23    fn log_n(&self) -> usize {
24        self.n().log2()
25    }
26    /// Returns the Torus precision the **allocation** can hold
27    /// ([`Self::max_size`] `* base2k`).
28    fn max_k(&self) -> TorusPrecision {
29        TorusPrecision(self.max_size() as u32 * self.base2k().as_u32())
30    }
31
32    /// Returns the Torus precision used by the object.
33    fn k(&self) -> TorusPrecision;
34
35    /// Returns the limb width generic core operations should process.
36    fn size(&self) -> usize {
37        self.k().div_ceil(self.base2k()) as usize
38    }
39
40    /// Returns the base-2-log of the limb width used for the RNS/CRT representation.
41    fn base2k(&self) -> Base2K;
42
43    /// Returns the allocated limb **capacity** of the backing container: the
44    /// physical width operations may compute at (`size() <= max_size()`).
45    /// Together with `k()` this fully describes an object: `k` is the claimed
46    /// precision and `max_size` the allocation. Layout/spec types without a
47    /// buffer report their metadata-derived width.
48    fn max_size(&self) -> usize;
49
50    /// Returns a plain-data [`LWELayout`] snapshot of the current parameters.
51    fn lwe_layout(&self) -> LWELayout {
52        LWELayout {
53            n: self.n(),
54            k: self.k(),
55            base2k: self.base2k(),
56        }
57    }
58}
59
60impl<T: LWEInfos + ?Sized> LWEInfos for &T {
61    fn n(&self) -> Degree {
62        (**self).n()
63    }
64
65    fn k(&self) -> TorusPrecision {
66        (**self).k()
67    }
68
69    fn base2k(&self) -> Base2K {
70        (**self).base2k()
71    }
72
73    fn max_size(&self) -> usize {
74        (**self).max_size()
75    }
76}
77
78impl<T: LWEInfos + ?Sized> LWEInfos for &mut T {
79    fn n(&self) -> Degree {
80        (**self).n()
81    }
82
83    fn k(&self) -> TorusPrecision {
84        (**self).k()
85    }
86
87    fn base2k(&self) -> Base2K {
88        (**self).base2k()
89    }
90
91    fn max_size(&self) -> usize {
92        (**self).max_size()
93    }
94}
95
96/// Trait for mutating LWE parameters in place.
97pub trait SetBase2k {
98    /// Sets the limb width `base2k`.
99    fn set_base2k(&mut self, base2k: Base2K);
100}
101
102/// Trait for mutating the Torus precision `k` in place.
103///
104/// `k` is a metadata label (the effective torus width); it is independent of
105/// the underlying buffer's allocation width `max_size`.
106pub trait SetK {
107    /// Sets the Torus precision `k`.
108    fn set_k(&mut self, k: TorusPrecision);
109}
110
111/// Plain-data snapshot of the parameters that describe an [`LWE`] ciphertext.
112#[derive(PartialEq, Eq, Copy, Clone, Debug)]
113pub struct LWELayout {
114    /// Ring degree (LWE dimension).
115    pub n: Degree,
116    /// Torus precision.
117    pub k: TorusPrecision,
118    /// Base-2-log of the limb width.
119    pub base2k: Base2K,
120}
121
122impl LWEInfos for LWELayout {
123    fn base2k(&self) -> Base2K {
124        self.base2k
125    }
126
127    fn n(&self) -> Degree {
128        self.n
129    }
130
131    fn max_size(&self) -> usize {
132        self.k.div_ceil(self.base2k) as usize
133    }
134
135    fn k(&self) -> TorusPrecision {
136        self.k
137    }
138}
139
140/// A scalar (non-polynomial) LWE ciphertext.
141///
142/// Stored as two separate [`VecZnx`] buffers:
143/// - `body`: degree-0 polynomial (n = 1) holding the scalar body `b`.
144/// - `mask`: degree-n polynomial (n = lwe_dim) holding the mask `(a_1, ..., a_n)`.
145///
146/// `D: Data` is the storage backend (e.g. `Vec<u8>`, `&[u8]`, `&mut [u8]`).
147#[derive(PartialEq, Eq, Clone)]
148pub struct LWE<D: Data, W: ZnxWord> {
149    pub(crate) body: VecZnx<D, W>,
150    pub(crate) mask: VecZnx<D, W>,
151    pub(crate) k: TorusPrecision,
152    pub(crate) base2k: Base2K,
153}
154
155pub type LWEBackendRef<'a, BE> = LWE<<BE as Backend>::BufRef<'a>, <BE as Backend>::ZnxWord>;
156pub type LWEBackendMut<'a, BE> = LWE<<BE as Backend>::BufMut<'a>, <BE as Backend>::ZnxWord>;
157
158impl<D: Data, W: ZnxWord> LWEInfos for LWE<D, W> {
159    fn base2k(&self) -> Base2K {
160        self.base2k
161    }
162
163    fn n(&self) -> Degree {
164        Degree(self.mask.n() as u32)
165    }
166
167    fn max_size(&self) -> usize {
168        self.mask.size().min(self.body.size())
169    }
170
171    fn k(&self) -> TorusPrecision {
172        self.k
173    }
174}
175
176impl<D: Data, W: ZnxWord> SetBase2k for LWE<D, W> {
177    fn set_base2k(&mut self, base2k: Base2K) {
178        self.base2k = base2k
179    }
180}
181
182impl<D: Data, W: ZnxWord> LWE<D, W> {
183    /// Returns a shared reference to the body [`VecZnx`] (n = 1).
184    pub fn body(&self) -> &VecZnx<D, W> {
185        &self.body
186    }
187
188    /// Returns a mutable reference to the body [`VecZnx`] (n = 1).
189    pub fn body_mut(&mut self) -> &mut VecZnx<D, W> {
190        &mut self.body
191    }
192
193    /// Returns a shared reference to the mask [`VecZnx`] (n = lwe_dim).
194    pub fn mask(&self) -> &VecZnx<D, W> {
195        &self.mask
196    }
197
198    /// Returns a mutable reference to the mask [`VecZnx`] (n = lwe_dim).
199    pub fn mask_mut(&mut self) -> &mut VecZnx<D, W> {
200        &mut self.mask
201    }
202
203    fn validate_shape(&self) -> std::io::Result<()> {
204        if self.base2k.as_u32() == 0 {
205            return Err(std::io::Error::new(
206                std::io::ErrorKind::InvalidData,
207                "LWE base2k must be non-zero",
208            ));
209        }
210        if self.body.n() != 1 {
211            return Err(std::io::Error::new(
212                std::io::ErrorKind::InvalidData,
213                format!("LWE body degree must be 1, got {}", self.body.n()),
214            ));
215        }
216        if self.body.cols() != 1 {
217            return Err(std::io::Error::new(
218                std::io::ErrorKind::InvalidData,
219                format!("LWE body cols must be 1, got {}", self.body.cols()),
220            ));
221        }
222        if self.mask.cols() != 1 {
223            return Err(std::io::Error::new(
224                std::io::ErrorKind::InvalidData,
225                format!("LWE mask cols must be 1, got {}", self.mask.cols()),
226            ));
227        }
228        if self.body.size() != self.mask.size() {
229            return Err(std::io::Error::new(
230                std::io::ErrorKind::InvalidData,
231                format!(
232                    "LWE body and mask sizes must match, got body.size={} mask.size={}",
233                    self.body.size(),
234                    self.mask.size()
235                ),
236            ));
237        }
238        Ok(())
239    }
240}
241
242impl<D: Data, W: ZnxWord> LWE<D, W> {
243    /// Zero-cost rename when both backends share the same `OwnedBuf`.
244    pub fn reinterpret<To>(self) -> LWE<To::OwnedBuf, To::ZnxWord>
245    where
246        To: Backend<OwnedBuf = D, ZnxWord = W>,
247    {
248        let body_shape = self.body.shape();
249        let body_data = self.body.into_data();
250        let mask_shape = self.mask.shape();
251        let mask_data = self.mask.into_data();
252        LWE {
253            body: VecZnx::from_data(body_data, body_shape.n(), body_shape.cols(), body_shape.size()),
254            mask: VecZnx::from_data(mask_data, mask_shape.n(), mask_shape.cols(), mask_shape.size()),
255            base2k: self.base2k,
256            k: self.k,
257        }
258    }
259}
260
261impl<D: HostDataRef, W: ZnxWord> fmt::Debug for LWE<D, W> {
262    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
263        write!(f, "{self}")
264    }
265}
266
267impl<D: HostDataRef, W: ZnxWord> fmt::Display for LWE<D, W> {
268    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
269        write!(
270            f,
271            "LWE: base2k={} k={}: body={} mask={}",
272            self.base2k().0,
273            self.k().0,
274            self.body,
275            self.mask
276        )
277    }
278}
279
280impl<D: HostDataMut, W: ZnxWord> FillUniform for LWE<D, W>
281where
282    VecZnx<D, W>: FillUniform,
283{
284    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
285        self.mask.fill_uniform(log_bound, source);
286    }
287}
288
289impl<W: ZnxWord> LWE<Vec<u8>, W> {
290    /// Allocates a new [`LWE`] with the given parameters.
291    pub(crate) fn alloc_from_infos<A>(infos: &A) -> Self
292    where
293        A: LWEInfos,
294    {
295        Self::alloc(infos.n(), infos.base2k(), infos.k())
296    }
297
298    /// Allocates a new [`LWE`] with the given parameters.
299    ///
300    /// * `n` -- LWE dimension (mask length).
301    /// * `base2k` -- base-2-log of the limb width.
302    /// * `k` -- torus precision.
303    pub(crate) fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
304        let size: usize = k.0.div_ceil(base2k.0) as usize;
305        LWE {
306            body: VecZnx::from_data(
307                poulpy_hal::layouts::HostBytesBackend::alloc_bytes(VecZnx::<Vec<u8>, W>::bytes_of(1, 1, size)),
308                1,
309                1,
310                size,
311            ),
312            mask: VecZnx::from_data(
313                poulpy_hal::layouts::HostBytesBackend::alloc_bytes(VecZnx::<Vec<u8>, W>::bytes_of(n.as_usize(), 1, size)),
314                n.as_usize(),
315                1,
316                size,
317            ),
318            base2k,
319            k,
320        }
321    }
322
323    /// Returns the byte count required for an [`LWE`] with the given parameters.
324    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
325    where
326        A: LWEInfos,
327    {
328        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
329    }
330
331    /// Returns the byte count required for an [`LWE`] with the given parameters.
332    ///
333    /// * `n` -- LWE dimension (mask length).
334    /// * `base2k` -- base-2-log of the limb width.
335    /// * `k` -- torus precision.
336    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
337        let size: usize = k.0.div_ceil(base2k.0) as usize;
338        VecZnx::<Vec<u8>, W>::bytes_of(1, 1, size) + VecZnx::<Vec<u8>, W>::bytes_of(n.as_usize(), 1, size)
339    }
340}
341
342pub trait LWEToBackendRef<BE: Backend> {
343    fn to_backend_ref(&self) -> LWEBackendRef<'_, BE>;
344}
345
346impl<BE: Backend, D: Data> LWEToBackendRef<BE> for LWE<D, BE::ZnxWord>
347where
348    VecZnx<D, BE::ZnxWord>: VecZnxToBackendRef<BE>,
349{
350    fn to_backend_ref(&self) -> LWEBackendRef<'_, BE> {
351        LWE {
352            base2k: self.base2k,
353            k: self.k,
354            body: self.body.to_backend_ref(),
355            mask: self.mask.to_backend_ref(),
356        }
357    }
358}
359
360pub trait LWEToBackendMut<BE: Backend>: LWEToBackendRef<BE> {
361    fn to_backend_mut(&mut self) -> LWEBackendMut<'_, BE>;
362}
363
364impl<BE: Backend, D: Data> LWEToBackendMut<BE> for LWE<D, BE::ZnxWord>
365where
366    VecZnx<D, BE::ZnxWord>: VecZnxToBackendRef<BE> + VecZnxToBackendMut<BE>,
367{
368    fn to_backend_mut(&mut self) -> LWEBackendMut<'_, BE> {
369        LWE {
370            base2k: self.base2k,
371            k: self.k,
372            body: self.body.to_backend_mut(),
373            mask: self.mask.to_backend_mut(),
374        }
375    }
376}
377
378impl<'b, BE: Backend + 'b> LWEToBackendRef<BE> for &mut LWE<BE::BufMut<'b>, BE::ZnxWord> {
379    fn to_backend_ref(&self) -> LWEBackendRef<'_, BE> {
380        LWE {
381            base2k: self.base2k,
382            k: self.k,
383            body: poulpy_hal::layouts::vec_znx_backend_ref_from_mut::<BE>(&self.body),
384            mask: poulpy_hal::layouts::vec_znx_backend_ref_from_mut::<BE>(&self.mask),
385        }
386    }
387}
388
389impl<'b, BE: Backend + 'b> LWEToBackendMut<BE> for &mut LWE<BE::BufMut<'b>, BE::ZnxWord> {
390    fn to_backend_mut(&mut self) -> LWEBackendMut<'_, BE> {
391        LWE {
392            base2k: self.base2k,
393            k: self.k,
394            body: poulpy_hal::layouts::vec_znx_backend_mut_from_mut::<BE>(&mut self.body),
395            mask: poulpy_hal::layouts::vec_znx_backend_mut_from_mut::<BE>(&mut self.mask),
396        }
397    }
398}
399
400impl<D: HostDataMut, W: ZnxWord> ReaderFrom for LWE<D, W> {
401    /// Deserialises an [`LWE`] in little-endian binary format.
402    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
403        self.set_base2k(Base2K(reader.read_u32::<LittleEndian>()?));
404        self.body.read_from(reader)?;
405        self.mask.read_from(reader)?;
406        self.validate_shape()
407    }
408}
409
410impl<D: HostDataRef, W: ZnxWord> WriterTo for LWE<D, W> {
411    /// Serialises the [`LWE`] in little-endian binary format.
412    fn write_to<Wr: std::io::Write>(&self, writer: &mut Wr) -> std::io::Result<()> {
413        writer.write_u32::<LittleEndian>(self.base2k.into())?;
414        self.body.write_to(writer)?;
415        self.mask.write_to(writer)
416    }
417}
poulpy_core/layouts/lwe.rs

poulpy_core/layouts/
lwe.rs
