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