1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, ReaderFrom, ScalarZnx, WriterTo, ZnxInfos, ZnxZero},
  3    source::Source,
  4};
  5
  6use crate::{
  7    dist::Distribution,
  8    layouts::{Base2K, Degree, GLWEInfos, LWEInfos, Rank, TorusPrecision},
  9};
 10
 11#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 12pub struct GLWESecretLayout {
 13    pub n: Degree,
 14    pub rank: Rank,
 15}
 16
 17impl LWEInfos for GLWESecretLayout {
 18    fn base2k(&self) -> Base2K {
 19        Base2K(0)
 20    }
 21
 22    fn k(&self) -> TorusPrecision {
 23        TorusPrecision(0)
 24    }
 25
 26    fn n(&self) -> Degree {
 27        self.n
 28    }
 29
 30    fn size(&self) -> usize {
 31        1
 32    }
 33}
 34impl GLWEInfos for GLWESecretLayout {
 35    fn rank(&self) -> Rank {
 36        self.rank
 37    }
 38}
 39
 40#[derive(PartialEq, Eq, Clone)]
 41pub struct GLWESecret<D: Data> {
 42    pub(crate) data: ScalarZnx<D>,
 43    pub(crate) dist: Distribution,
 44}
 45
 46impl<D: Data> LWEInfos for GLWESecret<D> {
 47    fn base2k(&self) -> Base2K {
 48        Base2K(0)
 49    }
 50
 51    fn k(&self) -> TorusPrecision {
 52        TorusPrecision(0)
 53    }
 54
 55    fn n(&self) -> Degree {
 56        Degree(self.data.n() as u32)
 57    }
 58
 59    fn size(&self) -> usize {
 60        1
 61    }
 62}
 63
 64impl<D: Data> GLWEInfos for GLWESecret<D> {
 65    fn rank(&self) -> Rank {
 66        Rank(self.data.cols() as u32)
 67    }
 68}
 69
 70impl GLWESecret<Vec<u8>> {
 71    pub fn alloc<A>(infos: &A) -> Self
 72    where
 73        A: GLWEInfos,
 74    {
 75        Self::alloc_with(infos.n(), infos.rank())
 76    }
 77
 78    pub fn alloc_with(n: Degree, rank: Rank) -> Self {
 79        Self {
 80            data: ScalarZnx::alloc(n.into(), rank.into()),
 81            dist: Distribution::NONE,
 82        }
 83    }
 84
 85    pub fn alloc_bytes<A>(infos: &A) -> usize
 86    where
 87        A: GLWEInfos,
 88    {
 89        Self::alloc_bytes_with(infos.n(), infos.rank())
 90    }
 91
 92    pub fn alloc_bytes_with(n: Degree, rank: Rank) -> usize {
 93        ScalarZnx::alloc_bytes(n.into(), rank.into())
 94    }
 95}
 96
 97impl<D: DataMut> GLWESecret<D> {
 98    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 99        (0..self.rank().into()).for_each(|i| {
100            self.data.fill_ternary_prob(i, prob, source);
101        });
102        self.dist = Distribution::TernaryProb(prob);
103    }
104
105    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
106        (0..self.rank().into()).for_each(|i| {
107            self.data.fill_ternary_hw(i, hw, source);
108        });
109        self.dist = Distribution::TernaryFixed(hw);
110    }
111
112    pub fn fill_binary_prob(&mut self, prob: f64, source: &mut Source) {
113        (0..self.rank().into()).for_each(|i| {
114            self.data.fill_binary_prob(i, prob, source);
115        });
116        self.dist = Distribution::BinaryProb(prob);
117    }
118
119    pub fn fill_binary_hw(&mut self, hw: usize, source: &mut Source) {
120        (0..self.rank().into()).for_each(|i| {
121            self.data.fill_binary_hw(i, hw, source);
122        });
123        self.dist = Distribution::BinaryFixed(hw);
124    }
125
126    pub fn fill_binary_block(&mut self, block_size: usize, source: &mut Source) {
127        (0..self.rank().into()).for_each(|i| {
128            self.data.fill_binary_block(i, block_size, source);
129        });
130        self.dist = Distribution::BinaryBlock(block_size);
131    }
132
133    pub fn fill_zero(&mut self) {
134        self.data.zero();
135        self.dist = Distribution::ZERO;
136    }
137}
138
139impl<D: DataMut> ReaderFrom for GLWESecret<D> {
140    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
141        match Distribution::read_from(reader) {
142            Ok(dist) => self.dist = dist,
143            Err(e) => return Err(e),
144        }
145        self.data.read_from(reader)
146    }
147}
148
149impl<D: DataRef> WriterTo for GLWESecret<D> {
150    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
151        match self.dist.write_to(writer) {
152            Ok(()) => {}
153            Err(e) => return Err(e),
154        }
155        self.data.write_to(writer)
156    }
157}