1use poulpy_hal::{
  2    layouts::{Backend, Data, DataMut, DataRef, Module, 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
 16pub trait LWESecretAlloc {
 17    fn alloc_lwe_secret(&self, n: Degree) -> LWESecret<Vec<u8>> {
 18        LWESecret {
 19            data: ScalarZnx::alloc(n.into(), 1),
 20            dist: Distribution::NONE,
 21        }
 22    }
 23}
 24
 25impl<B: Backend> LWESecretAlloc for Module<B> {}
 26
 27impl LWESecret<Vec<u8>> {
 28    pub fn alloc<B: Backend>(module: &Module<B>, n: Degree) -> Self
 29    where
 30        Module<B>: LWESecretAlloc,
 31    {
 32        module.alloc_lwe_secret(n)
 33    }
 34}
 35
 36impl<D: DataRef> LWESecret<D> {
 37    pub fn raw(&self) -> &[i64] {
 38        self.data.at(0, 0)
 39    }
 40
 41    pub fn dist(&self) -> Distribution {
 42        self.dist
 43    }
 44
 45    pub fn data(&self) -> &ScalarZnx<D> {
 46        &self.data
 47    }
 48}
 49
 50impl<D: Data> LWEInfos for LWESecret<D> {
 51    fn base2k(&self) -> Base2K {
 52        Base2K(0)
 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: DataMut> LWESecret<D> {
 68    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 69        self.data.fill_ternary_prob(0, prob, source);
 70        self.dist = Distribution::TernaryProb(prob);
 71    }
 72
 73    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 74        self.data.fill_ternary_hw(0, hw, source);
 75        self.dist = Distribution::TernaryFixed(hw);
 76    }
 77
 78    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
 79        self.data.fill_binary_prob(0, prob, source);
 80        self.dist = Distribution::BinaryProb(prob);
 81    }
 82
 83    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
 84        self.data.fill_binary_hw(0, hw, source);
 85        self.dist = Distribution::BinaryFixed(hw);
 86    }
 87
 88    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
 89        self.data.fill_binary_block(0, block_size, source);
 90        self.dist = Distribution::BinaryBlock(block_size);
 91    }
 92
 93    pub fn fill_zero(&mut self) {
 94        self.data.zero();
 95        self.dist = Distribution::ZERO;
 96    }
 97}
 98
 99pub trait LWESecretToRef {
100    fn to_ref(&self) -> LWESecret<&[u8]>;
101}
102
103impl<D: DataRef> LWESecretToRef for LWESecret<D> {
104    fn to_ref(&self) -> LWESecret<&[u8]> {
105        LWESecret {
106            dist: self.dist,
107            data: self.data.to_ref(),
108        }
109    }
110}
111
112pub trait LWESecretToMut {
113    fn to_mut(&mut self) -> LWESecret<&mut [u8]>;
114}
115
116impl<D: DataMut> LWESecretToMut for LWESecret<D> {
117    fn to_mut(&mut self) -> LWESecret<&mut [u8]> {
118        LWESecret {
119            dist: self.dist,
120            data: self.data.to_mut(),
121        }
122    }
123}