1use poulpy_hal::{
  2    layouts::{
  3        Backend, Data, DataMut, DataRef, Module, ReaderFrom, ScalarZnx, ScalarZnxToMut, ScalarZnxToRef, WriterTo, ZnxInfos,
  4        ZnxZero,
  5    },
  6    source::Source,
  7};
  8
  9use crate::{
 10    GetDistribution,
 11    dist::Distribution,
 12    layouts::{Base2K, Degree, GLWEInfos, GetDegree, LWEInfos, Rank, TorusPrecision},
 13};
 14
 15#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 16pub struct GLWESecretLayout {
 17    pub n: Degree,
 18    pub rank: Rank,
 19}
 20
 21impl LWEInfos for GLWESecretLayout {
 22    fn base2k(&self) -> Base2K {
 23        Base2K(0)
 24    }
 25
 26    fn k(&self) -> TorusPrecision {
 27        TorusPrecision(0)
 28    }
 29
 30    fn n(&self) -> Degree {
 31        self.n
 32    }
 33
 34    fn size(&self) -> usize {
 35        1
 36    }
 37}
 38impl GLWEInfos for GLWESecretLayout {
 39    fn rank(&self) -> Rank {
 40        self.rank
 41    }
 42}
 43
 44#[derive(PartialEq, Eq, Clone)]
 45pub struct GLWESecret<D: Data> {
 46    pub(crate) data: ScalarZnx<D>,
 47    pub(crate) dist: Distribution,
 48}
 49
 50impl<D: Data> LWEInfos for GLWESecret<D> {
 51    fn base2k(&self) -> Base2K {
 52        Base2K(0)
 53    }
 54
 55    fn k(&self) -> TorusPrecision {
 56        TorusPrecision(0)
 57    }
 58
 59    fn n(&self) -> Degree {
 60        Degree(self.data.n() as u32)
 61    }
 62
 63    fn size(&self) -> usize {
 64        1
 65    }
 66}
 67
 68impl<D: Data> GetDistribution for GLWESecret<D> {
 69    fn dist(&self) -> &Distribution {
 70        &self.dist
 71    }
 72}
 73
 74impl<D: Data> GLWEInfos for GLWESecret<D> {
 75    fn rank(&self) -> Rank {
 76        Rank(self.data.cols() as u32)
 77    }
 78}
 79
 80pub trait GLWESecretAlloc
 81where
 82    Self: GetDegree,
 83{
 84    fn alloc_glwe_secret(&self, rank: Rank) -> GLWESecret<Vec<u8>> {
 85        GLWESecret {
 86            data: ScalarZnx::alloc(self.ring_degree().into(), rank.into()),
 87            dist: Distribution::NONE,
 88        }
 89    }
 90
 91    fn alloc_glwe_secret_from_infos<A>(&self, infos: &A) -> GLWESecret<Vec<u8>>
 92    where
 93        A: GLWEInfos,
 94    {
 95        self.alloc_glwe_secret(infos.rank())
 96    }
 97
 98    fn bytes_of_glwe_secret(&self, rank: Rank) -> usize {
 99        ScalarZnx::bytes_of(self.ring_degree().into(), rank.into())
100    }
101
102    fn bytes_of_glwe_secret_from_infos<A>(&self, infos: &A) -> usize
103    where
104        A: GLWEInfos,
105    {
106        self.bytes_of_glwe_secret(infos.rank())
107    }
108}
109
110impl<B: Backend> GLWESecretAlloc for Module<B> where Self: GetDegree {}
111
112impl GLWESecret<Vec<u8>> {
113    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
114    where
115        A: GLWEInfos,
116        M: GLWESecretAlloc,
117    {
118        module.alloc_glwe_secret_from_infos(infos)
119    }
120
121    pub fn alloc<M>(module: &M, rank: Rank) -> Self
122    where
123        M: GLWESecretAlloc,
124    {
125        module.alloc_glwe_secret(rank)
126    }
127
128    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
129    where
130        A: GLWEInfos,
131        M: GLWESecretAlloc,
132    {
133        module.bytes_of_glwe_secret_from_infos(infos)
134    }
135
136    pub fn bytes_of<M>(module: &M, rank: Rank) -> usize
137    where
138        M: GLWESecretAlloc,
139    {
140        module.bytes_of_glwe_secret(rank)
141    }
142}
143
144impl<D: DataMut> GLWESecret<D> {
145    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
146        (0..self.rank().into()).for_each(|i| {
147            self.data.fill_ternary_prob(i, prob, source);
148        });
149        self.dist = Distribution::TernaryProb(prob);
150    }
151
152    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
153        (0..self.rank().into()).for_each(|i| {
154            self.data.fill_ternary_hw(i, hw, source);
155        });
156        self.dist = Distribution::TernaryFixed(hw);
157    }
158
159    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
160        (0..self.rank().into()).for_each(|i| {
161            self.data.fill_binary_prob(i, prob, source);
162        });
163        self.dist = Distribution::BinaryProb(prob);
164    }
165
166    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
167        (0..self.rank().into()).for_each(|i| {
168            self.data.fill_binary_hw(i, hw, source);
169        });
170        self.dist = Distribution::BinaryFixed(hw);
171    }
172
173    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
174        (0..self.rank().into()).for_each(|i| {
175            self.data.fill_binary_block(i, block_size, source);
176        });
177        self.dist = Distribution::BinaryBlock(block_size);
178    }
179
180    pub fn fill_zero(&mut self) {
181        self.data.zero();
182        self.dist = Distribution::ZERO;
183    }
184}
185
186pub trait GLWESecretToMut {
187    fn to_mut(&mut self) -> GLWESecret<&mut [u8]>;
188}
189
190impl<D: DataMut> GLWESecretToMut for GLWESecret<D> {
191    fn to_mut(&mut self) -> GLWESecret<&mut [u8]> {
192        GLWESecret {
193            dist: self.dist,
194            data: self.data.to_mut(),
195        }
196    }
197}
198
199pub trait GLWESecretToRef {
200    fn to_ref(&self) -> GLWESecret<&[u8]>;
201}
202
203impl<D: DataRef> GLWESecretToRef for GLWESecret<D> {
204    fn to_ref(&self) -> GLWESecret<&[u8]> {
205        GLWESecret {
206            data: self.data.to_ref(),
207            dist: self.dist,
208        }
209    }
210}
211
212impl<D: DataMut> ReaderFrom for GLWESecret<D> {
213    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
214        match Distribution::read_from(reader) {
215            Ok(dist) => self.dist = dist,
216            Err(e) => return Err(e),
217        }
218        self.data.read_from(reader)
219    }
220}
221
222impl<D: DataRef> WriterTo for GLWESecret<D> {
223    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
224        match self.dist.write_to(writer) {
225            Ok(()) => {}
226            Err(e) => return Err(e),
227        }
228        self.data.write_to(writer)
229    }
230}