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
 98impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B> {
 99    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
100    where
101        A: GLWEInfos,
102        M: GLWEPublicKeyPreparedFactory<B>,
103    {
104        module.alloc_glwe_public_key_prepared_from_infos(infos)
105    }
106
107    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
108    where
109        M: GLWEPublicKeyPreparedFactory<B>,
110    {
111        module.alloc_glwe_public_key_prepared(base2k, k, rank)
112    }
113
114    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
115    where
116        A: GLWEInfos,
117        M: GLWEPublicKeyPreparedFactory<B>,
118    {
119        module.bytes_of_glwe_public_key_prepared_from_infos(infos)
120    }
121
122    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
123    where
124        M: GLWEPublicKeyPreparedFactory<B>,
125    {
126        module.bytes_of_glwe_public_key_prepared(base2k, k, rank)
127    }
128}
129
130impl<D: DataMut, B: Backend> GLWEPublicKeyPrepared<D, B> {
131    pub fn prepare<O, M>(&mut self, module: &M, other: &O)
132    where
133        O: GLWEToRef + GetDistribution,
134        M: GLWEPublicKeyPreparedFactory<B>,
135    {
136        module.prepare_glwe_public_key(self, other);
137    }
138}
139
140impl<D: DataMut, B: Backend> GLWEPreparedToMut<B> for GLWEPublicKeyPrepared<D, B>
141where
142    GLWEPrepared<D, B>: GLWEPreparedToMut<B>,
143{
144    fn to_mut(&mut self) -> GLWEPrepared<&mut [u8], B> {
145        self.key.to_mut()
146    }
147}
148
149impl<D: DataRef, B: Backend> GLWEPreparedToRef<B> for GLWEPublicKeyPrepared<D, B>
150where
151    GLWEPrepared<D, B>: GLWEPreparedToRef<B>,
152{
153    fn to_ref(&self) -> GLWEPrepared<&[u8], B> {
154        self.key.to_ref()
155    }
156}