1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, ScalarZnx, ScalarZnxToMut, ScalarZnxToRef, ZnxInfos, ZnxView, ZnxZero},
  3    source::Source,
  4};
  5
  6use crate::{
  7    GetDistribution,
  8    dist::Distribution,
  9    layouts::{Base2K, Degree, LWEInfos, TorusPrecision},
 10};
 11
 12pub struct LWESecret<D: Data> {
 13    pub(crate) data: ScalarZnx<D>,
 14    pub(crate) dist: Distribution,
 15}
 16
 17impl LWESecret<Vec<u8>> {
 18    pub fn alloc(n: Degree) -> Self {
 19        LWESecret {
 20            data: ScalarZnx::alloc(n.into(), 1),
 21            dist: Distribution::NONE,
 22        }
 23    }
 24}
 25
 26impl<D: DataRef> GetDistribution for LWESecret<D> {
 27    fn dist(&self) -> &Distribution {
 28        &self.dist
 29    }
 30}
 31
 32impl<D: DataRef> LWESecret<D> {
 33    pub fn raw(&self) -> &[i64] {
 34        self.data.at(0, 0)
 35    }
 36
 37    pub fn dist(&self) -> Distribution {
 38        self.dist
 39    }
 40
 41    pub fn data(&self) -> &ScalarZnx<D> {
 42        &self.data
 43    }
 44}
 45
 46impl<D: Data> LWEInfos for LWESecret<D> {
 47    fn base2k(&self) -> Base2K {
 48        Base2K(0)
 49    }
 50    fn k(&self) -> TorusPrecision {
 51        TorusPrecision(0)
 52    }
 53
 54    fn n(&self) -> Degree {
 55        Degree(self.data.n() as u32)
 56    }
 57
 58    fn size(&self) -> usize {
 59        1
 60    }
 61}
 62
 63impl<D: DataMut> LWESecret<D> {
 64    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 65        self.data.fill_ternary_prob(0, prob, source);
 66        self.dist = Distribution::TernaryProb(prob);
 67    }
 68
 69    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 70        self.data.fill_ternary_hw(0, hw, source);
 71        self.dist = Distribution::TernaryFixed(hw);
 72    }
 73
 74    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
 75        self.data.fill_binary_prob(0, prob, source);
 76        self.dist = Distribution::BinaryProb(prob);
 77    }
 78
 79    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
 80        self.data.fill_binary_hw(0, hw, source);
 81        self.dist = Distribution::BinaryFixed(hw);
 82    }
 83
 84    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
 85        self.data.fill_binary_block(0, block_size, source);
 86        self.dist = Distribution::BinaryBlock(block_size);
 87    }
 88
 89    pub fn fill_zero(&mut self) {
 90        self.data.zero();
 91        self.dist = Distribution::ZERO;
 92    }
 93}
 94
 95pub trait LWESecretToRef {
 96    fn to_ref(&self) -> LWESecret<&[u8]>;
 97}
 98
 99impl<D: DataRef> LWESecretToRef for LWESecret<D> {
100    fn to_ref(&self) -> LWESecret<&[u8]> {
101        LWESecret {
102            dist: self.dist,
103            data: self.data.to_ref(),
104        }
105    }
106}
107
108pub trait LWESecretToMut {
109    fn to_mut(&mut self) -> LWESecret<&mut [u8]>;
110}
111
112impl<D: DataMut> LWESecretToMut for LWESecret<D> {
113    fn to_mut(&mut self) -> LWESecret<&mut [u8]> {
114        LWESecret {
115            dist: self.dist,
116            data: self.data.to_mut(),
117        }
118    }
119}