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