1use std::fmt;
  2
  3use poulpy_hal::{
  4    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, VecZnx, VecZnxToMut, VecZnxToRef, WriterTo, 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.base2k().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: VecZnx<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) -> &VecZnx<D> {
 94        &self.data
 95    }
 96}
 97
 98impl<D: DataMut> LWE<D> {
 99    pub fn data_mut(&mut self) -> &mut VecZnx<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!(f, "LWE: base2k={} k={}: {}", self.base2k().0, self.k().0, self.data)
113    }
114}
115
116impl<D: DataMut> FillUniform for LWE<D>
117where
118    VecZnx<D>: FillUniform,
119{
120    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
121        self.data.fill_uniform(log_bound, source);
122    }
123}
124
125impl LWE<Vec<u8>> {
126    pub fn alloc_from_infos<A>(infos: &A) -> Self
127    where
128        A: LWEInfos,
129    {
130        Self::alloc(infos.n(), infos.base2k(), infos.k())
131    }
132
133    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
134        LWE {
135            data: VecZnx::alloc((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize),
136            k,
137            base2k,
138        }
139    }
140
141    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
142    where
143        A: LWEInfos,
144    {
145        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
146    }
147
148    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
149        VecZnx::bytes_of((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize)
150    }
151}
152
153pub trait LWEToRef {
154    fn to_ref(&self) -> LWE<&[u8]>;
155}
156
157impl<D: DataRef> LWEToRef for LWE<D> {
158    fn to_ref(&self) -> LWE<&[u8]> {
159        LWE {
160            k: self.k,
161            base2k: self.base2k,
162            data: self.data.to_ref(),
163        }
164    }
165}
166
167pub trait LWEToMut {
168    #[allow(dead_code)]
169    fn to_mut(&mut self) -> LWE<&mut [u8]>;
170}
171
172impl<D: DataMut> LWEToMut for LWE<D> {
173    fn to_mut(&mut self) -> LWE<&mut [u8]> {
174        LWE {
175            k: self.k,
176            base2k: self.base2k,
177            data: self.data.to_mut(),
178        }
179    }
180}
181
182impl<D: DataMut> ReaderFrom for LWE<D> {
183    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
184        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
185        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
186        self.data.read_from(reader)
187    }
188}
189
190impl<D: DataRef> WriterTo for LWE<D> {
191    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
192        writer.write_u32::<LittleEndian>(self.k.into())?;
193        writer.write_u32::<LittleEndian>(self.base2k.into())?;
194        self.data.write_to(writer)
195    }
196}