1use poulpy_hal::{
  2    api::{SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, SvpPPol, SvpPPolToMut, SvpPPolToRef, ZnxInfos},
  4};
  5
  6use crate::{
  7    dist::Distribution,
  8    layouts::{
  9        Base2K, Degree, GLWEInfos, GLWESecret, GLWESecretToRef, GetDegree, GetDist, LWEInfos, Rank, TorusPrecision,
 10        prepared::SetDist,
 11    },
 12};
 13
 14pub struct GLWESecretPrepared<D: Data, B: Backend> {
 15    pub(crate) data: SvpPPol<D, B>,
 16    pub(crate) dist: Distribution,
 17}
 18
 19impl<D: DataRef, B: Backend> SetDist for GLWESecretPrepared<D, B> {
 20    fn set_dist(&mut self, dist: Distribution) {
 21        self.dist = dist
 22    }
 23}
 24
 25impl<D: Data, B: Backend> LWEInfos for GLWESecretPrepared<D, B> {
 26    fn base2k(&self) -> Base2K {
 27        Base2K(0)
 28    }
 29
 30    fn k(&self) -> TorusPrecision {
 31        TorusPrecision(0)
 32    }
 33
 34    fn n(&self) -> Degree {
 35        Degree(self.data.n() as u32)
 36    }
 37
 38    fn size(&self) -> usize {
 39        self.data.size()
 40    }
 41}
 42impl<D: Data, B: Backend> GLWEInfos for GLWESecretPrepared<D, B> {
 43    fn rank(&self) -> Rank {
 44        Rank(self.data.cols() as u32)
 45    }
 46}
 47
 48pub trait GLWESecretPreparedAlloc<B: Backend>
 49where
 50    Self: GetDegree + SvpPPolAllocBytes + SvpPPolAlloc<B>,
 51{
 52    fn glwe_secret_prepared_alloc(&self, rank: Rank) -> GLWESecretPrepared<Vec<u8>, B> {
 53        GLWESecretPrepared {
 54            data: self.svp_ppol_alloc(rank.into()),
 55            dist: Distribution::NONE,
 56        }
 57    }
 58    fn glwe_secret_prepared_alloc_from_infos<A>(&self, infos: &A) -> GLWESecretPrepared<Vec<u8>, B>
 59    where
 60        A: GLWEInfos,
 61    {
 62        assert_eq!(self.n(), infos.n());
 63        self.glwe_secret_prepared_alloc(infos.rank())
 64    }
 65
 66    fn glwe_secret_alloc_bytes(&self, rank: Rank) -> usize {
 67        self.svp_ppol_alloc_bytes(rank.into())
 68    }
 69    fn glwe_secret_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
 70    where
 71        A: GLWEInfos,
 72    {
 73        assert_eq!(self.n(), infos.n());
 74        self.glwe_secret_alloc_bytes(infos.rank())
 75    }
 76}
 77
 78impl<B: Backend> GLWESecretPreparedAlloc<B> for Module<B> where Self: GetDegree + SvpPPolAllocBytes + SvpPPolAlloc<B> {}
 79
 80impl<B: Backend> GLWESecretPrepared<Vec<u8>, B>
 81where
 82    Module<B>: GLWESecretPreparedAlloc<B>,
 83{
 84    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
 85    where
 86        A: GLWEInfos,
 87    {
 88        module.glwe_secret_prepared_alloc_from_infos(infos)
 89    }
 90
 91    pub fn alloc(module: &Module<B>, rank: Rank) -> Self {
 92        module.glwe_secret_prepared_alloc(rank)
 93    }
 94
 95    pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
 96    where
 97        A: GLWEInfos,
 98    {
 99        module.glwe_secret_alloc_bytes_from_infos(infos)
100    }
101
102    pub fn alloc_bytes(module: &Module<B>, rank: Rank) -> usize {
103        module.glwe_secret_alloc_bytes(rank)
104    }
105}
106
107impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
108    pub fn n(&self) -> Degree {
109        Degree(self.data.n() as u32)
110    }
111
112    pub fn rank(&self) -> Rank {
113        Rank(self.data.cols() as u32)
114    }
115}
116
117pub trait GLWESecretPrepare<B: Backend>
118where
119    Self: SvpPrepare<B>,
120{
121    fn glwe_secret_prepare<R, O>(&self, res: &mut R, other: &O)
122    where
123        R: GLWESecretPreparedToMut<B> + SetDist,
124        O: GLWESecretToRef + GetDist,
125    {
126        {
127            let mut res: GLWESecretPrepared<&mut [u8], _> = res.to_mut();
128            let other: GLWESecret<&[u8]> = other.to_ref();
129
130            for i in 0..res.rank().into() {
131                self.svp_prepare(&mut res.data, i, &other.data, i);
132            }
133        }
134
135        res.set_dist(other.get_dist());
136    }
137}
138
139impl<B: Backend> GLWESecretPrepare<B> for Module<B> where Self: SvpPrepare<B> {}
140
141pub trait GLWESecretPreparedToRef<B: Backend> {
142    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B>;
143}
144
145impl<D: DataRef, B: Backend> GLWESecretPreparedToRef<B> for GLWESecretPrepared<D, B> {
146    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B> {
147        GLWESecretPrepared {
148            data: self.data.to_ref(),
149            dist: self.dist,
150        }
151    }
152}
153
154pub trait GLWESecretPreparedToMut<B: Backend> {
155    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
156}
157
158impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
159    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
160        GLWESecretPrepared {
161            dist: self.dist,
162            data: self.data.to_mut(),
163        }
164    }
165}