1use backend::hal::{
  2    api::{VecZnxAlloc, VecZnxAllocBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, VecZnx, VecZnxToMut, VecZnxToRef, WriterTo},
  4};
  5
  6use crate::{GLWEOps, Infos, SetMetaData};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GLWECiphertext<D: Data> {
 10    pub data: VecZnx<D>,
 11    pub basek: usize,
 12    pub k: usize,
 13}
 14
 15impl GLWECiphertext<Vec<u8>> {
 16    pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self
 17    where
 18        Module<B>: VecZnxAlloc,
 19    {
 20        Self {
 21            data: module.vec_znx_alloc(rank + 1, k.div_ceil(basek)),
 22            basek,
 23            k,
 24        }
 25    }
 26
 27    pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
 28    where
 29        Module<B>: VecZnxAllocBytes,
 30    {
 31        module.vec_znx_alloc_bytes(rank + 1, k.div_ceil(basek))
 32    }
 33}
 34
 35impl<D: Data> Infos for GLWECiphertext<D> {
 36    type Inner = VecZnx<D>;
 37
 38    fn inner(&self) -> &Self::Inner {
 39        &self.data
 40    }
 41
 42    fn basek(&self) -> usize {
 43        self.basek
 44    }
 45
 46    fn k(&self) -> usize {
 47        self.k
 48    }
 49}
 50
 51impl<D: Data> GLWECiphertext<D> {
 52    pub fn rank(&self) -> usize {
 53        self.cols() - 1
 54    }
 55}
 56
 57impl<D: DataRef> GLWECiphertext<D> {
 58    pub fn clone(&self) -> GLWECiphertext<Vec<u8>> {
 59        GLWECiphertext {
 60            data: self.data.clone(),
 61            basek: self.basek(),
 62            k: self.k(),
 63        }
 64    }
 65}
 66
 67impl<D: DataMut> SetMetaData for GLWECiphertext<D> {
 68    fn set_k(&mut self, k: usize) {
 69        self.k = k
 70    }
 71
 72    fn set_basek(&mut self, basek: usize) {
 73        self.basek = basek
 74    }
 75}
 76
 77pub trait GLWECiphertextToRef: Infos {
 78    fn to_ref(&self) -> GLWECiphertext<&[u8]>;
 79}
 80
 81impl<D: DataRef> GLWECiphertextToRef for GLWECiphertext<D> {
 82    fn to_ref(&self) -> GLWECiphertext<&[u8]> {
 83        GLWECiphertext {
 84            data: self.data.to_ref(),
 85            basek: self.basek,
 86            k: self.k,
 87        }
 88    }
 89}
 90
 91pub trait GLWECiphertextToMut: Infos {
 92    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]>;
 93}
 94
 95impl<D: DataMut> GLWECiphertextToMut for GLWECiphertext<D> {
 96    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]> {
 97        GLWECiphertext {
 98            data: self.data.to_mut(),
 99            basek: self.basek,
100            k: self.k,
101        }
102    }
103}
104
105impl<D: DataMut> GLWEOps for GLWECiphertext<D> where GLWECiphertext<D>: GLWECiphertextToMut + Infos + SetMetaData {}
106
107use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
108
109impl<D: DataMut> ReaderFrom for GLWECiphertext<D> {
110    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
111        self.k = reader.read_u64::<LittleEndian>()? as usize;
112        self.basek = reader.read_u64::<LittleEndian>()? as usize;
113        self.data.read_from(reader)
114    }
115}
116
117impl<D: DataRef> WriterTo for GLWECiphertext<D> {
118    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
119        writer.write_u64::<LittleEndian>(self.k as u64)?;
120        writer.write_u64::<LittleEndian>(self.basek as u64)?;
121        self.data.write_to(writer)
122    }
123}