poulpy_core/layouts/compressed/
glwe.rs

1use poulpy_hal::{
2    api::{VecZnxCopy, VecZnxFillUniform},
3    layouts::{
4        Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, VecZnx, VecZnxToMut, VecZnxToRef, WriterTo, ZnxInfos,
5    },
6    source::Source,
7};
8
9use crate::layouts::{Base2K, Degree, GLWE, GLWEInfos, GLWEToMut, GetDegree, LWEInfos, Rank, SetGLWEInfos, TorusPrecision};
10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
11use std::fmt;
12
13#[derive(PartialEq, Eq, Clone)]
14pub struct GLWECompressed<D: Data> {
15    pub(crate) data: VecZnx<D>,
16    pub(crate) base2k: Base2K,
17    pub(crate) k: TorusPrecision,
18    pub(crate) rank: Rank,
19    pub(crate) seed: [u8; 32],
20}
21
22pub trait GLWECompressedSeedMut {
23    fn seed_mut(&mut self) -> &mut [u8; 32];
24}
25
26impl<D: DataMut> GLWECompressedSeedMut for GLWECompressed<D> {
27    fn seed_mut(&mut self) -> &mut [u8; 32] {
28        &mut self.seed
29    }
30}
31
32pub trait GLWECompressedSeed {
33    fn seed(&self) -> &[u8; 32];
34}
35
36impl<D: DataRef> GLWECompressedSeed for GLWECompressed<D> {
37    fn seed(&self) -> &[u8; 32] {
38        &self.seed
39    }
40}
41
42impl<D: Data> LWEInfos for GLWECompressed<D> {
43    fn base2k(&self) -> Base2K {
44        self.base2k
45    }
46
47    fn k(&self) -> TorusPrecision {
48        self.k
49    }
50
51    fn size(&self) -> usize {
52        self.data.size()
53    }
54
55    fn n(&self) -> Degree {
56        Degree(self.data.n() as u32)
57    }
58}
59impl<D: Data> GLWEInfos for GLWECompressed<D> {
60    fn rank(&self) -> Rank {
61        self.rank
62    }
63}
64
65impl<D: DataRef> fmt::Debug for GLWECompressed<D> {
66    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
67        write!(f, "{self}")
68    }
69}
70
71impl<D: DataRef> fmt::Display for GLWECompressed<D> {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        write!(
74            f,
75            "GLWECompressed: base2k={} k={} rank={} seed={:?}: {}",
76            self.base2k(),
77            self.k(),
78            self.rank(),
79            self.seed,
80            self.data
81        )
82    }
83}
84
85impl<D: DataMut> FillUniform for GLWECompressed<D> {
86    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
87        self.data.fill_uniform(log_bound, source);
88    }
89}
90
91impl GLWECompressed<Vec<u8>> {
92    pub fn alloc_from_infos<A>(infos: &A) -> Self
93    where
94        A: GLWEInfos,
95    {
96        Self::alloc(infos.n(), infos.base2k(), infos.k(), infos.rank())
97    }
98
99    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
100        GLWECompressed {
101            data: VecZnx::alloc(n.into(), 1, k.0.div_ceil(base2k.0) as usize),
102            base2k,
103            k,
104            rank,
105            seed: [0u8; 32],
106        }
107    }
108
109    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
110    where
111        A: GLWEInfos,
112    {
113        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
114    }
115
116    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
117        VecZnx::bytes_of(n.into(), 1, k.0.div_ceil(base2k.0) as usize)
118    }
119}
120
121impl<D: DataMut> ReaderFrom for GLWECompressed<D> {
122    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
123        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
124        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
125        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
126        reader.read_exact(&mut self.seed)?;
127        self.data.read_from(reader)
128    }
129}
130
131impl<D: DataRef> WriterTo for GLWECompressed<D> {
132    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
133        writer.write_u32::<LittleEndian>(self.k.into())?;
134        writer.write_u32::<LittleEndian>(self.base2k.into())?;
135        writer.write_u32::<LittleEndian>(self.rank.into())?;
136        writer.write_all(&self.seed)?;
137        self.data.write_to(writer)
138    }
139}
140
141pub trait GLWEDecompress
142where
143    Self: GetDegree + VecZnxFillUniform + VecZnxCopy,
144{
145    fn decompress_glwe<R, O>(&self, res: &mut R, other: &O)
146    where
147        R: GLWEToMut + SetGLWEInfos,
148        O: GLWECompressedToRef + GLWEInfos,
149    {
150        {
151            let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
152            let other: &GLWECompressed<&[u8]> = &other.to_ref();
153            assert_eq!(
154                res.n(),
155                self.ring_degree(),
156                "invalid receiver: res.n()={} != other.n()={}",
157                res.n(),
158                self.ring_degree()
159            );
160
161            assert_eq!(res.glwe_layout(), other.glwe_layout());
162
163            let mut source: Source = Source::new(other.seed);
164
165            self.vec_znx_copy(&mut res.data, 0, &other.data, 0);
166            (1..(other.rank() + 1).into()).for_each(|i| {
167                self.vec_znx_fill_uniform(other.base2k.into(), &mut res.data, i, &mut source);
168            });
169        }
170
171        res.set_base2k(other.base2k());
172        res.set_k(other.k());
173    }
174}
175
176impl<B: Backend> GLWEDecompress for Module<B> where Self: GetDegree + VecZnxFillUniform + VecZnxCopy {}
177
178impl<D: DataMut> GLWE<D> {
179    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
180    where
181        O: GLWECompressedToRef + GLWEInfos,
182        M: GLWEDecompress,
183    {
184        module.decompress_glwe(self, other);
185    }
186}
187
188pub trait GLWECompressedToRef {
189    fn to_ref(&self) -> GLWECompressed<&[u8]>;
190}
191
192impl<D: DataRef> GLWECompressedToRef for GLWECompressed<D> {
193    fn to_ref(&self) -> GLWECompressed<&[u8]> {
194        GLWECompressed {
195            seed: self.seed,
196            base2k: self.base2k,
197            k: self.k,
198            rank: self.rank,
199            data: self.data.to_ref(),
200        }
201    }
202}
203
204pub trait GLWECompressedToMut {
205    fn to_mut(&mut self) -> GLWECompressed<&mut [u8]>;
206}
207
208impl<D: DataMut> GLWECompressedToMut for GLWECompressed<D> {
209    fn to_mut(&mut self) -> GLWECompressed<&mut [u8]> {
210        GLWECompressed {
211            seed: self.seed,
212            base2k: self.base2k,
213            k: self.k,
214            rank: self.rank,
215            data: self.data.to_mut(),
216        }
217    }
218}