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>, n: usize, 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(n, 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>, n: usize, 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(n, 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.n(),
 81            other.basek(),
 82            other.k(),
 83            other.rows(),
 84            other.digits(),
 85            other.rank(),
 86        );
 87        ggsw_exec.prepare(module, other, scratch);
 88        ggsw_exec
 89    }
 90}
 91
 92impl<D: Data, B: Backend> Infos for GGSWCiphertextExec<D, B> {
 93    type Inner = VmpPMat<D, B>;
 94
 95    fn inner(&self) -> &Self::Inner {
 96        &self.data
 97    }
 98
 99    fn basek(&self) -> usize {
100        self.basek
101    }
102
103    fn k(&self) -> usize {
104        self.k
105    }
106}
107
108impl<D: Data, B: Backend> GGSWCiphertextExec<D, B> {
109    pub fn rank(&self) -> usize {
110        self.data.cols_out() - 1
111    }
112
113    pub fn digits(&self) -> usize {
114        self.digits
115    }
116}
117
118impl<DataSelf: DataMut, B: Backend> GGSWCiphertextExec<DataSelf, B> {
119    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGSWCiphertext<DataOther>, scratch: &mut Scratch<B>)
120    where
121        DataOther: DataRef,
122        Module<B>: GGSWLayoutFamily<B>,
123    {
124        module.vmp_prepare(&mut self.data, &other.data, scratch);
125        self.k = other.k;
126        self.basek = other.basek;
127        self.digits = other.digits;
128    }
129}