1use poulpy_hal::{
  2    api::{VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftApply},
  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(crate) 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> + VecZnxDftAllocBytes,
 56{
 57    fn glwe_public_key_prepared_alloc(&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 glwe_public_key_alloc_from_infos<A>(&self, infos: &A) -> GLWEPublicKeyPrepared<Vec<u8>, B>
 67    where
 68        A: GLWEInfos,
 69    {
 70        self.glwe_public_key_prepared_alloc(infos.base2k(), infos.k(), infos.rank())
 71    }
 72
 73    fn glwe_public_key_prepared_alloc_bytes(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
 74        self.vec_znx_dft_alloc_bytes((rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
 75    }
 76
 77    fn glwe_public_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
 78    where
 79        A: GLWEInfos,
 80    {
 81        self.glwe_public_key_prepared_alloc_bytes(infos.base2k(), infos.k(), infos.rank())
 82    }
 83}
 84
 85impl<B: Backend> GLWEPublicKeyPreparedAlloc<B> for Module<B> where Self: VecZnxDftAlloc<B> + VecZnxDftAllocBytes {}
 86
 87impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B>
 88where
 89    Module<B>: GLWEPublicKeyPreparedAlloc<B>,
 90{
 91    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
 92    where
 93        A: GLWEInfos,
 94    {
 95        module.glwe_public_key_alloc_from_infos(infos)
 96    }
 97
 98    pub fn alloc(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 99        module.glwe_public_key_prepared_alloc(base2k, k, rank)
100    }
101
102    pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
103    where
104        A: GLWEInfos,
105    {
106        module.glwe_public_key_prepared_alloc_bytes_from_infos(infos)
107    }
108
109    pub fn alloc_bytes(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
110        module.glwe_public_key_prepared_alloc_bytes(base2k, k, rank)
111    }
112}
113
114pub trait GLWEPublicKeyPrepare<B: Backend>
115where
116    Self: GetDegree + VecZnxDftApply<B>,
117{
118    fn glwe_public_key_prepare<R, O>(&self, res: &mut R, other: &O)
119    where
120        R: GLWEPublicKeyPreparedToMut<B> + SetDist,
121        O: GLWEPublicKeyToRef + GetDist,
122    {
123        {
124            let mut res: GLWEPublicKeyPrepared<&mut [u8], B> = res.to_mut();
125            let other: GLWEPublicKey<&[u8]> = other.to_ref();
126
127            assert_eq!(res.n(), self.n());
128            assert_eq!(other.n(), self.n());
129            assert_eq!(res.size(), other.size());
130            assert_eq!(res.k(), other.k());
131            assert_eq!(res.base2k(), other.base2k());
132
133            for i in 0..(res.rank() + 1).into() {
134                self.vec_znx_dft_apply(1, 0, &mut res.data, i, &other.data, i);
135            }
136        }
137
138        res.set_dist(other.get_dist());
139    }
140}
141
142impl<B: Backend> GLWEPublicKeyPrepare<B> for Module<B> where Self: GetDegree + VecZnxDftApply<B> {}
143
144impl<D: DataMut, B: Backend> GLWEPublicKeyPrepared<D, B>
145where
146    Module<B>: GLWEPublicKeyPrepare<B>,
147{
148    pub fn prepare<O>(&mut self, module: &Module<B>, other: &O)
149    where
150        O: GLWEPublicKeyToRef + GetDist,
151    {
152        module.glwe_public_key_prepare(self, other);
153    }
154}
155
156pub trait GLWEPublicKeyPreparedToMut<B: Backend> {
157    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B>;
158}
159
160impl<D: DataMut, B: Backend> GLWEPublicKeyPreparedToMut<B> for GLWEPublicKeyPrepared<D, B> {
161    fn to_mut(&mut self) -> GLWEPublicKeyPrepared<&mut [u8], B> {
162        GLWEPublicKeyPrepared {
163            dist: self.dist,
164            k: self.k,
165            base2k: self.base2k,
166            data: self.data.to_mut(),
167        }
168    }
169}
170
171pub trait GLWEPublicKeyPreparedToRef<B: Backend> {
172    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B>;
173}
174
175impl<D: DataRef, B: Backend> GLWEPublicKeyPreparedToRef<B> for GLWEPublicKeyPrepared<D, B> {
176    fn to_ref(&self) -> GLWEPublicKeyPrepared<&[u8], B> {
177        GLWEPublicKeyPrepared {
178            data: self.data.to_ref(),
179            dist: self.dist,
180            k: self.k,
181            base2k: self.base2k,
182        }
183    }
184}