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