1use poulpy_hal::{
  2    api::{SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, SvpPPol, SvpPPolToMut, SvpPPolToRef, ZnxInfos},
  4};
  5
  6use crate::{
  7    dist::Distribution,
  8    layouts::{
  9        Base2K, GLWEInfos, GLWESecret, GLWESecretToRef, GetDist, GetRingDegree, LWEInfos, Rank, RingDegree, 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) -> RingDegree {
 35        RingDegree(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: GetRingDegree + SvpPPolBytesOf + SvpPPolAlloc<B>,
 51{
 52    fn alloc_glwe_secret_prepared(&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 alloc_glwe_secret_prepared_from_infos<A>(&self, infos: &A) -> GLWESecretPrepared<Vec<u8>, B>
 59    where
 60        A: GLWEInfos,
 61    {
 62        assert_eq!(self.ring_degree(), infos.n());
 63        self.alloc_glwe_secret_prepared(infos.rank())
 64    }
 65
 66    fn bytes_of_glwe_secret_prepared(&self, rank: Rank) -> usize {
 67        self.bytes_of_svp_ppol(rank.into())
 68    }
 69    fn bytes_of_glwe_secret_prepared_from_infos<A>(&self, infos: &A) -> usize
 70    where
 71        A: GLWEInfos,
 72    {
 73        assert_eq!(self.ring_degree(), infos.n());
 74        self.bytes_of_glwe_secret_prepared(infos.rank())
 75    }
 76}
 77
 78impl<B: Backend> GLWESecretPreparedAlloc<B> for Module<B> where Self: GetRingDegree + SvpPPolBytesOf + SvpPPolAlloc<B> {}
 79
 80impl<B: Backend> GLWESecretPrepared<Vec<u8>, B> {
 81    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
 82    where
 83        A: GLWEInfos,
 84        M: GLWESecretPreparedAlloc<B>,
 85    {
 86        module.alloc_glwe_secret_prepared_from_infos(infos)
 87    }
 88
 89    pub fn alloc<M>(module: &M, rank: Rank) -> Self
 90    where
 91        M: GLWESecretPreparedAlloc<B>,
 92    {
 93        module.alloc_glwe_secret_prepared(rank)
 94    }
 95
 96    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
 97    where
 98        A: GLWEInfos,
 99        M: GLWESecretPreparedAlloc<B>,
100    {
101        module.bytes_of_glwe_secret_prepared_from_infos(infos)
102    }
103
104    pub fn bytes_of<M>(module: &M, rank: Rank) -> usize
105    where
106        M: GLWESecretPreparedAlloc<B>,
107    {
108        module.bytes_of_glwe_secret_prepared(rank)
109    }
110}
111
112impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
113    pub fn n(&self) -> RingDegree {
114        RingDegree(self.data.n() as u32)
115    }
116
117    pub fn rank(&self) -> Rank {
118        Rank(self.data.cols() as u32)
119    }
120}
121
122pub trait GLWESecretPrepare<B: Backend>
123where
124    Self: SvpPrepare<B>,
125{
126    fn prepare_glwe_secret<R, O>(&self, res: &mut R, other: &O)
127    where
128        R: GLWESecretPreparedToMut<B> + SetDist,
129        O: GLWESecretToRef + GetDist,
130    {
131        {
132            let mut res: GLWESecretPrepared<&mut [u8], _> = res.to_mut();
133            let other: GLWESecret<&[u8]> = other.to_ref();
134
135            for i in 0..res.rank().into() {
136                self.svp_prepare(&mut res.data, i, &other.data, i);
137            }
138        }
139
140        res.set_dist(other.get_dist());
141    }
142}
143
144impl<B: Backend> GLWESecretPrepare<B> for Module<B> where Self: SvpPrepare<B> {}
145
146impl<D: DataMut, B: Backend> GLWESecretPrepared<D, B> {
147    pub fn prepare<M, O>(&mut self, module: &M, other: &O)
148    where
149        M: GLWESecretPrepare<B>,
150        O: GLWESecretToRef + GetDist,
151    {
152        module.prepare_glwe_secret(self, other);
153    }
154}
155
156pub trait GLWESecretPreparedToRef<B: Backend> {
157    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B>;
158}
159
160impl<D: DataRef, B: Backend> GLWESecretPreparedToRef<B> for GLWESecretPrepared<D, B> {
161    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B> {
162        GLWESecretPrepared {
163            data: self.data.to_ref(),
164            dist: self.dist,
165        }
166    }
167}
168
169pub trait GLWESecretPreparedToMut<B: Backend> {
170    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
171}
172
173impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
174    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
175        GLWESecretPrepared {
176            dist: self.dist,
177            data: self.data.to_mut(),
178        }
179    }
180}