1use backend::hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
  4};
  5
  6use crate::{GGSWCiphertext, GGSWLayoutFamily, Infos};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GGSWCiphertextExec<D: Data, B: Backend> {
 10    pub(crate) data: VmpPMat<D, B>,
 11    pub(crate) basek: usize,
 12    pub(crate) k: usize,
 13    pub(crate) digits: usize,
 14}
 15
 16impl<B: Backend> GGSWCiphertextExec<Vec<u8>, B> {
 17    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
 18    where
 19        Module<B>: GGSWLayoutFamily<B>,
 20    {
 21        let size: usize = k.div_ceil(basek);
 22        debug_assert!(digits > 0, "invalid ggsw: `digits` == 0");
 23
 24        debug_assert!(
 25            size > digits,
 26            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 27            size,
 28            digits
 29        );
 30
 31        assert!(
 32            rows * digits <= size,
 33            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 34            rows,
 35            digits,
 36            size
 37        );
 38
 39        Self {
 40            data: module.vmp_pmat_alloc(rows, rank + 1, rank + 1, k.div_ceil(basek)),
 41            basek,
 42            k: k,
 43            digits,
 44        }
 45    }
 46
 47    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
 48    where
 49        Module<B>: GGSWLayoutFamily<B>,
 50    {
 51        let size: usize = k.div_ceil(basek);
 52        debug_assert!(
 53            size > digits,
 54            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 55            size,
 56            digits
 57        );
 58
 59        assert!(
 60            rows * digits <= size,
 61            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 62            rows,
 63            digits,
 64            size
 65        );
 66
 67        module.vmp_pmat_alloc_bytes(rows, rank + 1, rank + 1, size)
 68    }
 69
 70    pub fn from<DataOther: DataRef>(
 71        module: &Module<B>,
 72        other: &GGSWCiphertext<DataOther>,
 73        scratch: &mut Scratch<B>,
 74    ) -> GGSWCiphertextExec<Vec<u8>, B>
 75    where
 76        Module<B>: GGSWLayoutFamily<B>,
 77    {
 78        let mut ggsw_exec: GGSWCiphertextExec<Vec<u8>, B> = Self::alloc(
 79            module,
 80            other.basek(),
 81            other.k(),
 82            other.rows(),
 83            other.digits(),
 84            other.rank(),
 85        );
 86        ggsw_exec.prepare(module, other, scratch);
 87        ggsw_exec
 88    }
 89}
 90
 91impl<D: Data, B: Backend> Infos for GGSWCiphertextExec<D, B> {
 92    type Inner = VmpPMat<D, B>;
 93
 94    fn inner(&self) -> &Self::Inner {
 95        &self.data
 96    }
 97
 98    fn basek(&self) -> usize {
 99        self.basek
100    }
101
102    fn k(&self) -> usize {
103        self.k
104    }
105}
106
107impl<D: Data, B: Backend> GGSWCiphertextExec<D, B> {
108    pub fn rank(&self) -> usize {
109        self.data.cols_out() - 1
110    }
111
112    pub fn digits(&self) -> usize {
113        self.digits
114    }
115}
116
117impl<DataSelf: DataMut, B: Backend> GGSWCiphertextExec<DataSelf, B> {
118    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGSWCiphertext<DataOther>, scratch: &mut Scratch<B>)
119    where
120        DataOther: DataRef,
121        Module<B>: GGSWLayoutFamily<B>,
122    {
123        module.vmp_prepare(&mut self.data, &other.data, scratch);
124        self.k = other.k;
125        self.basek = other.basek;
126        self.digits = other.digits;
127    }
128}