1use crate::{
  2    alloc_aligned,
  3    layouts::{
  4        Data, DataMut, DataRef, DataView, DataViewMut, DigestU64, FillUniform, ReaderFrom, Reset, ToOwnedDeep, VecZnx, WriterTo,
  5        ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
  6    },
  7    source::Source,
  8};
  9use std::{
 10    fmt,
 11    hash::{DefaultHasher, Hasher},
 12};
 13
 14use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 15use rand::RngCore;
 16
 17#[repr(C)]
 18#[derive(PartialEq, Eq, Clone, Hash)]
 19pub struct MatZnx<D: Data> {
 20    data: D,
 21    n: usize,
 22    size: usize,
 23    rows: usize,
 24    cols_in: usize,
 25    cols_out: usize,
 26}
 27
 28impl<D: DataRef> DigestU64 for MatZnx<D> {
 29    fn digest_u64(&self) -> u64 {
 30        let mut h: DefaultHasher = DefaultHasher::new();
 31        h.write(self.data.as_ref());
 32        h.write_usize(self.n);
 33        h.write_usize(self.size);
 34        h.write_usize(self.rows);
 35        h.write_usize(self.cols_in);
 36        h.write_usize(self.cols_out);
 37        h.finish()
 38    }
 39}
 40
 41impl<D: DataRef> ToOwnedDeep for MatZnx<D> {
 42    type Owned = MatZnx<Vec<u8>>;
 43    fn to_owned_deep(&self) -> Self::Owned {
 44        MatZnx {
 45            data: self.data.as_ref().to_vec(),
 46            n: self.n,
 47            size: self.size,
 48            rows: self.rows,
 49            cols_in: self.cols_in,
 50            cols_out: self.cols_out,
 51        }
 52    }
 53}
 54
 55impl<D: DataRef> fmt::Debug for MatZnx<D> {
 56    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 57        write!(f, "{}", self)
 58    }
 59}
 60
 61impl<D: Data> ZnxInfos for MatZnx<D> {
 62    fn cols(&self) -> usize {
 63        self.cols_in
 64    }
 65
 66    fn rows(&self) -> usize {
 67        self.rows
 68    }
 69
 70    fn n(&self) -> usize {
 71        self.n
 72    }
 73
 74    fn size(&self) -> usize {
 75        self.size
 76    }
 77
 78    fn poly_count(&self) -> usize {
 79        self.rows() * self.cols_in() * self.cols_out() * self.size()
 80    }
 81}
 82
 83impl<D: Data> ZnxSliceSize for MatZnx<D> {
 84    fn sl(&self) -> usize {
 85        self.n() * self.cols_out()
 86    }
 87}
 88
 89impl<D: Data> DataView for MatZnx<D> {
 90    type D = D;
 91    fn data(&self) -> &Self::D {
 92        &self.data
 93    }
 94}
 95
 96impl<D: Data> DataViewMut for MatZnx<D> {
 97    fn data_mut(&mut self) -> &mut Self::D {
 98        &mut self.data
 99    }
100}
101
102impl<D: DataRef> ZnxView for MatZnx<D> {
103    type Scalar = i64;
104}
105
106impl<D: Data> MatZnx<D> {
107    pub fn cols_in(&self) -> usize {
108        self.cols_in
109    }
110
111    pub fn cols_out(&self) -> usize {
112        self.cols_out
113    }
114}
115
116impl MatZnx<Vec<u8>> {
117    pub fn alloc_bytes(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
118        rows * cols_in * VecZnx::<Vec<u8>>::alloc_bytes(n, cols_out, size)
119    }
120
121    pub fn alloc(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
122        let data: Vec<u8> = alloc_aligned(Self::alloc_bytes(n, rows, cols_in, cols_out, size));
123        Self {
124            data,
125            n,
126            size,
127            rows,
128            cols_in,
129            cols_out,
130        }
131    }
132
133    pub fn from_bytes(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
134        let data: Vec<u8> = bytes.into();
135        assert!(data.len() == Self::alloc_bytes(n, rows, cols_in, cols_out, size));
136        Self {
137            data,
138            n,
139            size,
140            rows,
141            cols_in,
142            cols_out,
143        }
144    }
145}
146
147impl<D: DataRef> MatZnx<D> {
148    pub fn at(&self, row: usize, col: usize) -> VecZnx<&[u8]> {
149        #[cfg(debug_assertions)]
150        {
151            assert!(row < self.rows(), "rows: {} >= {}", row, self.rows());
152            assert!(col < self.cols_in(), "cols: {} >= {}", col, self.cols_in());
153        }
154
155        let self_ref: MatZnx<&[u8]> = self.to_ref();
156        let nb_bytes: usize = VecZnx::<Vec<u8>>::alloc_bytes(self.n, self.cols_out, self.size);
157        let start: usize = nb_bytes * self.cols() * row + col * nb_bytes;
158        let end: usize = start + nb_bytes;
159
160        VecZnx {
161            data: &self_ref.data[start..end],
162            n: self.n,
163            cols: self.cols_out,
164            size: self.size,
165            max_size: self.size,
166        }
167    }
168}
169
170impl<D: DataMut> MatZnx<D> {
171    pub fn at_mut(&mut self, row: usize, col: usize) -> VecZnx<&mut [u8]> {
172        #[cfg(debug_assertions)]
173        {
174            assert!(row < self.rows(), "rows: {} >= {}", row, self.rows());
175            assert!(col < self.cols_in(), "cols: {} >= {}", col, self.cols_in());
176        }
177
178        let n: usize = self.n();
179        let cols_out: usize = self.cols_out();
180        let cols_in: usize = self.cols_in();
181        let size: usize = self.size();
182
183        let self_ref: MatZnx<&mut [u8]> = self.to_mut();
184        let nb_bytes: usize = VecZnx::<Vec<u8>>::alloc_bytes(n, cols_out, size);
185        let start: usize = nb_bytes * cols_in * row + col * nb_bytes;
186        let end: usize = start + nb_bytes;
187
188        VecZnx {
189            data: &mut self_ref.data[start..end],
190            n,
191            cols: cols_out,
192            size,
193            max_size: size,
194        }
195    }
196}
197
198impl<D: DataMut> FillUniform for MatZnx<D> {
199    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
200        match log_bound {
201            64 => source.fill_bytes(self.data.as_mut()),
202            0 => panic!("invalid log_bound, cannot be zero"),
203            _ => {
204                let mask: u64 = (1u64 << log_bound) - 1;
205                for x in self.raw_mut().iter_mut() {
206                    let r = source.next_u64() & mask;
207                    *x = ((r << (64 - log_bound)) as i64) >> (64 - log_bound);
208                }
209            }
210        }
211    }
212}
213
214impl<D: DataMut> Reset for MatZnx<D> {
215    fn reset(&mut self) {
216        self.zero();
217        self.n = 0;
218        self.size = 0;
219        self.rows = 0;
220        self.cols_in = 0;
221        self.cols_out = 0;
222    }
223}
224
225pub type MatZnxOwned = MatZnx<Vec<u8>>;
226pub type MatZnxMut<'a> = MatZnx<&'a mut [u8]>;
227pub type MatZnxRef<'a> = MatZnx<&'a [u8]>;
228
229pub trait MatZnxToRef {
230    fn to_ref(&self) -> MatZnx<&[u8]>;
231}
232
233impl<D: DataRef> MatZnxToRef for MatZnx<D> {
234    fn to_ref(&self) -> MatZnx<&[u8]> {
235        MatZnx {
236            data: self.data.as_ref(),
237            n: self.n,
238            rows: self.rows,
239            cols_in: self.cols_in,
240            cols_out: self.cols_out,
241            size: self.size,
242        }
243    }
244}
245
246pub trait MatZnxToMut {
247    fn to_mut(&mut self) -> MatZnx<&mut [u8]>;
248}
249
250impl<D: DataMut> MatZnxToMut for MatZnx<D> {
251    fn to_mut(&mut self) -> MatZnx<&mut [u8]> {
252        MatZnx {
253            data: self.data.as_mut(),
254            n: self.n,
255            rows: self.rows,
256            cols_in: self.cols_in,
257            cols_out: self.cols_out,
258            size: self.size,
259        }
260    }
261}
262
263impl<D: Data> MatZnx<D> {
264    pub fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
265        Self {
266            data,
267            n,
268            rows,
269            cols_in,
270            cols_out,
271            size,
272        }
273    }
274}
275
276impl<D: DataMut> ReaderFrom for MatZnx<D> {
277    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
278        self.n = reader.read_u64::<LittleEndian>()? as usize;
279        self.size = reader.read_u64::<LittleEndian>()? as usize;
280        self.rows = reader.read_u64::<LittleEndian>()? as usize;
281        self.cols_in = reader.read_u64::<LittleEndian>()? as usize;
282        self.cols_out = reader.read_u64::<LittleEndian>()? as usize;
283        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
284        let buf: &mut [u8] = self.data.as_mut();
285        if buf.len() != len {
286            return Err(std::io::Error::new(
287                std::io::ErrorKind::UnexpectedEof,
288                format!("self.data.len()={} != read len={}", buf.len(), len),
289            ));
290        }
291        reader.read_exact(&mut buf[..len])?;
292        Ok(())
293    }
294}
295
296impl<D: DataRef> WriterTo for MatZnx<D> {
297    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
298        writer.write_u64::<LittleEndian>(self.n as u64)?;
299        writer.write_u64::<LittleEndian>(self.size as u64)?;
300        writer.write_u64::<LittleEndian>(self.rows as u64)?;
301        writer.write_u64::<LittleEndian>(self.cols_in as u64)?;
302        writer.write_u64::<LittleEndian>(self.cols_out as u64)?;
303        let buf: &[u8] = self.data.as_ref();
304        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
305        writer.write_all(buf)?;
306        Ok(())
307    }
308}
309
310impl<D: DataRef> fmt::Display for MatZnx<D> {
311    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
312        writeln!(
313            f,
314            "MatZnx(n={}, rows={}, cols_in={}, cols_out={}, size={})",
315            self.n, self.rows, self.cols_in, self.cols_out, self.size
316        )?;
317
318        for row_i in 0..self.rows {
319            writeln!(f, "Row {}:", row_i)?;
320            for col_i in 0..self.cols_in {
321                writeln!(f, "cols_in {}:", col_i)?;
322                writeln!(f, "{}:", self.at(row_i, col_i))?;
323            }
324        }
325        Ok(())
326    }
327}
328
329impl<D: DataMut> ZnxZero for MatZnx<D> {
330    fn zero(&mut self) {
331        self.raw_mut().fill(0)
332    }
333
334    fn zero_at(&mut self, i: usize, j: usize) {
335        self.at_mut(i, j).zero();
336    }
337}