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