Skip to main content

poulpy_core/layouts/compressed/
lwe.rs

1use 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/// Seed-compressed LWE ciphertext layout.
16///
17/// Stores only the body (constant term) of an [`LWE`] ciphertext; the
18/// mask coefficients are regenerated deterministically from a 32-byte
19/// PRNG seed during decompression.
20#[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    /// Allocates a new compressed LWE by copying parameters from an existing info provider.
72    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    /// Allocates a new compressed LWE with the given parameters.
80    ///
81    /// The ring degree is fixed to 1 (scalar LWE). The number of limbs
82    /// is `ceil(k / base2k)`.
83    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
175// module-only API: decompression is provided by `LWEDecompress` on `Module`.
176
177pub 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}