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