1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  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    #[allow(clippy::too_many_arguments)]
 21    pub fn alloc(
 22        module: &Module<B>,
 23        n: usize,
 24        basek: usize,
 25        k: usize,
 26        rows: usize,
 27        digits: usize,
 28        rank_in: usize,
 29        rank_out: usize,
 30    ) -> Self
 31    where
 32        Module<B>: VmpPMatAlloc<B>,
 33    {
 34        let size: usize = k.div_ceil(basek);
 35        debug_assert!(
 36            size > digits,
 37            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 38            size,
 39            digits
 40        );
 41
 42        assert!(
 43            rows * digits <= size,
 44            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 45            rows,
 46            digits,
 47            size
 48        );
 49
 50        Self {
 51            data: module.vmp_pmat_alloc(n, rows, rank_in, rank_out + 1, size),
 52            basek,
 53            k,
 54            digits,
 55        }
 56    }
 57
 58    #[allow(clippy::too_many_arguments)]
 59    pub fn bytes_of(
 60        module: &Module<B>,
 61        n: usize,
 62        basek: usize,
 63        k: usize,
 64        rows: usize,
 65        digits: usize,
 66        rank_in: usize,
 67        rank_out: usize,
 68    ) -> usize
 69    where
 70        Module<B>: VmpPMatAllocBytes,
 71    {
 72        let size: usize = k.div_ceil(basek);
 73        debug_assert!(
 74            size > digits,
 75            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 76            size,
 77            digits
 78        );
 79
 80        assert!(
 81            rows * digits <= size,
 82            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 83            rows,
 84            digits,
 85            size
 86        );
 87
 88        module.vmp_pmat_alloc_bytes(n, rows, rank_in, rank_out + 1, rows)
 89    }
 90}
 91
 92impl<D: Data, B: Backend> Infos for GGLWECiphertextPrepared<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> GGLWECiphertextPrepared<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    pub fn rank_in(&self) -> usize {
118        self.data.cols_in()
119    }
120
121    pub fn rank_out(&self) -> usize {
122        self.data.cols_out() - 1
123    }
124}
125
126impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWECiphertext<DR>> for GGLWECiphertextPrepared<D, B>
127where
128    Module<B>: VmpPrepare<B>,
129{
130    fn prepare(&mut self, module: &Module<B>, other: &GGLWECiphertext<DR>, scratch: &mut Scratch<B>) {
131        module.vmp_prepare(&mut self.data, &other.data, scratch);
132        self.basek = other.basek;
133        self.k = other.k;
134        self.digits = other.digits;
135    }
136}
137
138impl<D: DataRef, B: Backend> PrepareAlloc<B, GGLWECiphertextPrepared<Vec<u8>, B>> for GGLWECiphertext<D>
139where
140    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
141{
142    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GGLWECiphertextPrepared<Vec<u8>, B> {
143        let mut atk_prepared: GGLWECiphertextPrepared<Vec<u8>, B> = GGLWECiphertextPrepared::alloc(
144            module,
145            self.n(),
146            self.basek(),
147            self.k(),
148            self.rows(),
149            self.digits(),
150            self.rank_in(),
151            self.rank_out(),
152        );
153        atk_prepared.prepare(module, self, scratch);
154        atk_prepared
155    }
156}