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