1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, ReaderFrom, ScalarZnx, WriterTo, ZnxInfos, ZnxZero},
  3    source::Source,
  4};
  5
  6use crate::dist::Distribution;
  7
  8#[derive(PartialEq, Eq, Clone)]
  9pub struct GLWESecret<D: Data> {
 10    pub(crate) data: ScalarZnx<D>,
 11    pub(crate) dist: Distribution,
 12}
 13
 14impl GLWESecret<Vec<u8>> {
 15    pub fn alloc(n: usize, rank: usize) -> Self {
 16        Self {
 17            data: ScalarZnx::alloc(n, rank),
 18            dist: Distribution::NONE,
 19        }
 20    }
 21
 22    pub fn bytes_of(n: usize, rank: usize) -> usize {
 23        ScalarZnx::alloc_bytes(n, rank)
 24    }
 25}
 26
 27impl<D: Data> GLWESecret<D> {
 28    pub fn n(&self) -> usize {
 29        self.data.n()
 30    }
 31
 32    pub fn log_n(&self) -> usize {
 33        self.data.log_n()
 34    }
 35
 36    pub fn rank(&self) -> usize {
 37        self.data.cols()
 38    }
 39}
 40
 41impl<D: DataMut> GLWESecret<D> {
 42    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 43        (0..self.rank()).for_each(|i| {
 44            self.data.fill_ternary_prob(i, prob, source);
 45        });
 46        self.dist = Distribution::TernaryProb(prob);
 47    }
 48
 49    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 50        (0..self.rank()).for_each(|i| {
 51            self.data.fill_ternary_hw(i, hw, source);
 52        });
 53        self.dist = Distribution::TernaryFixed(hw);
 54    }
 55
 56    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
 57        (0..self.rank()).for_each(|i| {
 58            self.data.fill_binary_prob(i, prob, source);
 59        });
 60        self.dist = Distribution::BinaryProb(prob);
 61    }
 62
 63    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
 64        (0..self.rank()).for_each(|i| {
 65            self.data.fill_binary_hw(i, hw, source);
 66        });
 67        self.dist = Distribution::BinaryFixed(hw);
 68    }
 69
 70    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
 71        (0..self.rank()).for_each(|i| {
 72            self.data.fill_binary_block(i, block_size, source);
 73        });
 74        self.dist = Distribution::BinaryBlock(block_size);
 75    }
 76
 77    pub fn fill_zero(&mut self) {
 78        self.data.zero();
 79        self.dist = Distribution::ZERO;
 80    }
 81}
 82
 83impl<D: DataMut> ReaderFrom for GLWESecret<D> {
 84    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 85        match Distribution::read_from(reader) {
 86            Ok(dist) => self.dist = dist,
 87            Err(e) => return Err(e),
 88        }
 89        self.data.read_from(reader)
 90    }
 91}
 92
 93impl<D: DataRef> WriterTo for GLWESecret<D> {
 94    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
 95        match self.dist.write_to(writer) {
 96            Ok(()) => {}
 97            Err(e) => return Err(e),
 98        }
 99        self.data.write_to(writer)
100    }
101}