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