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, BuildError, 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) -> LWECiphertextLayout {
 22        LWECiphertextLayout {
 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 LWECiphertextLayout {
 37    pub n: Degree,
 38    pub k: TorusPrecision,
 39    pub base2k: Base2K,
 40}
 41
 42impl LWEInfos for LWECiphertextLayout {
 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 LWECiphertext<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 LWECiphertext<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 LWECiphertext<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> LWECiphertext<D> {
 90    pub fn data(&self) -> &Zn<D> {
 91        &self.data
 92    }
 93}
 94
 95impl<D: DataMut> LWECiphertext<D> {
 96    pub fn data_mut(&mut self) -> &Zn<D> {
 97        &mut self.data
 98    }
 99}
100
101impl<D: DataRef> fmt::Debug for LWECiphertext<D> {
102    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
103        write!(f, "{self}")
104    }
105}
106
107impl<D: DataRef> fmt::Display for LWECiphertext<D> {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        write!(
110            f,
111            "LWECiphertext: base2k={} k={}: {}",
112            self.base2k().0,
113            self.k().0,
114            self.data
115        )
116    }
117}
118
119impl<D: DataMut> FillUniform for LWECiphertext<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 LWECiphertext<Vec<u8>> {
129    pub fn alloc<A>(infos: &A) -> Self
130    where
131        A: LWEInfos,
132    {
133        Self::alloc_with(infos.n(), infos.base2k(), infos.k())
134    }
135
136    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
137        Self {
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 alloc_bytes<A>(infos: &A) -> usize
145    where
146        A: LWEInfos,
147    {
148        Self::alloc_bytes_with(infos.n(), infos.base2k(), infos.k())
149    }
150
151    pub fn alloc_bytes_with(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
152        Zn::alloc_bytes((n + 1).into(), 1, k.0.div_ceil(base2k.0) as usize)
153    }
154}
155
156impl LWECiphertextBuilder<Vec<u8>> {
157    #[inline]
158    pub fn layout<A>(mut self, layout: A) -> Self
159    where
160        A: LWEInfos,
161    {
162        self.data = Some(Zn::alloc((layout.n() + 1).into(), 1, layout.size()));
163        self.base2k = Some(layout.base2k());
164        self.k = Some(layout.k());
165        self
166    }
167}
168
169pub struct LWECiphertextBuilder<D: Data> {
170    data: Option<Zn<D>>,
171    base2k: Option<Base2K>,
172    k: Option<TorusPrecision>,
173}
174
175impl<D: Data> LWECiphertext<D> {
176    #[inline]
177    pub fn builder() -> LWECiphertextBuilder<D> {
178        LWECiphertextBuilder {
179            data: None,
180            base2k: None,
181            k: None,
182        }
183    }
184}
185
186impl<D: Data> LWECiphertextBuilder<D> {
187    #[inline]
188    pub fn data(mut self, data: Zn<D>) -> Self {
189        self.data = Some(data);
190        self
191    }
192    #[inline]
193    pub fn base2k(mut self, base2k: Base2K) -> Self {
194        self.base2k = Some(base2k);
195        self
196    }
197    #[inline]
198    pub fn k(mut self, k: TorusPrecision) -> Self {
199        self.k = Some(k);
200        self
201    }
202
203    pub fn build(self) -> Result<LWECiphertext<D>, BuildError> {
204        let data: Zn<D> = self.data.ok_or(BuildError::MissingData)?;
205        let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
206        let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
207
208        if base2k.0 == 0 {
209            return Err(BuildError::ZeroBase2K);
210        }
211
212        if k.0 == 0 {
213            return Err(BuildError::ZeroTorusPrecision);
214        }
215
216        if data.n() == 0 {
217            return Err(BuildError::ZeroDegree);
218        }
219
220        if data.cols() == 0 {
221            return Err(BuildError::ZeroCols);
222        }
223
224        if data.size() == 0 {
225            return Err(BuildError::ZeroLimbs);
226        }
227
228        Ok(LWECiphertext { data, base2k, k })
229    }
230}
231
232pub trait LWECiphertextToRef {
233    fn to_ref(&self) -> LWECiphertext<&[u8]>;
234}
235
236impl<D: DataRef> LWECiphertextToRef for LWECiphertext<D> {
237    fn to_ref(&self) -> LWECiphertext<&[u8]> {
238        LWECiphertext {
239            k: self.k,
240            base2k: self.base2k,
241            data: self.data.to_ref(),
242        }
243    }
244}
245
246pub trait LWECiphertextToMut {
247    #[allow(dead_code)]
248    fn to_mut(&mut self) -> LWECiphertext<&mut [u8]>;
249}
250
251impl<D: DataMut> LWECiphertextToMut for LWECiphertext<D> {
252    fn to_mut(&mut self) -> LWECiphertext<&mut [u8]> {
253        LWECiphertext {
254            k: self.k,
255            base2k: self.base2k,
256            data: self.data.to_mut(),
257        }
258    }
259}
260
261impl<D: DataMut> ReaderFrom for LWECiphertext<D> {
262    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
263        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
264        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
265        self.data.read_from(reader)
266    }
267}
268
269impl<D: DataRef> WriterTo for LWECiphertext<D> {
270    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
271        writer.write_u32::<LittleEndian>(self.k.into())?;
272        writer.write_u32::<LittleEndian>(self.base2k.into())?;
273        self.data.write_to(writer)
274    }
275}