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, GLWEInfos, GetRingDegree, LWEInfos, Rank, RingDegree, 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: RingDegree,
 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) -> RingDegree {
 47        RingDegree(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) -> RingDegree {
 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: GetRingDegree,
 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.ring_degree().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.ring_degree().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<B: Backend> GLWEPublicKeyAlloc for Module<B> where Self: GetRingDegree {}
126
127impl GLWEPublicKey<Vec<u8>> {
128    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
129    where
130        A: GLWEInfos,
131        M: GLWEPublicKeyAlloc,
132    {
133        module.alloc_glwe_public_key_from_infos(infos)
134    }
135
136    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
137    where
138        M: GLWEPublicKeyAlloc,
139    {
140        module.alloc_glwe_public_key(base2k, k, rank)
141    }
142
143    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
144    where
145        A: GLWEInfos,
146        M: GLWEPublicKeyAlloc,
147    {
148        module.bytes_of_glwe_public_key_from_infos(infos)
149    }
150
151    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
152    where
153        M: GLWEPublicKeyAlloc,
154    {
155        module.bytes_of_glwe_public_key(base2k, k, rank)
156    }
157}
158
159impl<D: DataMut> ReaderFrom for GLWEPublicKey<D> {
160    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
161        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
162        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
163        match Distribution::read_from(reader) {
164            Ok(dist) => self.dist = dist,
165            Err(e) => return Err(e),
166        }
167        self.data.read_from(reader)
168    }
169}
170
171impl<D: DataRef> WriterTo for GLWEPublicKey<D> {
172    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
173        writer.write_u32::<LittleEndian>(self.k.0)?;
174        writer.write_u32::<LittleEndian>(self.base2k.0)?;
175        match self.dist.write_to(writer) {
176            Ok(()) => {}
177            Err(e) => return Err(e),
178        }
179        self.data.write_to(writer)
180    }
181}
182
183pub trait GLWEPublicKeyToRef {
184    fn to_ref(&self) -> GLWEPublicKey<&[u8]>;
185}
186
187impl<D: DataRef> GLWEPublicKeyToRef for GLWEPublicKey<D> {
188    fn to_ref(&self) -> GLWEPublicKey<&[u8]> {
189        GLWEPublicKey {
190            data: self.data.to_ref(),
191            base2k: self.base2k,
192            k: self.k,
193            dist: self.dist,
194        }
195    }
196}
197
198pub trait GLWEPublicKeyToMut {
199    fn to_mut(&mut self) -> GLWEPublicKey<&mut [u8]>;
200}
201
202impl<D: DataMut> GLWEPublicKeyToMut for GLWEPublicKey<D> {
203    fn to_mut(&mut self) -> GLWEPublicKey<&mut [u8]> {
204        GLWEPublicKey {
205            base2k: self.base2k,
206            k: self.k,
207            dist: self.dist,
208            data: self.data.to_mut(),
209        }
210    }
211}