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            for i in 0..res.rank().into() {
 90                self.svp_prepare(&mut res.data, i, &other.data, i);
 91            }
 92        }
 93
 94        *res.dist_mut() = *other.dist();
 95    }
 96}
 97
 98impl<B: Backend> GLWESecretPreparedFactory<B> for Module<B> where
 99    Self: GetDegree + SvpPPolBytesOf + SvpPPolAlloc<B> + SvpPrepare<B>
100{
101}
102
103impl<B: Backend> GLWESecretPrepared<Vec<u8>, B> {
104    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
105    where
106        A: GLWEInfos,
107        M: GLWESecretPreparedFactory<B>,
108    {
109        module.alloc_glwe_secret_prepared_from_infos(infos)
110    }
111
112    pub fn alloc<M>(module: &M, rank: Rank) -> Self
113    where
114        M: GLWESecretPreparedFactory<B>,
115    {
116        module.alloc_glwe_secret_prepared(rank)
117    }
118
119    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
120    where
121        A: GLWEInfos,
122        M: GLWESecretPreparedFactory<B>,
123    {
124        module.bytes_of_glwe_secret_prepared_from_infos(infos)
125    }
126
127    pub fn bytes_of<M>(module: &M, rank: Rank) -> usize
128    where
129        M: GLWESecretPreparedFactory<B>,
130    {
131        module.bytes_of_glwe_secret_prepared(rank)
132    }
133}
134
135impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
136    pub fn n(&self) -> Degree {
137        Degree(self.data.n() as u32)
138    }
139
140    pub fn rank(&self) -> Rank {
141        Rank(self.data.cols() as u32)
142    }
143}
144
145impl<D: DataMut, B: Backend> GLWESecretPrepared<D, B> {
146    pub fn prepare<M, O>(&mut self, module: &M, other: &O)
147    where
148        M: GLWESecretPreparedFactory<B>,
149        O: GLWESecretToRef + GetDistribution,
150    {
151        module.prepare_glwe_secret(self, other);
152    }
153}
154
155pub trait GLWESecretPreparedToRef<B: Backend> {
156    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B>;
157}
158
159impl<D: DataRef, B: Backend> GLWESecretPreparedToRef<B> for GLWESecretPrepared<D, B> {
160    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B> {
161        GLWESecretPrepared {
162            data: self.data.to_ref(),
163            dist: self.dist,
164        }
165    }
166}
167
168pub trait GLWESecretPreparedToMut<B: Backend>
169where
170    Self: GLWESecretPreparedToRef<B>,
171{
172    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
173}
174
175impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
176    fn to_mut(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
177        GLWESecretPrepared {
178            dist: self.dist,
179            data: self.data.to_mut(),
180        }
181    }
182}