1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, ZnxInfos},
  4    oep::VmpPMatAllocBytesImpl,
  5};
  6
  7use crate::layouts::{
  8    Base2K, BuildError, Degree, Digits, GGLWECiphertext, GGLWELayoutInfos, GLWEInfos, LWEInfos, Rank, Rows, TorusPrecision,
  9    prepared::{Prepare, PrepareAlloc},
 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) digits: Digits,
 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> GGLWELayoutInfos 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 digits(&self) -> Digits {
 54        self.digits
 55    }
 56
 57    fn rows(&self) -> Rows {
 58        Rows(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    digits: Option<Digits>,
 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            digits: 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: GGLWELayoutInfos,
 86        B: VmpPMatAllocBytesImpl<B>,
 87    {
 88        self.data = Some(VmpPMat::alloc(
 89            infos.n().into(),
 90            infos.rows().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.digits = Some(infos.digits());
 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 digits(mut self, digits: Digits) -> Self {
121        self.digits = Some(digits);
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 digits: Digits = self.digits.ok_or(BuildError::MissingDigits)?;
130
131        if base2k == 0_u32 {
132            return Err(BuildError::ZeroBase2K);
133        }
134
135        if digits == 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            digits,
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: GGLWELayoutInfos,
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.rows(),
176            infos.digits(),
177            infos.rank_in(),
178            infos.rank_out(),
179        )
180    }
181
182    pub fn alloc_with(
183        module: &Module<B>,
184        base2k: Base2K,
185        k: TorusPrecision,
186        rows: Rows,
187        digits: Digits,
188        rank_in: Rank,
189        rank_out: Rank,
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 > digits.0,
197            "invalid gglwe: ceil(k/base2k): {size} <= digits: {}",
198            digits.0
199        );
200
201        assert!(
202            rows.0 * digits.0 <= size as u32,
203            "invalid gglwe: rows: {} * digits:{} > ceil(k/base2k): {size}",
204            rows.0,
205            digits.0,
206        );
207
208        Self {
209            data: module.vmp_pmat_alloc(rows.into(), rank_in.into(), (rank_out + 1).into(), size),
210            k,
211            base2k,
212            digits,
213        }
214    }
215
216    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
217    where
218        A: GGLWELayoutInfos,
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.rows(),
227            infos.digits(),
228            infos.rank_in(),
229            infos.rank_out(),
230        )
231    }
232
233    pub fn alloc_bytes_with(
234        module: &Module<B>,
235        base2k: Base2K,
236        k: TorusPrecision,
237        rows: Rows,
238        digits: Digits,
239        rank_in: Rank,
240        rank_out: Rank,
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 > digits.0,
248            "invalid gglwe: ceil(k/base2k): {size} <= digits: {}",
249            digits.0
250        );
251
252        assert!(
253            rows.0 * digits.0 <= size as u32,
254            "invalid gglwe: rows: {} * digits:{} > ceil(k/base2k): {size}",
255            rows.0,
256            digits.0,
257        );
258
259        module.vmp_pmat_alloc_bytes(rows.into(), rank_in.into(), (rank_out + 1).into(), size)
260    }
261}
262
263impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWECiphertext<DR>> for GGLWECiphertextPrepared<D, B>
264where
265    Module<B>: VmpPrepare<B>,
266{
267    fn prepare(&mut self, module: &Module<B>, other: &GGLWECiphertext<DR>, scratch: &mut Scratch<B>) {
268        module.vmp_prepare(&mut self.data, &other.data, scratch);
269        self.k = other.k;
270        self.base2k = other.base2k;
271        self.digits = other.digits;
272    }
273}
274
275impl<D: DataRef, B: Backend> PrepareAlloc<B, GGLWECiphertextPrepared<Vec<u8>, B>> for GGLWECiphertext<D>
276where
277    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
278{
279    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GGLWECiphertextPrepared<Vec<u8>, B> {
280        let mut atk_prepared: GGLWECiphertextPrepared<Vec<u8>, B> = GGLWECiphertextPrepared::alloc(module, self);
281        atk_prepared.prepare(module, self, scratch);
282        atk_prepared
283    }
284}