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