1use poulpy_hal::layouts::{Data, DataMut, DataRef, ReaderFrom, VecZnx, VecZnxToMut, VecZnxToRef, WriterTo, ZnxInfos};
  2
  3use crate::{
  4    GetDistribution, GetDistributionMut,
  5    dist::Distribution,
  6    layouts::{Base2K, Degree, GLWEInfos, LWEInfos, Rank, TorusPrecision},
  7};
  8use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  9
 10#[derive(PartialEq, Eq)]
 11pub struct GLWEPublicKey<D: Data> {
 12    pub(crate) data: VecZnx<D>,
 13    pub(crate) base2k: Base2K,
 14    pub(crate) k: TorusPrecision,
 15    pub(crate) dist: Distribution,
 16}
 17
 18impl<D: DataMut> GetDistributionMut for GLWEPublicKey<D> {
 19    fn dist_mut(&mut self) -> &mut Distribution {
 20        &mut self.dist
 21    }
 22}
 23
 24impl<D: DataRef> GetDistribution for GLWEPublicKey<D> {
 25    fn dist(&self) -> &Distribution {
 26        &self.dist
 27    }
 28}
 29
 30#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 31pub struct GLWEPublicKeyLayout {
 32    pub n: Degree,
 33    pub base2k: Base2K,
 34    pub k: TorusPrecision,
 35    pub rank: Rank,
 36}
 37
 38impl<D: Data> LWEInfos for GLWEPublicKey<D> {
 39    fn base2k(&self) -> Base2K {
 40        self.base2k
 41    }
 42
 43    fn k(&self) -> TorusPrecision {
 44        self.k
 45    }
 46
 47    fn n(&self) -> Degree {
 48        Degree(self.data.n() as u32)
 49    }
 50
 51    fn size(&self) -> usize {
 52        self.data.size()
 53    }
 54}
 55
 56impl<D: Data> GLWEInfos for GLWEPublicKey<D> {
 57    fn rank(&self) -> Rank {
 58        Rank(self.data.cols() as u32 - 1)
 59    }
 60}
 61
 62impl LWEInfos for GLWEPublicKeyLayout {
 63    fn base2k(&self) -> Base2K {
 64        self.base2k
 65    }
 66
 67    fn k(&self) -> TorusPrecision {
 68        self.k
 69    }
 70
 71    fn n(&self) -> Degree {
 72        self.n
 73    }
 74
 75    fn size(&self) -> usize {
 76        self.k.0.div_ceil(self.base2k.0) as usize
 77    }
 78}
 79
 80impl GLWEInfos for GLWEPublicKeyLayout {
 81    fn rank(&self) -> Rank {
 82        self.rank
 83    }
 84}
 85
 86impl GLWEPublicKey<Vec<u8>> {
 87    pub fn alloc_from_infos<A>(infos: &A) -> Self
 88    where
 89        A: GLWEInfos,
 90    {
 91        Self::alloc(infos.n(), infos.base2k(), infos.k(), infos.rank())
 92    }
 93
 94    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 95        GLWEPublicKey {
 96            data: VecZnx::alloc(n.into(), (rank + 1).into(), k.0.div_ceil(base2k.0) as usize),
 97            base2k,
 98            k,
 99            dist: Distribution::NONE,
100        }
101    }
102
103    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
104    where
105        A: GLWEInfos,
106    {
107        Self::bytes_of(infos.n(), infos.base2k(), infos.k(), infos.rank())
108    }
109
110    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
111        VecZnx::bytes_of(n.into(), (rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
112    }
113}
114
115impl<D: DataMut> ReaderFrom for GLWEPublicKey<D> {
116    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
117        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
118        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
119        match Distribution::read_from(reader) {
120            Ok(dist) => self.dist = dist,
121            Err(e) => return Err(e),
122        }
123        self.data.read_from(reader)
124    }
125}
126
127impl<D: DataRef> WriterTo for GLWEPublicKey<D> {
128    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
129        writer.write_u32::<LittleEndian>(self.k.0)?;
130        writer.write_u32::<LittleEndian>(self.base2k.0)?;
131        match self.dist.write_to(writer) {
132            Ok(()) => {}
133            Err(e) => return Err(e),
134        }
135        self.data.write_to(writer)
136    }
137}
138
139pub trait GLWEPublicKeyToRef {
140    fn to_ref(&self) -> GLWEPublicKey<&[u8]>;
141}
142
143impl<D: DataRef> GLWEPublicKeyToRef for GLWEPublicKey<D> {
144    fn to_ref(&self) -> GLWEPublicKey<&[u8]> {
145        GLWEPublicKey {
146            data: self.data.to_ref(),
147            base2k: self.base2k,
148            k: self.k,
149            dist: self.dist,
150        }
151    }
152}
153
154pub trait GLWEPublicKeyToMut {
155    fn to_mut(&mut self) -> GLWEPublicKey<&mut [u8]>;
156}
157
158impl<D: DataMut> GLWEPublicKeyToMut for GLWEPublicKey<D> {
159    fn to_mut(&mut self) -> GLWEPublicKey<&mut [u8]> {
160        GLWEPublicKey {
161            base2k: self.base2k,
162            k: self.k,
163            dist: self.dist,
164            data: self.data.to_mut(),
165        }
166    }
167}