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}