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, Degree, GLWEInfos, GetDegree, LWEInfos, Rank, TorusPrecision},
 12};
 13
 14#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 15pub struct GLWESecretLayout {
 16    pub n: Degree,
 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) -> Degree {
 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) -> Degree {
 59        Degree(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: GetDegree,
 76{
 77    fn alloc_glwe_secret(&self, rank: Rank) -> GLWESecret<Vec<u8>> {
 78        GLWESecret {
 79            data: ScalarZnx::alloc(self.n().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.n().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 GLWESecret<Vec<u8>> {
104    pub fn alloc_from_infos<A, B: Backend>(module: Module<B>, infos: &A) -> Self
105    where
106        A: GLWEInfos,
107        Module<B>: GLWESecretAlloc,
108    {
109        module.alloc_glwe_secret_from_infos(infos)
110    }
111
112    pub fn alloc<B: Backend>(module: Module<B>, rank: Rank) -> Self
113    where
114        Module<B>: GLWESecretAlloc,
115    {
116        module.alloc_glwe_secret(rank)
117    }
118
119    pub fn bytes_of_from_infos<A, B: Backend>(module: Module<B>, infos: &A) -> usize
120    where
121        A: GLWEInfos,
122        Module<B>: GLWESecretAlloc,
123    {
124        module.bytes_of_glwe_secret_from_infos(infos)
125    }
126
127    pub fn bytes_of<B: Backend>(module: Module<B>, rank: Rank) -> usize
128    where
129        Module<B>: GLWESecretAlloc,
130    {
131        module.bytes_of_glwe_secret(rank)
132    }
133}
134
135impl<D: DataMut> GLWESecret<D> {
136    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
137        (0..self.rank().into()).for_each(|i| {
138            self.data.fill_ternary_prob(i, prob, source);
139        });
140        self.dist = Distribution::TernaryProb(prob);
141    }
142
143    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
144        (0..self.rank().into()).for_each(|i| {
145            self.data.fill_ternary_hw(i, hw, source);
146        });
147        self.dist = Distribution::TernaryFixed(hw);
148    }
149
150    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
151        (0..self.rank().into()).for_each(|i| {
152            self.data.fill_binary_prob(i, prob, source);
153        });
154        self.dist = Distribution::BinaryProb(prob);
155    }
156
157    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
158        (0..self.rank().into()).for_each(|i| {
159            self.data.fill_binary_hw(i, hw, source);
160        });
161        self.dist = Distribution::BinaryFixed(hw);
162    }
163
164    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
165        (0..self.rank().into()).for_each(|i| {
166            self.data.fill_binary_block(i, block_size, source);
167        });
168        self.dist = Distribution::BinaryBlock(block_size);
169    }
170
171    pub fn fill_zero(&mut self) {
172        self.data.zero();
173        self.dist = Distribution::ZERO;
174    }
175}
176
177pub trait GLWESecretToMut {
178    fn to_mut(&mut self) -> GLWESecret<&mut [u8]>;
179}
180
181impl<D: DataMut> GLWESecretToMut for GLWESecret<D> {
182    fn to_mut(&mut self) -> GLWESecret<&mut [u8]> {
183        GLWESecret {
184            dist: self.dist,
185            data: self.data.to_mut(),
186        }
187    }
188}
189
190pub trait GLWESecretToRef {
191    fn to_ref(&self) -> GLWESecret<&[u8]>;
192}
193
194impl<D: DataRef> GLWESecretToRef for GLWESecret<D> {
195    fn to_ref(&self) -> GLWESecret<&[u8]> {
196        GLWESecret {
197            data: self.data.to_ref(),
198            dist: self.dist,
199        }
200    }
201}
202
203impl<D: DataMut> ReaderFrom for GLWESecret<D> {
204    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
205        match Distribution::read_from(reader) {
206            Ok(dist) => self.dist = dist,
207            Err(e) => return Err(e),
208        }
209        self.data.read_from(reader)
210    }
211}
212
213impl<D: DataRef> WriterTo for GLWESecret<D> {
214    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
215        match self.dist.write_to(writer) {
216            Ok(()) => {}
217            Err(e) => return Err(e),
218        }
219        self.data.write_to(writer)
220    }
221}