1use std::fmt;
  2
  3use poulpy_hal::layouts::{Data, DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos};
  4
  5use crate::layouts::{Base2K, Degree, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, Rank, SetGLWEInfos, TorusPrecision};
  6
  7#[derive(PartialEq, Eq, Copy, Clone, Debug)]
  8pub struct GLWEPlaintextLayout {
  9    pub n: Degree,
 10    pub base2k: Base2K,
 11    pub k: TorusPrecision,
 12}
 13
 14impl LWEInfos for GLWEPlaintextLayout {
 15    fn base2k(&self) -> Base2K {
 16        self.base2k
 17    }
 18
 19    fn k(&self) -> TorusPrecision {
 20        self.k
 21    }
 22
 23    fn n(&self) -> Degree {
 24        self.n
 25    }
 26}
 27
 28impl GLWEInfos for GLWEPlaintextLayout {
 29    fn rank(&self) -> Rank {
 30        Rank(0)
 31    }
 32}
 33
 34pub struct GLWEPlaintext<D: Data> {
 35    pub data: VecZnx<D>,
 36    pub base2k: Base2K,
 37    pub k: TorusPrecision,
 38}
 39
 40impl<D: DataMut> SetGLWEInfos for GLWEPlaintext<D> {
 41    fn set_base2k(&mut self, base2k: Base2K) {
 42        self.base2k = base2k
 43    }
 44
 45    fn set_k(&mut self, k: TorusPrecision) {
 46        self.k = k
 47    }
 48}
 49
 50impl<D: Data> LWEInfos for GLWEPlaintext<D> {
 51    fn base2k(&self) -> Base2K {
 52        self.base2k
 53    }
 54
 55    fn k(&self) -> TorusPrecision {
 56        self.k
 57    }
 58
 59    fn size(&self) -> usize {
 60        self.data.size()
 61    }
 62
 63    fn n(&self) -> Degree {
 64        Degree(self.data.n() as u32)
 65    }
 66}
 67
 68impl<D: Data> GLWEInfos for GLWEPlaintext<D> {
 69    fn rank(&self) -> Rank {
 70        Rank(self.data.cols() as u32 - 1)
 71    }
 72}
 73
 74impl<D: DataRef> fmt::Display for GLWEPlaintext<D> {
 75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 76        write!(f, "GLWEPlaintext: base2k={} k={}: {}", self.base2k().0, self.k().0, self.data)
 77    }
 78}
 79
 80impl GLWEPlaintext<Vec<u8>> {
 81    pub fn alloc_from_infos<A>(infos: &A) -> Self
 82    where
 83        A: GLWEInfos,
 84    {
 85        Self::alloc(infos.n(), infos.base2k(), infos.k())
 86    }
 87
 88    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
 89        GLWEPlaintext {
 90            data: VecZnx::alloc(n.into(), 1, k.0.div_ceil(base2k.0) as usize),
 91            base2k,
 92            k,
 93        }
 94    }
 95
 96    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
 97    where
 98        A: GLWEInfos,
 99    {
100        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
101    }
102
103    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
104        VecZnx::bytes_of(n.into(), 1, k.0.div_ceil(base2k.0) as usize)
105    }
106}
107
108impl<D: DataRef> GLWEToRef for GLWEPlaintext<D> {
109    fn to_ref(&self) -> GLWE<&[u8]> {
110        GLWE {
111            k: self.k,
112            base2k: self.base2k,
113            data: self.data.to_ref(),
114        }
115    }
116}
117
118impl<D: DataMut> GLWEToMut for GLWEPlaintext<D> {
119    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
120        GLWE {
121            k: self.k,
122            base2k: self.base2k,
123            data: self.data.to_mut(),
124        }
125    }
126}
127
128pub trait GLWEPlaintextToRef {
129    fn to_ref(&self) -> GLWEPlaintext<&[u8]>;
130}
131
132impl<D: DataRef> GLWEPlaintextToRef for GLWEPlaintext<D> {
133    fn to_ref(&self) -> GLWEPlaintext<&[u8]> {
134        GLWEPlaintext {
135            data: self.data.to_ref(),
136            base2k: self.base2k,
137            k: self.k,
138        }
139    }
140}
141
142pub trait GLWEPlaintextToMut {
143    fn to_ref(&mut self) -> GLWEPlaintext<&mut [u8]>;
144}
145
146impl<D: DataMut> GLWEPlaintextToMut for GLWEPlaintext<D> {
147    fn to_ref(&mut self) -> GLWEPlaintext<&mut [u8]> {
148        GLWEPlaintext {
149            base2k: self.base2k,
150            k: self.k,
151            data: self.data.to_mut(),
152        }
153    }
154}
155
156impl<D: DataMut> GLWEPlaintext<D> {
157    pub fn data_mut(&mut self) -> &mut VecZnx<D> {
158        &mut self.data
159    }
160}
161
162impl<D: DataRef> GLWEPlaintext<D> {
163    pub fn data(&self) -> &VecZnx<D> {
164        &self.data
165    }
166}