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, 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> GLWEInfos for GLWESecret<D> {
 68    fn rank(&self) -> Rank {
 69        Rank(self.data.cols() as u32)
 70    }
 71}
 72
 73pub trait GLWESecretAlloc
 74where
 75    Self: GetRingDegree,
 76{
 77    fn alloc_glwe_secret(&self, rank: Rank) -> GLWESecret<Vec<u8>> {
 78        GLWESecret {
 79            data: ScalarZnx::alloc(self.ring_degree().into(), rank.into()),
 80            dist: Distribution::NONE,
 81        }
 82    }
 83
 84    fn alloc_glwe_secret_from_infos<A>(&self, infos: &A) -> GLWESecret<Vec<u8>>
 85    where
 86        A: GLWEInfos,
 87    {
 88        self.alloc_glwe_secret(infos.rank())
 89    }
 90
 91    fn bytes_of_glwe_secret(&self, rank: Rank) -> usize {
 92        ScalarZnx::bytes_of(self.ring_degree().into(), rank.into())
 93    }
 94
 95    fn bytes_of_glwe_secret_from_infos<A>(&self, infos: &A) -> usize
 96    where
 97        A: GLWEInfos,
 98    {
 99        self.bytes_of_glwe_secret(infos.rank())
100    }
101}
102
103impl<B: Backend> GLWESecretAlloc for Module<B> where Self: GetRingDegree {}
104
105impl GLWESecret<Vec<u8>> {
106    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
107    where
108        A: GLWEInfos,
109        M: GLWESecretAlloc,
110    {
111        module.alloc_glwe_secret_from_infos(infos)
112    }
113
114    pub fn alloc<M>(module: &M, rank: Rank) -> Self
115    where
116        M: GLWESecretAlloc,
117    {
118        module.alloc_glwe_secret(rank)
119    }
120
121    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
122    where
123        A: GLWEInfos,
124        M: GLWESecretAlloc,
125    {
126        module.bytes_of_glwe_secret_from_infos(infos)
127    }
128
129    pub fn bytes_of<M>(module: &M, rank: Rank) -> usize
130    where
131        M: GLWESecretAlloc,
132    {
133        module.bytes_of_glwe_secret(rank)
134    }
135}
136
137impl<D: DataMut> GLWESecret<D> {
138    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
139        (0..self.rank().into()).for_each(|i| {
140            self.data.fill_ternary_prob(i, prob, source);
141        });
142        self.dist = Distribution::TernaryProb(prob);
143    }
144
145    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
146        (0..self.rank().into()).for_each(|i| {
147            self.data.fill_ternary_hw(i, hw, source);
148        });
149        self.dist = Distribution::TernaryFixed(hw);
150    }
151
152    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
153        (0..self.rank().into()).for_each(|i| {
154            self.data.fill_binary_prob(i, prob, source);
155        });
156        self.dist = Distribution::BinaryProb(prob);
157    }
158
159    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
160        (0..self.rank().into()).for_each(|i| {
161            self.data.fill_binary_hw(i, hw, source);
162        });
163        self.dist = Distribution::BinaryFixed(hw);
164    }
165
166    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
167        (0..self.rank().into()).for_each(|i| {
168            self.data.fill_binary_block(i, block_size, source);
169        });
170        self.dist = Distribution::BinaryBlock(block_size);
171    }
172
173    pub fn fill_zero(&mut self) {
174        self.data.zero();
175        self.dist = Distribution::ZERO;
176    }
177}
178
179pub trait GLWESecretToMut {
180    fn to_mut(&mut self) -> GLWESecret<&mut [u8]>;
181}
182
183impl<D: DataMut> GLWESecretToMut for GLWESecret<D> {
184    fn to_mut(&mut self) -> GLWESecret<&mut [u8]> {
185        GLWESecret {
186            dist: self.dist,
187            data: self.data.to_mut(),
188        }
189    }
190}
191
192pub trait GLWESecretToRef {
193    fn to_ref(&self) -> GLWESecret<&[u8]>;
194}
195
196impl<D: DataRef> GLWESecretToRef for GLWESecret<D> {
197    fn to_ref(&self) -> GLWESecret<&[u8]> {
198        GLWESecret {
199            data: self.data.to_ref(),
200            dist: self.dist,
201        }
202    }
203}
204
205impl<D: DataMut> ReaderFrom for GLWESecret<D> {
206    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
207        match Distribution::read_from(reader) {
208            Ok(dist) => self.dist = dist,
209            Err(e) => return Err(e),
210        }
211        self.data.read_from(reader)
212    }
213}
214
215impl<D: DataRef> WriterTo for GLWESecret<D> {
216    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
217        match self.dist.write_to(writer) {
218            Ok(()) => {}
219            Err(e) => return Err(e),
220        }
221        self.data.write_to(writer)
222    }
223}