1use poulpy_hal::{
  2    api::{VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf},
  3    layouts::{Backend, Data, DataMut, DataRef, Module},
  4};
  5
  6use crate::{
  7    GetDistribution, GetDistributionMut,
  8    dist::Distribution,
  9    layouts::{
 10        Base2K, Degree, GLWEInfos, GLWEPrepared, GLWEPreparedFactory, GLWEPreparedToMut, GLWEPreparedToRef, GLWEToRef, GetDegree,
 11        LWEInfos, Rank, TorusPrecision,
 12    },
 13};
 14
 15#[derive(PartialEq, Eq)]
 16pub struct GLWEPublicKeyPrepared<D: Data, B: Backend> {
 17    pub(crate) key: GLWEPrepared<D, B>,
 18    pub(crate) dist: Distribution,
 19}
 20
 21impl<D: DataRef, BE: Backend> GetDistribution for GLWEPublicKeyPrepared<D, BE> {
 22    fn dist(&self) -> &Distribution {
 23        &self.dist
 24    }
 25}
 26
 27impl<D: DataMut, BE: Backend> GetDistributionMut for GLWEPublicKeyPrepared<D, BE> {
 28    fn dist_mut(&mut self) -> &mut Distribution {
 29        &mut self.dist
 30    }
 31}
 32
 33impl<D: Data, B: Backend> LWEInfos for GLWEPublicKeyPrepared<D, B> {
 34    fn base2k(&self) -> Base2K {
 35        self.key.base2k()
 36    }
 37
 38    fn k(&self) -> TorusPrecision {
 39        self.key.k()
 40    }
 41
 42    fn size(&self) -> usize {
 43        self.key.size()
 44    }
 45
 46    fn n(&self) -> Degree {
 47        self.key.n()
 48    }
 49}
 50
 51impl<D: Data, B: Backend> GLWEInfos for GLWEPublicKeyPrepared<D, B> {
 52    fn rank(&self) -> Rank {
 53        self.key.rank()
 54    }
 55}
 56
 57pub trait GLWEPublicKeyPreparedFactory<B: Backend>
 58where
 59    Self: GetDegree + GLWEPreparedFactory<B>,
 60{
 61    fn alloc_glwe_public_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> GLWEPublicKeyPrepared<Vec<u8>, B> {
 62        GLWEPublicKeyPrepared {
 63            key: self.alloc_glwe_prepared(base2k, k, rank),
 64            dist: Distribution::NONE,
 65        }
 66    }
 67
 68    fn alloc_glwe_public_key_prepared_from_infos<A>(&self, infos: &A) -> GLWEPublicKeyPrepared<Vec<u8>, B>
 69    where
 70        A: GLWEInfos,
 71    {
 72        self.alloc_glwe_public_key_prepared(infos.base2k(), infos.k(), infos.rank())
 73    }
 74
 75    fn bytes_of_glwe_public_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
 76        self.bytes_of_glwe_prepared(base2k, k, rank)
 77    }
 78
 79    fn bytes_of_glwe_public_key_prepared_from_infos<A>(&self, infos: &A) -> usize
 80    where
 81        A: GLWEInfos,
 82    {
 83        self.bytes_of_glwe_public_key_prepared(infos.base2k(), infos.k(), infos.rank())
 84    }
 85
 86    fn prepare_glwe_public_key<R, O>(&self, res: &mut R, other: &O)
 87    where
 88        R: GLWEPreparedToMut<B> + GetDistributionMut,
 89        O: GLWEToRef + GetDistribution,
 90    {
 91        self.prepare_glwe(res, other);
 92        *res.dist_mut() = *other.dist();
 93    }
 94}
 95
 96impl<B: Backend> GLWEPublicKeyPreparedFactory<B> for Module<B> where Self: VecZnxDftAlloc<B> + VecZnxDftBytesOf + VecZnxDftApply<B>
 97{}
 98
 99impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B> {
100    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
101    where
102        A: GLWEInfos,
103        M: GLWEPublicKeyPreparedFactory<B>,
104    {
105        module.alloc_glwe_public_key_prepared_from_infos(infos)
106    }
107
108    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
109    where
110        M: GLWEPublicKeyPreparedFactory<B>,
111    {
112        module.alloc_glwe_public_key_prepared(base2k, k, rank)
113    }
114
115    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
116    where
117        A: GLWEInfos,
118        M: GLWEPublicKeyPreparedFactory<B>,
119    {
120        module.bytes_of_glwe_public_key_prepared_from_infos(infos)
121    }
122
123    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
124    where
125        M: GLWEPublicKeyPreparedFactory<B>,
126    {
127        module.bytes_of_glwe_public_key_prepared(base2k, k, rank)
128    }
129}
130
131impl<D: DataMut, B: Backend> GLWEPublicKeyPrepared<D, B> {
132    pub fn prepare<O, M>(&mut self, module: &M, other: &O)
133    where
134        O: GLWEToRef + GetDistribution,
135        M: GLWEPublicKeyPreparedFactory<B>,
136    {
137        module.prepare_glwe_public_key(self, other);
138    }
139}
140
141impl<D: DataMut, B: Backend> GLWEPreparedToMut<B> for GLWEPublicKeyPrepared<D, B>
142where
143    GLWEPrepared<D, B>: GLWEPreparedToMut<B>,
144{
145    fn to_mut(&mut self) -> GLWEPrepared<&mut [u8], B> {
146        self.key.to_mut()
147    }
148}
149
150impl<D: DataRef, B: Backend> GLWEPreparedToRef<B> for GLWEPublicKeyPrepared<D, B>
151where
152    GLWEPrepared<D, B>: GLWEPreparedToRef<B>,
153{
154    fn to_ref(&self) -> GLWEPrepared<&[u8], B> {
155        self.key.to_ref()
156    }
157}