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 GLWESecretPreparedFactory<B: Backend>
 53where
 54    Self: GetDegree + SvpPPolBytesOf + SvpPPolAlloc<B> + SvpPrepare<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    fn prepare_glwe_secret<R, O>(&self, res: &mut R, other: &O)
 82    where
 83        R: GLWESecretPreparedToMut<B> + GetDistributionMut,
 84        O: GLWESecretToRef + GetDistribution,
 85    {
 86        {
 87            let mut res: GLWESecretPrepared<&mut [u8], _> = res.to_mut();
 88            let other: GLWESecret<&[u8]> = other.to_ref();
 89
 90            for i in 0..res.rank().into() {
 91                self.svp_prepare(&mut res.data, i, &other.data, i);
 92            }
 93        }
 94
 95        *res.dist_mut() = *other.dist();
 96    }
 97}
 98
 99impl<B: Backend> GLWESecretPreparedFactory<B> for Module<B> where
100    Self: GetDegree + SvpPPolBytesOf + SvpPPolAlloc<B> + SvpPrepare<B>
101{
102}
103
104impl<B: Backend> GLWESecretPrepared<Vec<u8>, B> {
105    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
106    where
107        A: GLWEInfos,
108        M: GLWESecretPreparedFactory<B>,
109    {
110        module.alloc_glwe_secret_prepared_from_infos(infos)
111    }
112
113    pub fn alloc<M>(module: &M, rank: Rank) -> Self
114    where
115        M: GLWESecretPreparedFactory<B>,
116    {
117        module.alloc_glwe_secret_prepared(rank)
118    }
119
120    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
121    where
122        A: GLWEInfos,
123        M: GLWESecretPreparedFactory<B>,
124    {
125        module.bytes_of_glwe_secret_prepared_from_infos(infos)
126    }
127
128    pub fn bytes_of<M>(module: &M, rank: Rank) -> usize
129    where
130        M: GLWESecretPreparedFactory<B>,
131    {
132        module.bytes_of_glwe_secret_prepared(rank)
133    }
134}
135
136impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
137    pub fn n(&self) -> Degree {
138        Degree(self.data.n() as u32)
139    }
140
141    pub fn rank(&self) -> Rank {
142        Rank(self.data.cols() as u32)
143    }
144}
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: GLWESecretPreparedFactory<B>,
150        O: GLWESecretToRef + GetDistribution,
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>
170where
171    Self: GLWESecretPreparedToRef<B>,
172{
173    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
174}
175
176impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
177    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
178        GLWESecretPrepared {
179            dist: self.dist,
180            data: self.data.to_mut(),
181        }
182    }
183}