1use backend::hal::{
  2    api::{MatZnxAlloc, MatZnxAllocBytes, VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, Scratch, VmpPMat, WriterTo},
  4};
  5
  6use crate::{GLWECiphertext, Infos};
  7
  8pub trait GGSWLayoutFamily<B: Backend> = VmpPMatAlloc<B> + VmpPMatAllocBytes + VmpPMatPrepare<B>;
  9
 10#[derive(PartialEq, Eq)]
 11pub struct GGSWCiphertext<D: Data> {
 12    pub(crate) data: MatZnx<D>,
 13    pub(crate) basek: usize,
 14    pub(crate) k: usize,
 15    pub(crate) digits: usize,
 16}
 17
 18impl<D: DataRef> GGSWCiphertext<D> {
 19    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
 20        GLWECiphertext {
 21            data: self.data.at(row, col),
 22            basek: self.basek,
 23            k: self.k,
 24        }
 25    }
 26}
 27
 28impl<D: DataMut> GGSWCiphertext<D> {
 29    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
 30        GLWECiphertext {
 31            data: self.data.at_mut(row, col),
 32            basek: self.basek,
 33            k: self.k,
 34        }
 35    }
 36}
 37
 38impl GGSWCiphertext<Vec<u8>> {
 39    pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
 40    where
 41        Module<B>: MatZnxAlloc,
 42    {
 43        let size: usize = k.div_ceil(basek);
 44        debug_assert!(digits > 0, "invalid ggsw: `digits` == 0");
 45
 46        debug_assert!(
 47            size > digits,
 48            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 49            size,
 50            digits
 51        );
 52
 53        assert!(
 54            rows * digits <= size,
 55            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 56            rows,
 57            digits,
 58            size
 59        );
 60
 61        Self {
 62            data: module.mat_znx_alloc(rows, rank + 1, rank + 1, k.div_ceil(basek)),
 63            basek,
 64            k: k,
 65            digits,
 66        }
 67    }
 68
 69    pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
 70    where
 71        Module<B>: MatZnxAllocBytes,
 72    {
 73        let size: usize = k.div_ceil(basek);
 74        debug_assert!(
 75            size > digits,
 76            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
 77            size,
 78            digits
 79        );
 80
 81        assert!(
 82            rows * digits <= size,
 83            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
 84            rows,
 85            digits,
 86            size
 87        );
 88
 89        module.mat_znx_alloc_bytes(rows, rank + 1, rank + 1, size)
 90    }
 91}
 92
 93impl<D: Data> Infos for GGSWCiphertext<D> {
 94    type Inner = MatZnx<D>;
 95
 96    fn inner(&self) -> &Self::Inner {
 97        &self.data
 98    }
 99
100    fn basek(&self) -> usize {
101        self.basek
102    }
103
104    fn k(&self) -> usize {
105        self.k
106    }
107}
108
109impl<D: Data> GGSWCiphertext<D> {
110    pub fn rank(&self) -> usize {
111        self.data.cols_out() - 1
112    }
113
114    pub fn digits(&self) -> usize {
115        self.digits
116    }
117}
118
119use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
120
121impl<D: DataMut> ReaderFrom for GGSWCiphertext<D> {
122    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
123        self.k = reader.read_u64::<LittleEndian>()? as usize;
124        self.basek = reader.read_u64::<LittleEndian>()? as usize;
125        self.digits = reader.read_u64::<LittleEndian>()? as usize;
126        self.data.read_from(reader)
127    }
128}
129
130impl<D: DataRef> WriterTo for GGSWCiphertext<D> {
131    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
132        writer.write_u64::<LittleEndian>(self.k as u64)?;
133        writer.write_u64::<LittleEndian>(self.basek as u64)?;
134        writer.write_u64::<LittleEndian>(self.digits as u64)?;
135        self.data.write_to(writer)
136    }
137}
138
139#[derive(PartialEq, Eq)]
140pub struct GGSWCiphertextExec<D: Data, B: Backend> {
141    pub(crate) data: VmpPMat<D, B>,
142    pub(crate) basek: usize,
143    pub(crate) k: usize,
144    pub(crate) digits: usize,
145}
146
147impl<B: Backend> GGSWCiphertextExec<Vec<u8>, B> {
148    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
149    where
150        Module<B>: GGSWLayoutFamily<B>,
151    {
152        let size: usize = k.div_ceil(basek);
153        debug_assert!(digits > 0, "invalid ggsw: `digits` == 0");
154
155        debug_assert!(
156            size > digits,
157            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
158            size,
159            digits
160        );
161
162        assert!(
163            rows * digits <= size,
164            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
165            rows,
166            digits,
167            size
168        );
169
170        Self {
171            data: module.vmp_pmat_alloc(rows, rank + 1, rank + 1, k.div_ceil(basek)),
172            basek,
173            k: k,
174            digits,
175        }
176    }
177
178    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
179    where
180        Module<B>: GGSWLayoutFamily<B>,
181    {
182        let size: usize = k.div_ceil(basek);
183        debug_assert!(
184            size > digits,
185            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
186            size,
187            digits
188        );
189
190        assert!(
191            rows * digits <= size,
192            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
193            rows,
194            digits,
195            size
196        );
197
198        module.vmp_pmat_alloc_bytes(rows, rank + 1, rank + 1, size)
199    }
200
201    pub fn from<DataOther: DataRef>(
202        module: &Module<B>,
203        other: &GGSWCiphertext<DataOther>,
204        scratch: &mut Scratch<B>,
205    ) -> GGSWCiphertextExec<Vec<u8>, B>
206    where
207        Module<B>: GGSWLayoutFamily<B>,
208    {
209        let mut ggsw_exec: GGSWCiphertextExec<Vec<u8>, B> = Self::alloc(
210            module,
211            other.basek(),
212            other.k(),
213            other.rows(),
214            other.digits(),
215            other.rank(),
216        );
217        ggsw_exec.prepare(module, other, scratch);
218        ggsw_exec
219    }
220}
221
222impl<D: Data, B: Backend> Infos for GGSWCiphertextExec<D, B> {
223    type Inner = VmpPMat<D, B>;
224
225    fn inner(&self) -> &Self::Inner {
226        &self.data
227    }
228
229    fn basek(&self) -> usize {
230        self.basek
231    }
232
233    fn k(&self) -> usize {
234        self.k
235    }
236}
237
238impl<D: Data, B: Backend> GGSWCiphertextExec<D, B> {
239    pub fn rank(&self) -> usize {
240        self.data.cols_out() - 1
241    }
242
243    pub fn digits(&self) -> usize {
244        self.digits
245    }
246}
247
248impl<DataSelf: DataMut, B: Backend> GGSWCiphertextExec<DataSelf, B> {
249    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGSWCiphertext<DataOther>, scratch: &mut Scratch<B>)
250    where
251        DataOther: DataRef,
252        Module<B>: GGSWLayoutFamily<B>,
253    {
254        module.vmp_prepare(&mut self.data, &other.data, scratch);
255        self.k = other.k;
256        self.basek = other.basek;
257        self.digits = other.digits;
258    }
259}