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