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!(
 77            f,
 78            "GLWEPlaintext: base2k={} k={}: {}",
 79            self.base2k().0,
 80            self.k().0,
 81            self.data
 82        )
 83    }
 84}
 85
 86impl GLWEPlaintext<Vec<u8>> {
 87    pub fn alloc_from_infos<A>(infos: &A) -> Self
 88    where
 89        A: GLWEInfos,
 90    {
 91        Self::alloc(infos.n(), infos.base2k(), infos.k())
 92    }
 93
 94    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision) -> Self {
 95        GLWEPlaintext {
 96            data: VecZnx::alloc(n.into(), 1, k.0.div_ceil(base2k.0) as usize),
 97            base2k,
 98            k,
 99        }
100    }
101
102    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
103    where
104        A: GLWEInfos,
105    {
106        Self::bytes_of(infos.n(), infos.base2k(), infos.k())
107    }
108
109    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
110        VecZnx::bytes_of(n.into(), 1, k.0.div_ceil(base2k.0) as usize)
111    }
112}
113
114impl<D: DataRef> GLWEToRef for GLWEPlaintext<D> {
115    fn to_ref(&self) -> GLWE<&[u8]> {
116        GLWE {
117            k: self.k,
118            base2k: self.base2k,
119            data: self.data.to_ref(),
120        }
121    }
122}
123
124impl<D: DataMut> GLWEToMut for GLWEPlaintext<D> {
125    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
126        GLWE {
127            k: self.k,
128            base2k: self.base2k,
129            data: self.data.to_mut(),
130        }
131    }
132}
133
134pub trait GLWEPlaintextToRef {
135    fn to_ref(&self) -> GLWEPlaintext<&[u8]>;
136}
137
138impl<D: DataRef> GLWEPlaintextToRef for GLWEPlaintext<D> {
139    fn to_ref(&self) -> GLWEPlaintext<&[u8]> {
140        GLWEPlaintext {
141            data: self.data.to_ref(),
142            base2k: self.base2k,
143            k: self.k,
144        }
145    }
146}
147
148pub trait GLWEPlaintextToMut {
149    fn to_ref(&mut self) -> GLWEPlaintext<&mut [u8]>;
150}
151
152impl<D: DataMut> GLWEPlaintextToMut for GLWEPlaintext<D> {
153    fn to_ref(&mut self) -> GLWEPlaintext<&mut [u8]> {
154        GLWEPlaintext {
155            base2k: self.base2k,
156            k: self.k,
157            data: self.data.to_mut(),
158        }
159    }
160}
161
162impl<D: DataMut> GLWEPlaintext<D> {
163    pub fn data_mut(&mut self) -> &mut VecZnx<D> {
164        &mut self.data
165    }
166}
167
168impl<D: DataRef> GLWEPlaintext<D> {
169    pub fn data(&self) -> &VecZnx<D> {
170        &self.data
171    }
172}