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