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