1use backend::hal::{
  2    api::{MatZnxAlloc, MatZnxAllocBytes, VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
  4};
  5
  6use crate::{GLWECiphertext, Infos};
  7
  8pub trait GGSWLayoutFamily<B: Backend> = VmpPMatAlloc<B> + VmpPMatAllocBytes + VmpPMatPrepare<B>;
  9
 10#[derive(PartialEq, Eq)]
 11pub struct GGSWCiphertext<D: Data> {
 12    pub(crate) data: MatZnx<D>,
 13    pub(crate) basek: usize,
 14    pub(crate) k: usize,
 15    pub(crate) digits: usize,
 16}
 17
 18impl<D: DataRef> GGSWCiphertext<D> {
 19    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
 20        GLWECiphertext {
 21            data: self.data.at(row, col),
 22            basek: self.basek,
 23            k: self.k,
 24        }
 25    }
 26}
 27
 28impl<D: DataMut> GGSWCiphertext<D> {
 29    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
 30        GLWECiphertext {
 31            data: self.data.at_mut(row, col),
 32            basek: self.basek,
 33            k: self.k,
 34        }
 35    }
 36}
 37
 38impl GGSWCiphertext<Vec<u8>> {
 39    pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
 40    where
 41        Module<B>: MatZnxAlloc,
 42    {
 43        let size: usize = k.div_ceil(basek);
 44        debug_assert!(digits > 0, "invalid ggsw: `digits` == 0");
 45
 46        debug_assert!(
 47            size > digits,
 48            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 49            size,
 50            digits
 51        );
 52
 53        assert!(
 54            rows * digits <= size,
 55            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 56            rows,
 57            digits,
 58            size
 59        );
 60
 61        Self {
 62            data: module.mat_znx_alloc(rows, rank + 1, rank + 1, k.div_ceil(basek)),
 63            basek,
 64            k: k,
 65            digits,
 66        }
 67    }
 68
 69    pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
 70    where
 71        Module<B>: MatZnxAllocBytes,
 72    {
 73        let size: usize = k.div_ceil(basek);
 74        debug_assert!(
 75            size > digits,
 76            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 77            size,
 78            digits
 79        );
 80
 81        assert!(
 82            rows * digits <= size,
 83            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 84            rows,
 85            digits,
 86            size
 87        );
 88
 89        module.mat_znx_alloc_bytes(rows, rank + 1, rank + 1, size)
 90    }
 91}
 92
 93impl<D: Data> Infos for GGSWCiphertext<D> {
 94    type Inner = MatZnx<D>;
 95
 96    fn inner(&self) -> &Self::Inner {
 97        &self.data
 98    }
 99
100    fn basek(&self) -> usize {
101        self.basek
102    }
103
104    fn k(&self) -> usize {
105        self.k
106    }
107}
108
109impl<D: Data> GGSWCiphertext<D> {
110    pub fn rank(&self) -> usize {
111        self.data.cols_out() - 1
112    }
113
114    pub fn digits(&self) -> usize {
115        self.digits
116    }
117}
118
119use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
120
121impl<D: DataMut> ReaderFrom for GGSWCiphertext<D> {
122    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
123        self.k = reader.read_u64::<LittleEndian>()? as usize;
124        self.basek = reader.read_u64::<LittleEndian>()? as usize;
125        self.digits = reader.read_u64::<LittleEndian>()? as usize;
126        self.data.read_from(reader)
127    }
128}
129
130impl<D: DataRef> WriterTo for GGSWCiphertext<D> {
131    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
132        writer.write_u64::<LittleEndian>(self.k as u64)?;
133        writer.write_u64::<LittleEndian>(self.basek as u64)?;
134        writer.write_u64::<LittleEndian>(self.digits as u64)?;
135        self.data.write_to(writer)
136    }
137}