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