poulpy_core/layouts/
lwe.rs

1use std::fmt;
2
3use poulpy_hal::{
4    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo, Zn, ZnToMut, ZnToRef, ZnxInfos},
5    source::Source,
6};
7
8use crate::layouts::{Base2K, Degree, TorusPrecision};
9use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
10
11pub trait LWEInfos {
12    fn n(&self) -> Degree;
13    fn log_n(&self) -> usize {
14        (u64::BITS - (self.n().as_usize() as u64 - 1).leading_zeros()) as usize
15    }
16    fn k(&self) -> TorusPrecision;
17    fn max_k(&self) -> TorusPrecision {
18        TorusPrecision(self.k().0 * self.size() as u32)
19    }
20    fn base2k(&self) -> Base2K;
21    fn size(&self) -> usize {
22        self.k().0.div_ceil(self.base2k().0) as usize
23    }
24    fn lwe_layout(&self) -> LWELayout {
25        LWELayout {
26            n: self.n(),
27            k: self.k(),
28            base2k: self.base2k(),
29        }
30    }
31}
32
33pub trait SetLWEInfos {
34    fn set_k(&mut self, k: TorusPrecision);
35    fn set_base2k(&mut self, base2k: Base2K);
36}
37
38#[derive(PartialEq, Eq, Copy, Clone, Debug)]
39pub struct LWELayout {
40    pub n: Degree,
41    pub k: TorusPrecision,
42    pub base2k: Base2K,
43}
44
45impl LWEInfos for LWELayout {
46    fn base2k(&self) -> Base2K {
47        self.base2k
48    }
49
50    fn k(&self) -> TorusPrecision {
51        self.k
52    }
53
54    fn n(&self) -> Degree {
55        self.n
56    }
57}
58#[derive(PartialEq, Eq, Clone)]
59pub struct LWE<D: Data> {
60    pub(crate) data: Zn<D>,
61    pub(crate) k: TorusPrecision,
62    pub(crate) base2k: Base2K,
63}
64
65impl<D: Data> LWEInfos for LWE<D> {
66    fn base2k(&self) -> Base2K {
67        self.base2k
68    }
69
70    fn k(&self) -> TorusPrecision {
71        self.k
72    }
73    fn n(&self) -> Degree {
74        Degree(self.data.n() as u32 - 1)
75    }
76
77    fn size(&self) -> usize {
78        self.data.size()
79    }
80}
81
82impl<D: Data> SetLWEInfos for LWE<D> {
83    fn set_base2k(&mut self, base2k: Base2K) {
84        self.base2k = base2k
85    }
86
87    fn set_k(&mut self, k: TorusPrecision) {
88        self.k = k
89    }
90}
91
92impl<D: DataRef> LWE<D> {
93    pub fn data(&self) -> &Zn<D> {
94        &self.data
95    }
96}
97
98impl<D: DataMut> LWE<D> {
99    pub fn data_mut(&mut self) -> &Zn<D> {
100        &mut self.data
101    }
102}
103
104impl<D: DataRef> fmt::Debug for LWE<D> {
105    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
106        write!(f, "{self}")
107    }
108}
109
110impl<D: DataRef> fmt::Display for LWE<D> {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        write!(
113            f,
114            "LWE: base2k={} k={}: {}",
115            self.base2k().0,
116            self.k().0,
117            self.data
118        )
119    }
120}
121
122impl<D: DataMut> FillUniform for LWE<D>
123where
124    Zn<D>: FillUniform,
125{
126    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
127        self.data.fill_uniform(log_bound, source);
128    }
129}
130
131impl LWE<Vec<u8>> {
132    pub fn alloc_from_infos<A>(infos: &A) -> Self
133    where
134        A: LWEInfos,
135    {
136        Self::alloc(infos.n(), infos.base2k(), infos.k())
137    }
138
139    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
140        LWE {
141            data: Zn::alloc((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize),
142            k,
143            base2k,
144        }
145    }
146
147    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
148    where
149        A: LWEInfos,
150    {
151        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
152    }
153
154    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
155        Zn::bytes_of((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize)
156    }
157}
158
159pub trait LWEToRef {
160    fn to_ref(&self) -> LWE<&[u8]>;
161}
162
163impl<D: DataRef> LWEToRef for LWE<D> {
164    fn to_ref(&self) -> LWE<&[u8]> {
165        LWE {
166            k: self.k,
167            base2k: self.base2k,
168            data: self.data.to_ref(),
169        }
170    }
171}
172
173pub trait LWEToMut {
174    #[allow(dead_code)]
175    fn to_mut(&mut self) -> LWE<&mut [u8]>;
176}
177
178impl<D: DataMut> LWEToMut for LWE<D> {
179    fn to_mut(&mut self) -> LWE<&mut [u8]> {
180        LWE {
181            k: self.k,
182            base2k: self.base2k,
183            data: self.data.to_mut(),
184        }
185    }
186}
187
188impl<D: DataMut> ReaderFrom for LWE<D> {
189    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
190        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
191        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
192        self.data.read_from(reader)
193    }
194}
195
196impl<D: DataRef> WriterTo for LWE<D> {
197    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
198        writer.write_u32::<LittleEndian>(self.k.into())?;
199        writer.write_u32::<LittleEndian>(self.base2k.into())?;
200        self.data.write_to(writer)
201    }
202}