1use poulpy_hal::{
  2    layouts::{
  3        Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, ToOwnedDeep, VecZnx, VecZnxToMut, VecZnxToRef,
  4        WriterTo, ZnxInfos,
  5    },
  6    source::Source,
  7};
  8
  9use crate::layouts::{Base2K, GetRingDegree, LWEInfos, Rank, RingDegree, TorusPrecision};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11use std::fmt;
 12
 13pub trait GLWEInfos
 14where
 15    Self: LWEInfos,
 16{
 17    fn rank(&self) -> Rank;
 18    fn glwe_layout(&self) -> GLWELayout {
 19        GLWELayout {
 20            n: self.n(),
 21            base2k: self.base2k(),
 22            k: self.k(),
 23            rank: self.rank(),
 24        }
 25    }
 26}
 27
 28pub trait SetGLWEInfos {
 29    fn set_k(&mut self, k: TorusPrecision);
 30    fn set_base2k(&mut self, base2k: Base2K);
 31}
 32
 33#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 34pub struct GLWELayout {
 35    pub n: RingDegree,
 36    pub base2k: Base2K,
 37    pub k: TorusPrecision,
 38    pub rank: Rank,
 39}
 40
 41impl LWEInfos for GLWELayout {
 42    fn n(&self) -> RingDegree {
 43        self.n
 44    }
 45
 46    fn base2k(&self) -> Base2K {
 47        self.base2k
 48    }
 49
 50    fn k(&self) -> TorusPrecision {
 51        self.k
 52    }
 53}
 54
 55impl GLWEInfos for GLWELayout {
 56    fn rank(&self) -> Rank {
 57        self.rank
 58    }
 59}
 60
 61#[derive(PartialEq, Eq, Clone)]
 62pub struct GLWE<D: Data> {
 63    pub(crate) data: VecZnx<D>,
 64    pub(crate) base2k: Base2K,
 65    pub(crate) k: TorusPrecision,
 66}
 67
 68impl<D: DataMut> SetGLWEInfos for GLWE<D> {
 69    fn set_base2k(&mut self, base2k: Base2K) {
 70        self.base2k = base2k
 71    }
 72
 73    fn set_k(&mut self, k: TorusPrecision) {
 74        self.k = k
 75    }
 76}
 77
 78impl<D: DataRef> GLWE<D> {
 79    pub fn data(&self) -> &VecZnx<D> {
 80        &self.data
 81    }
 82}
 83
 84impl<D: DataMut> GLWE<D> {
 85    pub fn data_mut(&mut self) -> &mut VecZnx<D> {
 86        &mut self.data
 87    }
 88}
 89
 90impl<D: Data> LWEInfos for GLWE<D> {
 91    fn base2k(&self) -> Base2K {
 92        self.base2k
 93    }
 94
 95    fn k(&self) -> TorusPrecision {
 96        self.k
 97    }
 98
 99    fn n(&self) -> RingDegree {
100        RingDegree(self.data.n() as u32)
101    }
102
103    fn size(&self) -> usize {
104        self.data.size()
105    }
106}
107
108impl<D: Data> GLWEInfos for GLWE<D> {
109    fn rank(&self) -> Rank {
110        Rank(self.data.cols() as u32 - 1)
111    }
112}
113
114impl<D: DataRef> ToOwnedDeep for GLWE<D> {
115    type Owned = GLWE<Vec<u8>>;
116    fn to_owned_deep(&self) -> Self::Owned {
117        GLWE {
118            data: self.data.to_owned_deep(),
119            k: self.k,
120            base2k: self.base2k,
121        }
122    }
123}
124
125impl<D: DataRef> fmt::Debug for GLWE<D> {
126    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
127        write!(f, "{self}")
128    }
129}
130
131impl<D: DataRef> fmt::Display for GLWE<D> {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        write!(
134            f,
135            "GLWE: base2k={} k={}: {}",
136            self.base2k().0,
137            self.k().0,
138            self.data
139        )
140    }
141}
142
143impl<D: DataMut> FillUniform for GLWE<D> {
144    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
145        self.data.fill_uniform(log_bound, source);
146    }
147}
148
149pub trait GLWEAlloc
150where
151    Self: GetRingDegree,
152{
153    fn alloc_glwe(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> GLWE<Vec<u8>> {
154        GLWE {
155            data: VecZnx::alloc(
156                self.ring_degree().into(),
157                (rank + 1).into(),
158                k.0.div_ceil(base2k.0) as usize,
159            ),
160            base2k,
161            k,
162        }
163    }
164
165    fn alloc_glwe_from_infos<A>(&self, infos: &A) -> GLWE<Vec<u8>>
166    where
167        A: GLWEInfos,
168    {
169        self.alloc_glwe(infos.base2k(), infos.k(), infos.rank())
170    }
171
172    fn bytes_of_glwe(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize {
173        VecZnx::bytes_of(
174            self.ring_degree().into(),
175            (rank + 1).into(),
176            k.0.div_ceil(base2k.0) as usize,
177        )
178    }
179
180    fn bytes_of_glwe_from_infos<A>(&self, infos: &A) -> usize
181    where
182        A: GLWEInfos,
183    {
184        self.bytes_of_glwe(infos.base2k(), infos.k(), infos.rank())
185    }
186}
187
188impl<B: Backend> GLWEAlloc for Module<B> where Self: GetRingDegree {}
189
190impl GLWE<Vec<u8>> {
191    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
192    where
193        A: GLWEInfos,
194        M: GLWEAlloc,
195    {
196        module.alloc_glwe_from_infos(infos)
197    }
198
199    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
200    where
201        M: GLWEAlloc,
202    {
203        module.alloc_glwe(base2k, k, rank)
204    }
205
206    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
207    where
208        A: GLWEInfos,
209        M: GLWEAlloc,
210    {
211        module.bytes_of_glwe_from_infos(infos)
212    }
213
214    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> usize
215    where
216        M: GLWEAlloc,
217    {
218        module.bytes_of_glwe(base2k, k, rank)
219    }
220}
221
222impl<D: DataMut> ReaderFrom for GLWE<D> {
223    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
224        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
225        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
226        self.data.read_from(reader)
227    }
228}
229
230impl<D: DataRef> WriterTo for GLWE<D> {
231    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
232        writer.write_u32::<LittleEndian>(self.k.0)?;
233        writer.write_u32::<LittleEndian>(self.base2k.0)?;
234        self.data.write_to(writer)
235    }
236}
237
238pub trait GLWEToRef {
239    fn to_ref(&self) -> GLWE<&[u8]>;
240}
241
242impl<D: DataRef> GLWEToRef for GLWE<D> {
243    fn to_ref(&self) -> GLWE<&[u8]> {
244        GLWE {
245            k: self.k,
246            base2k: self.base2k,
247            data: self.data.to_ref(),
248        }
249    }
250}
251
252pub trait GLWEToMut {
253    fn to_mut(&mut self) -> GLWE<&mut [u8]>;
254}
255
256impl<D: DataMut> GLWEToMut for GLWE<D> {
257    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
258        GLWE {
259            k: self.k,
260            base2k: self.base2k,
261            data: self.data.to_mut(),
262        }
263    }
264}