1use poulpy_hal::layouts::{Data, DataMut, DataRef, ReaderFrom, VecZnx, WriterTo, ZnxInfos};
  2
  3use crate::{
  4    dist::Distribution,
  5    layouts::{Base2K, BuildError, Degree, GLWEInfos, LWEInfos, Rank, TorusPrecision},
  6};
  7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  8
  9#[derive(PartialEq, Eq)]
 10pub struct GLWEPublicKey<D: Data> {
 11    pub(crate) data: VecZnx<D>,
 12    pub(crate) base2k: Base2K,
 13    pub(crate) k: TorusPrecision,
 14    pub(crate) dist: Distribution,
 15}
 16
 17#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 18pub struct GLWEPublicKeyLayout {
 19    pub n: Degree,
 20    pub base2k: Base2K,
 21    pub k: TorusPrecision,
 22    pub rank: Rank,
 23}
 24
 25impl<D: Data> LWEInfos for GLWEPublicKey<D> {
 26    fn base2k(&self) -> Base2K {
 27        self.base2k
 28    }
 29
 30    fn k(&self) -> TorusPrecision {
 31        self.k
 32    }
 33
 34    fn n(&self) -> Degree {
 35        Degree(self.data.n() as u32)
 36    }
 37
 38    fn size(&self) -> usize {
 39        self.data.size()
 40    }
 41}
 42
 43impl<D: Data> GLWEInfos for GLWEPublicKey<D> {
 44    fn rank(&self) -> Rank {
 45        Rank(self.data.cols() as u32 - 1)
 46    }
 47}
 48
 49impl LWEInfos for GLWEPublicKeyLayout {
 50    fn base2k(&self) -> Base2K {
 51        self.base2k
 52    }
 53
 54    fn k(&self) -> TorusPrecision {
 55        self.k
 56    }
 57
 58    fn n(&self) -> Degree {
 59        self.n
 60    }
 61
 62    fn size(&self) -> usize {
 63        self.k.0.div_ceil(self.base2k.0) as usize
 64    }
 65}
 66
 67impl GLWEInfos for GLWEPublicKeyLayout {
 68    fn rank(&self) -> Rank {
 69        self.rank
 70    }
 71}
 72
 73pub struct GLWEPublicKeyBuilder<D: Data> {
 74    data: Option<VecZnx<D>>,
 75    base2k: Option<Base2K>,
 76    k: Option<TorusPrecision>,
 77}
 78
 79impl<D: Data> GLWEPublicKey<D> {
 80    #[inline]
 81    pub fn builder() -> GLWEPublicKeyBuilder<D> {
 82        GLWEPublicKeyBuilder {
 83            data: None,
 84            base2k: None,
 85            k: None,
 86        }
 87    }
 88}
 89
 90impl GLWEPublicKeyBuilder<Vec<u8>> {
 91    #[inline]
 92    pub fn layout<A>(mut self, layout: &A) -> Self
 93    where
 94        A: GLWEInfos,
 95    {
 96        self.data = Some(VecZnx::alloc(
 97            layout.n().into(),
 98            (layout.rank() + 1).into(),
 99            layout.size(),
100        ));
101        self.base2k = Some(layout.base2k());
102        self.k = Some(layout.k());
103        self
104    }
105}
106
107impl<D: Data> GLWEPublicKeyBuilder<D> {
108    #[inline]
109    pub fn data(mut self, data: VecZnx<D>) -> Self {
110        self.data = Some(data);
111        self
112    }
113    #[inline]
114    pub fn base2k(mut self, base2k: Base2K) -> Self {
115        self.base2k = Some(base2k);
116        self
117    }
118    #[inline]
119    pub fn k(mut self, k: TorusPrecision) -> Self {
120        self.k = Some(k);
121        self
122    }
123
124    pub fn build(self) -> Result<GLWEPublicKey<D>, BuildError> {
125        let data: VecZnx<D> = self.data.ok_or(BuildError::MissingData)?;
126        let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
127        let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
128
129        if base2k == 0_u32 {
130            return Err(BuildError::ZeroBase2K);
131        }
132
133        if k == 0_u32 {
134            return Err(BuildError::ZeroTorusPrecision);
135        }
136
137        if data.n() == 0 {
138            return Err(BuildError::ZeroDegree);
139        }
140
141        if data.cols() == 0 {
142            return Err(BuildError::ZeroCols);
143        }
144
145        if data.size() == 0 {
146            return Err(BuildError::ZeroLimbs);
147        }
148
149        Ok(GLWEPublicKey {
150            data,
151            base2k,
152            k,
153            dist: Distribution::NONE,
154        })
155    }
156}
157
158impl GLWEPublicKey<Vec<u8>> {
159    pub fn alloc<A>(infos: &A) -> Self
160    where
161        A: GLWEInfos,
162    {
163        Self::alloc_with(infos.n(), infos.base2k(), infos.k(), infos.rank())
164    }
165
166    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
167        Self {
168            data: VecZnx::alloc(n.into(), (rank + 1).into(), k.0.div_ceil(base2k.0) as usize),
169            base2k,
170            k,
171            dist: Distribution::NONE,
172        }
173    }
174
175    pub fn alloc_bytes<A>(infos: &A) -> usize
176    where
177        A: GLWEInfos,
178    {
179        Self::alloc_bytes_with(infos.n(), infos.base2k(), infos.k(), infos.rank())
180    }
181
182    pub fn alloc_bytes_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
183        VecZnx::alloc_bytes(n.into(), (rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
184    }
185}
186
187impl<D: DataMut> ReaderFrom for GLWEPublicKey<D> {
188    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
189        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
190        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
191        match Distribution::read_from(reader) {
192            Ok(dist) => self.dist = dist,
193            Err(e) => return Err(e),
194        }
195        self.data.read_from(reader)
196    }
197}
198
199impl<D: DataRef> WriterTo for GLWEPublicKey<D> {
200    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
201        writer.write_u32::<LittleEndian>(self.k.0)?;
202        writer.write_u32::<LittleEndian>(self.base2k.0)?;
203        match self.dist.write_to(writer) {
204            Ok(()) => {}
205            Err(e) => return Err(e),
206        }
207        self.data.write_to(writer)
208    }
209}