poulpy_core/layouts/
lwe.rs1use 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}