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