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