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, 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
 20impl<D: Data, B: Backend> LWEInfos for GLWEPublicKeyPrepared<D, B> {
 21    fn base2k(&self) -> Base2K {
 22        self.base2k
 23    }
 24
 25    fn k(&self) -> TorusPrecision {
 26        self.k
 27    }
 28
 29    fn size(&self) -> usize {
 30        self.data.size()
 31    }
 32
 33    fn n(&self) -> Degree {
 34        Degree(self.data.n() as u32)
 35    }
 36}
 37
 38impl<D: Data, B: Backend> GLWEInfos for GLWEPublicKeyPrepared<D, B> {
 39    fn rank(&self) -> Rank {
 40        Rank(self.data.cols() as u32 - 1)
 41    }
 42}
 43
 44pub struct GLWEPublicKeyPreparedBuilder<D: Data, B: Backend> {
 45    data: Option<VecZnxDft<D, B>>,
 46    base2k: Option<Base2K>,
 47    k: Option<TorusPrecision>,
 48}
 49
 50impl<D: Data, B: Backend> GLWEPublicKeyPrepared<D, B> {
 51    #[inline]
 52    pub fn builder() -> GLWEPublicKeyPreparedBuilder<D, B> {
 53        GLWEPublicKeyPreparedBuilder {
 54            data: None,
 55            base2k: None,
 56            k: None,
 57        }
 58    }
 59}
 60
 61impl<B: Backend> GLWEPublicKeyPreparedBuilder<Vec<u8>, B> {
 62    #[inline]
 63    pub fn layout<A>(mut self, layout: &A) -> Self
 64    where
 65        A: GLWEInfos,
 66        B: VecZnxDftAllocBytesImpl<B>,
 67    {
 68        self.data = Some(VecZnxDft::alloc(
 69            layout.n().into(),
 70            (layout.rank() + 1).into(),
 71            layout.size(),
 72        ));
 73        self.base2k = Some(layout.base2k());
 74        self.k = Some(layout.k());
 75        self
 76    }
 77}
 78
 79impl<D: Data, B: Backend> GLWEPublicKeyPreparedBuilder<D, B> {
 80    #[inline]
 81    pub fn data(mut self, data: VecZnxDft<D, B>) -> Self {
 82        self.data = Some(data);
 83        self
 84    }
 85    #[inline]
 86    pub fn base2k(mut self, base2k: Base2K) -> Self {
 87        self.base2k = Some(base2k);
 88        self
 89    }
 90    #[inline]
 91    pub fn k(mut self, k: TorusPrecision) -> Self {
 92        self.k = Some(k);
 93        self
 94    }
 95
 96    pub fn build(self) -> Result<GLWEPublicKeyPrepared<D, B>, BuildError> {
 97        let data: VecZnxDft<D, B> = self.data.ok_or(BuildError::MissingData)?;
 98        let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
 99        let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
100
101        if base2k == 0_u32 {
102            return Err(BuildError::ZeroBase2K);
103        }
104
105        if k == 0_u32 {
106            return Err(BuildError::ZeroTorusPrecision);
107        }
108
109        if data.n() == 0 {
110            return Err(BuildError::ZeroDegree);
111        }
112
113        if data.cols() == 0 {
114            return Err(BuildError::ZeroCols);
115        }
116
117        if data.size() == 0 {
118            return Err(BuildError::ZeroLimbs);
119        }
120
121        Ok(GLWEPublicKeyPrepared {
122            data,
123            base2k,
124            k,
125            dist: Distribution::NONE,
126        })
127    }
128}
129
130impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B> {
131    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
132    where
133        A: GLWEInfos,
134        Module<B>: VecZnxDftAlloc<B>,
135    {
136        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
137        Self::alloc_with(module, infos.base2k(), infos.k(), infos.rank())
138    }
139
140    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
141    where
142        Module<B>: VecZnxDftAlloc<B>,
143    {
144        Self {
145            data: module.vec_znx_dft_alloc((rank + 1).into(), k.0.div_ceil(base2k.0) as usize),
146            base2k,
147            k,
148            dist: Distribution::NONE,
149        }
150    }
151
152    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
153    where
154        A: GLWEInfos,
155        Module<B>: VecZnxDftAllocBytes,
156    {
157        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
158        Self::alloc_bytes_with(module, infos.base2k(), infos.k(), infos.rank())
159    }
160
161    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
162    where
163        Module<B>: VecZnxDftAllocBytes,
164    {
165        module.vec_znx_dft_alloc_bytes((rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
166    }
167}
168
169impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEPublicKeyPrepared<Vec<u8>, B>> for GLWEPublicKey<D>
170where
171    Module<B>: VecZnxDftAlloc<B> + VecZnxDftApply<B>,
172{
173    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEPublicKeyPrepared<Vec<u8>, B> {
174        let mut pk_prepared: GLWEPublicKeyPrepared<Vec<u8>, B> = GLWEPublicKeyPrepared::alloc(module, self);
175        pk_prepared.prepare(module, self, scratch);
176        pk_prepared
177    }
178}
179
180impl<B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWEPublicKeyPrepared<Vec<u8>, B> {
181    fn prepare_scratch_space(_module: &Module<B>, _infos: &A) -> usize {
182        0
183    }
184}
185
186impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEPublicKey<DR>> for GLWEPublicKeyPrepared<DM, B>
187where
188    Module<B>: VecZnxDftApply<B>,
189{
190    fn prepare(&mut self, module: &Module<B>, other: &GLWEPublicKey<DR>, _scratch: &mut Scratch<B>) {
191        #[cfg(debug_assertions)]
192        {
193            assert_eq!(self.n(), other.n());
194            assert_eq!(self.size(), other.size());
195        }
196
197        (0..(self.rank() + 1).into()).for_each(|i| {
198            module.vec_znx_dft_apply(1, 0, &mut self.data, i, &other.data, i);
199        });
200        self.k = other.k();
201        self.base2k = other.base2k();
202        self.dist = other.dist;
203    }
204}
205
206pub trait GLWEPublicKeyPreparedToMut<B: Backend> {
207    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B>;
208}
209
210impl<D: DataMut, B: Backend> GLWEPublicKeyPreparedToMut<B> for GLWEPublicKeyPrepared<D, B> {
211    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B> {
212        GLWEPublicKeyPrepared {
213            dist: self.dist,
214            k: self.k,
215            base2k: self.base2k,
216            data: self.data.to_mut(),
217        }
218    }
219}
220
221pub trait GLWEPublicKeyPreparedToRef<B: Backend> {
222    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B>;
223}
224
225impl<D: DataRef, B: Backend> GLWEPublicKeyPreparedToRef<B> for GLWEPublicKeyPrepared<D, B> {
226    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B> {
227        GLWEPublicKeyPrepared {
228            data: self.data.to_ref(),
229            dist: self.dist,
230            k: self.k,
231            base2k: self.base2k,
232        }
233    }
234}