1use poulpy_hal::{
  2    api::{VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, ZnxInfos},
  4};
  5
  6use crate::{
  7    GetDistribution, GetDistributionMut,
  8    dist::Distribution,
  9    layouts::{Base2K, Degree, GLWEInfos, GLWEPublicKey, GLWEPublicKeyToRef, GetDegree, 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: DataRef, BE: Backend> GetDistribution for GLWEPublicKeyPrepared<D, BE> {
 21    fn dist(&self) -> &Distribution {
 22        &self.dist
 23    }
 24}
 25
 26impl<D: DataMut, BE: Backend> GetDistributionMut for GLWEPublicKeyPrepared<D, BE> {
 27    fn dist_mut(&mut self) -> &mut Distribution {
 28        &mut self.dist
 29    }
 30}
 31
 32impl<D: Data, B: Backend> LWEInfos for GLWEPublicKeyPrepared<D, B> {
 33    fn base2k(&self) -> Base2K {
 34        self.base2k
 35    }
 36
 37    fn k(&self) -> TorusPrecision {
 38        self.k
 39    }
 40
 41    fn size(&self) -> usize {
 42        self.data.size()
 43    }
 44
 45    fn n(&self) -> Degree {
 46        Degree(self.data.n() as u32)
 47    }
 48}
 49
 50impl<D: Data, B: Backend> GLWEInfos for GLWEPublicKeyPrepared<D, B> {
 51    fn rank(&self) -> Rank {
 52        Rank(self.data.cols() as u32 - 1)
 53    }
 54}
 55
 56pub trait GLWEPublicKeyPreparedAlloc<B: Backend>
 57where
 58    Self: GetDegree + VecZnxDftAlloc<B> + VecZnxDftBytesOf,
 59{
 60    fn alloc_glwe_public_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> GLWEPublicKeyPrepared<Vec<u8>, B> {
 61        GLWEPublicKeyPrepared {
 62            data: self.vec_znx_dft_alloc((rank + 1).into(), k.0.div_ceil(base2k.0) as usize),
 63            base2k,
 64            k,
 65            dist: Distribution::NONE,
 66        }
 67    }
 68
 69    fn alloc_glwe_public_key_prepared_from_infos<A>(&self, infos: &A) -> GLWEPublicKeyPrepared<Vec<u8>, B>
 70    where
 71        A: GLWEInfos,
 72    {
 73        self.alloc_glwe_public_key_prepared(infos.base2k(), infos.k(), infos.rank())
 74    }
 75
 76    fn bytes_of_glwe_public_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
 77        self.bytes_of_vec_znx_dft((rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
 78    }
 79
 80    fn bytes_of_glwe_public_key_prepared_from_infos<A>(&self, infos: &A) -> usize
 81    where
 82        A: GLWEInfos,
 83    {
 84        self.bytes_of_glwe_public_key_prepared(infos.base2k(), infos.k(), infos.rank())
 85    }
 86}
 87
 88impl<B: Backend> GLWEPublicKeyPreparedAlloc<B> for Module<B> where Self: VecZnxDftAlloc<B> + VecZnxDftBytesOf {}
 89
 90impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B> {
 91    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
 92    where
 93        A: GLWEInfos,
 94        M: GLWEPublicKeyPreparedAlloc<B>,
 95    {
 96        module.alloc_glwe_public_key_prepared_from_infos(infos)
 97    }
 98
 99    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
100    where
101        M: GLWEPublicKeyPreparedAlloc<B>,
102    {
103        module.alloc_glwe_public_key_prepared(base2k, k, rank)
104    }
105
106    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
107    where
108        A: GLWEInfos,
109        M: GLWEPublicKeyPreparedAlloc<B>,
110    {
111        module.bytes_of_glwe_public_key_prepared_from_infos(infos)
112    }
113
114    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
115    where
116        M: GLWEPublicKeyPreparedAlloc<B>,
117    {
118        module.bytes_of_glwe_public_key_prepared(base2k, k, rank)
119    }
120}
121
122pub trait GLWEPublicKeyPrepare<B: Backend>
123where
124    Self: GetDegree + VecZnxDftApply<B>,
125{
126    fn prepare_glwe_public_key<R, O>(&self, res: &mut R, other: &O)
127    where
128        R: GLWEPublicKeyPreparedToMut<B> + GetDistributionMut,
129        O: GLWEPublicKeyToRef + GetDistribution,
130    {
131        {
132            let mut res: GLWEPublicKeyPrepared<&mut [u8], B> = res.to_mut();
133            let other: GLWEPublicKey<&[u8]> = other.to_ref();
134
135            assert_eq!(res.n(), self.ring_degree());
136            assert_eq!(other.n(), self.ring_degree());
137            assert_eq!(res.size(), other.size());
138            assert_eq!(res.k(), other.k());
139            assert_eq!(res.base2k(), other.base2k());
140
141            for i in 0..(res.rank() + 1).into() {
142                self.vec_znx_dft_apply(1, 0, &mut res.data, i, &other.data, i);
143            }
144        }
145
146        *res.dist_mut() = *other.dist();
147    }
148}
149
150impl<B: Backend> GLWEPublicKeyPrepare<B> for Module<B> where Self: GetDegree + VecZnxDftApply<B> {}
151
152impl<D: DataMut, B: Backend> GLWEPublicKeyPrepared<D, B> {
153    pub fn prepare<O, M>(&mut self, module: &M, other: &O)
154    where
155        O: GLWEPublicKeyToRef + GetDistribution,
156        M: GLWEPublicKeyPrepare<B>,
157    {
158        module.prepare_glwe_public_key(self, other);
159    }
160}
161
162pub trait GLWEPublicKeyPreparedToMut<B: Backend> {
163    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B>;
164}
165
166impl<D: DataMut, B: Backend> GLWEPublicKeyPreparedToMut<B> for GLWEPublicKeyPrepared<D, B> {
167    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B> {
168        GLWEPublicKeyPrepared {
169            dist: self.dist,
170            k: self.k,
171            base2k: self.base2k,
172            data: self.data.to_mut(),
173        }
174    }
175}
176
177pub trait GLWEPublicKeyPreparedToRef<B: Backend> {
178    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B>;
179}
180
181impl<D: DataRef, B: Backend> GLWEPublicKeyPreparedToRef<B> for GLWEPublicKeyPrepared<D, B> {
182    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B> {
183        GLWEPublicKeyPrepared {
184            data: self.data.to_ref(),
185            dist: self.dist,
186            k: self.k,
187            base2k: self.base2k,
188        }
189    }
190}