1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare, VmpPrepareTmpBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, VmpPMatToRef, ZnxInfos},
  4    oep::VmpPMatAllocBytesImpl,
  5};
  6
  7use crate::layouts::{
  8    Base2K, BuildError, Degree, Dnum, Dsize, GGSWCiphertext, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  9    prepared::{Prepare, PrepareAlloc, PrepareScratchSpace},
 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) dsize: Dsize,
 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 dsize(&self) -> Dsize {
 46        self.dsize
 47    }
 48
 49    fn dnum(&self) -> Dnum {
 50        Dnum(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    dsize: Option<Dsize>,
 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            dsize: 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.dsize().0,
 82            "invalid ggsw: ceil(k/base2k): {} <= dsize: {}",
 83            infos.size(),
 84            infos.dsize()
 85        );
 86
 87        assert!(
 88            infos.dnum().0 * infos.dsize().0 <= infos.size() as u32,
 89            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {}",
 90            infos.dnum(),
 91            infos.dsize(),
 92            infos.size(),
 93        );
 94
 95        self.data = Some(VmpPMat::alloc(
 96            infos.n().into(),
 97            infos.dnum().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.dsize = Some(infos.dsize());
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 dsize(mut self, dsize: Dsize) -> Self {
128        self.dsize = Some(dsize);
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 dsize: Dsize = self.dsize.ok_or(BuildError::MissingDigits)?;
137
138        if base2k == 0_u32 {
139            return Err(BuildError::ZeroBase2K);
140        }
141
142        if dsize == 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            dsize,
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.dnum(),
182            infos.dsize(),
183            infos.rank(),
184        )
185    }
186
187    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, 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 > dsize.0,
194            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
195            dsize.0
196        );
197
198        assert!(
199            dnum.0 * dsize.0 <= size as u32,
200            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
201            dnum.0,
202            dsize.0,
203        );
204
205        Self {
206            data: module.vmp_pmat_alloc(
207                dnum.into(),
208                (rank + 1).into(),
209                (rank + 1).into(),
210                k.0.div_ceil(base2k.0) as usize,
211            ),
212            k,
213            base2k,
214            dsize,
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.dnum(),
228            infos.dsize(),
229            infos.rank(),
230        )
231    }
232
233    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> usize
234    where
235        Module<B>: VmpPMatAllocBytes,
236    {
237        let size: usize = k.0.div_ceil(base2k.0) as usize;
238        debug_assert!(
239            size as u32 > dsize.0,
240            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
241            dsize.0
242        );
243
244        assert!(
245            dnum.0 * dsize.0 <= size as u32,
246            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
247            dnum.0,
248            dsize.0,
249        );
250
251        module.vmp_pmat_alloc_bytes(dnum.into(), (rank + 1).into(), (rank + 1).into(), size)
252    }
253}
254
255impl<D: DataRef, B: Backend> GGSWCiphertextPrepared<D, B> {
256    pub fn data(&self) -> &VmpPMat<D, B> {
257        &self.data
258    }
259}
260
261impl<DR: DataRef, B: Backend, A: GGSWInfos> PrepareScratchSpace<B, A> for GGSWCiphertextPrepared<DR, B>
262where
263    Module<B>: VmpPrepareTmpBytes,
264{
265    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
266        module.vmp_prepare_tmp_bytes(
267            infos.dnum().into(),
268            (infos.rank() + 1).into(),
269            (infos.rank() + 1).into(),
270            infos.size(),
271        )
272    }
273}
274
275impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGSWCiphertext<DR>> for GGSWCiphertextPrepared<D, B>
276where
277    Module<B>: VmpPrepare<B>,
278{
279    fn prepare(&mut self, module: &Module<B>, other: &GGSWCiphertext<DR>, scratch: &mut Scratch<B>) {
280        module.vmp_prepare(&mut self.data, &other.data, scratch);
281        self.k = other.k;
282        self.base2k = other.base2k;
283        self.dsize = other.dsize;
284    }
285}
286
287impl<D: DataRef, B: Backend> PrepareAlloc<B, GGSWCiphertextPrepared<Vec<u8>, B>> for GGSWCiphertext<D>
288where
289    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
290{
291    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GGSWCiphertextPrepared<Vec<u8>, B> {
292        let mut ggsw_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = GGSWCiphertextPrepared::alloc(module, self);
293        ggsw_prepared.prepare(module, self, scratch);
294        ggsw_prepared
295    }
296}
297
298pub trait GGSWCiphertextPreparedToRef<B: Backend> {
299    fn to_ref(&self) -> GGSWCiphertextPrepared<&[u8], B>;
300}
301
302impl<D: DataRef, B: Backend> GGSWCiphertextPreparedToRef<B> for GGSWCiphertextPrepared<D, B> {
303    fn to_ref(&self) -> GGSWCiphertextPrepared<&[u8], B> {
304        GGSWCiphertextPrepared::builder()
305            .base2k(self.base2k())
306            .dsize(self.dsize())
307            .k(self.k())
308            .data(self.data.to_ref())
309            .build()
310            .unwrap()
311    }
312}