1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare, VmpPrepareTmpBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, ZnxInfos},
  4    oep::VmpPMatAllocBytesImpl,
  5};
  6
  7use crate::layouts::{
  8    Base2K, BuildError, Degree, Dnum, Dsize, GGLWECiphertext, GGLWEInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  9    prepared::{Prepare, PrepareAlloc, PrepareScratchSpace},
 10};
 11
 12#[derive(PartialEq, Eq)]
 13pub struct GGLWECiphertextPrepared<D: Data, B: Backend> {
 14    pub(crate) data: VmpPMat<D, B>,
 15    pub(crate) k: TorusPrecision,
 16    pub(crate) base2k: Base2K,
 17    pub(crate) dsize: Dsize,
 18}
 19
 20impl<D: Data, B: Backend> LWEInfos for GGLWECiphertextPrepared<D, B> {
 21    fn n(&self) -> Degree {
 22        Degree(self.data.n() as u32)
 23    }
 24
 25    fn base2k(&self) -> Base2K {
 26        self.base2k
 27    }
 28
 29    fn k(&self) -> TorusPrecision {
 30        self.k
 31    }
 32
 33    fn size(&self) -> usize {
 34        self.data.size()
 35    }
 36}
 37
 38impl<D: Data, B: Backend> GLWEInfos for GGLWECiphertextPrepared<D, B> {
 39    fn rank(&self) -> Rank {
 40        self.rank_out()
 41    }
 42}
 43
 44impl<D: Data, B: Backend> GGLWEInfos for GGLWECiphertextPrepared<D, B> {
 45    fn rank_in(&self) -> Rank {
 46        Rank(self.data.cols_in() as u32)
 47    }
 48
 49    fn rank_out(&self) -> Rank {
 50        Rank(self.data.cols_out() as u32 - 1)
 51    }
 52
 53    fn dsize(&self) -> Dsize {
 54        self.dsize
 55    }
 56
 57    fn dnum(&self) -> Dnum {
 58        Dnum(self.data.rows() as u32)
 59    }
 60}
 61
 62pub struct GGLWECiphertextPreparedBuilder<D: Data, B: Backend> {
 63    data: Option<VmpPMat<D, B>>,
 64    base2k: Option<Base2K>,
 65    k: Option<TorusPrecision>,
 66    dsize: Option<Dsize>,
 67}
 68
 69impl<D: Data, B: Backend> GGLWECiphertextPrepared<D, B> {
 70    #[inline]
 71    pub fn builder() -> GGLWECiphertextPreparedBuilder<D, B> {
 72        GGLWECiphertextPreparedBuilder {
 73            data: None,
 74            base2k: None,
 75            k: None,
 76            dsize: None,
 77        }
 78    }
 79}
 80
 81impl<B: Backend> GGLWECiphertextPreparedBuilder<Vec<u8>, B> {
 82    #[inline]
 83    pub fn layout<A>(mut self, infos: &A) -> Self
 84    where
 85        A: GGLWEInfos,
 86        B: VmpPMatAllocBytesImpl<B>,
 87    {
 88        self.data = Some(VmpPMat::alloc(
 89            infos.n().into(),
 90            infos.dnum().into(),
 91            infos.rank_in().into(),
 92            (infos.rank_out() + 1).into(),
 93            infos.size(),
 94        ));
 95        self.base2k = Some(infos.base2k());
 96        self.k = Some(infos.k());
 97        self.dsize = Some(infos.dsize());
 98        self
 99    }
100}
101
102impl<D: Data, B: Backend> GGLWECiphertextPreparedBuilder<D, B> {
103    #[inline]
104    pub fn data(mut self, data: VmpPMat<D, B>) -> Self {
105        self.data = Some(data);
106        self
107    }
108    #[inline]
109    pub fn base2k(mut self, base2k: Base2K) -> Self {
110        self.base2k = Some(base2k);
111        self
112    }
113    #[inline]
114    pub fn k(mut self, k: TorusPrecision) -> Self {
115        self.k = Some(k);
116        self
117    }
118
119    #[inline]
120    pub fn dsize(mut self, dsize: Dsize) -> Self {
121        self.dsize = Some(dsize);
122        self
123    }
124
125    pub fn build(self) -> Result<GGLWECiphertextPrepared<D, B>, BuildError> {
126        let data: VmpPMat<D, B> = self.data.ok_or(BuildError::MissingData)?;
127        let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
128        let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
129        let dsize: Dsize = self.dsize.ok_or(BuildError::MissingDigits)?;
130
131        if base2k == 0_u32 {
132            return Err(BuildError::ZeroBase2K);
133        }
134
135        if dsize == 0_u32 {
136            return Err(BuildError::ZeroBase2K);
137        }
138
139        if k == 0_u32 {
140            return Err(BuildError::ZeroTorusPrecision);
141        }
142
143        if data.n() == 0 {
144            return Err(BuildError::ZeroDegree);
145        }
146
147        if data.cols() == 0 {
148            return Err(BuildError::ZeroCols);
149        }
150
151        if data.size() == 0 {
152            return Err(BuildError::ZeroLimbs);
153        }
154
155        Ok(GGLWECiphertextPrepared {
156            data,
157            base2k,
158            k,
159            dsize,
160        })
161    }
162}
163
164impl<B: Backend> GGLWECiphertextPrepared<Vec<u8>, B> {
165    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
166    where
167        A: GGLWEInfos,
168        Module<B>: VmpPMatAlloc<B>,
169    {
170        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
171        Self::alloc_with(
172            module,
173            infos.base2k(),
174            infos.k(),
175            infos.rank_in(),
176            infos.rank_out(),
177            infos.dnum(),
178            infos.dsize(),
179        )
180    }
181
182    pub fn alloc_with(
183        module: &Module<B>,
184        base2k: Base2K,
185        k: TorusPrecision,
186        rank_in: Rank,
187        rank_out: Rank,
188        dnum: Dnum,
189        dsize: Dsize,
190    ) -> Self
191    where
192        Module<B>: VmpPMatAlloc<B>,
193    {
194        let size: usize = k.0.div_ceil(base2k.0) as usize;
195        debug_assert!(
196            size as u32 > dsize.0,
197            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
198            dsize.0
199        );
200
201        assert!(
202            dnum.0 * dsize.0 <= size as u32,
203            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
204            dnum.0,
205            dsize.0,
206        );
207
208        Self {
209            data: module.vmp_pmat_alloc(dnum.into(), rank_in.into(), (rank_out + 1).into(), size),
210            k,
211            base2k,
212            dsize,
213        }
214    }
215
216    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
217    where
218        A: GGLWEInfos,
219        Module<B>: VmpPMatAllocBytes,
220    {
221        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
222        Self::alloc_bytes_with(
223            module,
224            infos.base2k(),
225            infos.k(),
226            infos.rank_in(),
227            infos.rank_out(),
228            infos.dnum(),
229            infos.dsize(),
230        )
231    }
232
233    pub fn alloc_bytes_with(
234        module: &Module<B>,
235        base2k: Base2K,
236        k: TorusPrecision,
237        rank_in: Rank,
238        rank_out: Rank,
239        dnum: Dnum,
240        dsize: Dsize,
241    ) -> usize
242    where
243        Module<B>: VmpPMatAllocBytes,
244    {
245        let size: usize = k.0.div_ceil(base2k.0) as usize;
246        debug_assert!(
247            size as u32 > dsize.0,
248            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
249            dsize.0
250        );
251
252        assert!(
253            dnum.0 * dsize.0 <= size as u32,
254            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
255            dnum.0,
256            dsize.0,
257        );
258
259        module.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
260    }
261}
262
263impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GGLWECiphertextPrepared<DR, B>
264where
265    Module<B>: VmpPrepareTmpBytes,
266{
267    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
268        module.vmp_prepare_tmp_bytes(
269            infos.dnum().into(),
270            infos.rank_in().into(),
271            (infos.rank() + 1).into(),
272            infos.size(),
273        )
274    }
275}
276
277impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWECiphertext<DR>> for GGLWECiphertextPrepared<D, B>
278where
279    Module<B>: VmpPrepare<B>,
280{
281    fn prepare(&mut self, module: &Module<B>, other: &GGLWECiphertext<DR>, scratch: &mut Scratch<B>) {
282        module.vmp_prepare(&mut self.data, &other.data, scratch);
283        self.k = other.k;
284        self.base2k = other.base2k;
285        self.dsize = other.dsize;
286    }
287}
288
289impl<D: DataRef, B: Backend> PrepareAlloc<B, GGLWECiphertextPrepared<Vec<u8>, B>> for GGLWECiphertext<D>
290where
291    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
292{
293    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GGLWECiphertextPrepared<Vec<u8>, B> {
294        let mut atk_prepared: GGLWECiphertextPrepared<Vec<u8>, B> = GGLWECiphertextPrepared::alloc(module, self);
295        atk_prepared.prepare(module, self, scratch);
296        atk_prepared
297    }
298}