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 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
128impl LWE<Vec<u8>> {
129    pub fn alloc_from_infos<A>(infos: &A) -> Self
130    where
131        A: LWEInfos,
132    {
133        Self::alloc(infos.n(), infos.base2k(), infos.k())
134    }
135
136    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
137        LWE {
138            data: Zn::alloc((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize),
139            k,
140            base2k,
141        }
142    }
143
144    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
145    where
146        A: LWEInfos,
147    {
148        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
149    }
150
151    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
152        Zn::bytes_of((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize)
153    }
154}
155
156pub trait LWEToRef {
157    fn to_ref(&self) -> LWE<&[u8]>;
158}
159
160impl<D: DataRef> LWEToRef for LWE<D> {
161    fn to_ref(&self) -> LWE<&[u8]> {
162        LWE {
163            k: self.k,
164            base2k: self.base2k,
165            data: self.data.to_ref(),
166        }
167    }
168}
169
170pub trait LWEToMut {
171    #[allow(dead_code)]
172    fn to_mut(&mut self) -> LWE<&mut [u8]>;
173}
174
175impl<D: DataMut> LWEToMut for LWE<D> {
176    fn to_mut(&mut self) -> LWE<&mut [u8]> {
177        LWE {
178            k: self.k,
179            base2k: self.base2k,
180            data: self.data.to_mut(),
181        }
182    }
183}
184
185impl<D: DataMut> ReaderFrom for LWE<D> {
186    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
187        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
188        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
189        self.data.read_from(reader)
190    }
191}
192
193impl<D: DataRef> WriterTo for LWE<D> {
194    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
195        writer.write_u32::<LittleEndian>(self.k.into())?;
196        writer.write_u32::<LittleEndian>(self.base2k.into())?;
197        self.data.write_to(writer)
198    }
199}