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