1use crate::{
  2    alloc_aligned,
  3    layouts::{
  4        Data, DataMut, DataRef, DataView, DataViewMut, DigestU64, FillUniform, ReaderFrom, 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 bytes_of(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
118        rows * cols_in * VecZnx::<Vec<u8>>::bytes_of(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::bytes_of(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::bytes_of(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>>::bytes_of(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>>::bytes_of(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
214pub type MatZnxOwned = MatZnx<Vec<u8>>;
215pub type MatZnxMut<'a> = MatZnx<&'a mut [u8]>;
216pub type MatZnxRef<'a> = MatZnx<&'a [u8]>;
217
218pub trait MatZnxToRef {
219    fn to_ref(&self) -> MatZnx<&[u8]>;
220}
221
222impl<D: DataRef> MatZnxToRef for MatZnx<D> {
223    fn to_ref(&self) -> MatZnx<&[u8]> {
224        MatZnx {
225            data: self.data.as_ref(),
226            n: self.n,
227            rows: self.rows,
228            cols_in: self.cols_in,
229            cols_out: self.cols_out,
230            size: self.size,
231        }
232    }
233}
234
235pub trait MatZnxToMut {
236    fn to_mut(&mut self) -> MatZnx<&mut [u8]>;
237}
238
239impl<D: DataMut> MatZnxToMut for MatZnx<D> {
240    fn to_mut(&mut self) -> MatZnx<&mut [u8]> {
241        MatZnx {
242            data: self.data.as_mut(),
243            n: self.n,
244            rows: self.rows,
245            cols_in: self.cols_in,
246            cols_out: self.cols_out,
247            size: self.size,
248        }
249    }
250}
251
252impl<D: Data> MatZnx<D> {
253    pub fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
254        Self {
255            data,
256            n,
257            rows,
258            cols_in,
259            cols_out,
260            size,
261        }
262    }
263}
264
265impl<D: DataMut> ReaderFrom for MatZnx<D> {
266    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
267        self.n = reader.read_u64::<LittleEndian>()? as usize;
268        self.size = reader.read_u64::<LittleEndian>()? as usize;
269        self.rows = reader.read_u64::<LittleEndian>()? as usize;
270        self.cols_in = reader.read_u64::<LittleEndian>()? as usize;
271        self.cols_out = reader.read_u64::<LittleEndian>()? as usize;
272        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
273        let buf: &mut [u8] = self.data.as_mut();
274        if buf.len() != len {
275            return Err(std::io::Error::new(
276                std::io::ErrorKind::UnexpectedEof,
277                format!("self.data.len()={} != read len={}", buf.len(), len),
278            ));
279        }
280        reader.read_exact(&mut buf[..len])?;
281        Ok(())
282    }
283}
284
285impl<D: DataRef> WriterTo for MatZnx<D> {
286    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
287        writer.write_u64::<LittleEndian>(self.n as u64)?;
288        writer.write_u64::<LittleEndian>(self.size as u64)?;
289        writer.write_u64::<LittleEndian>(self.rows as u64)?;
290        writer.write_u64::<LittleEndian>(self.cols_in as u64)?;
291        writer.write_u64::<LittleEndian>(self.cols_out as u64)?;
292        let buf: &[u8] = self.data.as_ref();
293        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
294        writer.write_all(buf)?;
295        Ok(())
296    }
297}
298
299impl<D: DataRef> fmt::Display for MatZnx<D> {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        writeln!(
302            f,
303            "MatZnx(n={}, rows={}, cols_in={}, cols_out={}, size={})",
304            self.n, self.rows, self.cols_in, self.cols_out, self.size
305        )?;
306
307        for row_i in 0..self.rows {
308            writeln!(f, "Row {row_i}:")?;
309            for col_i in 0..self.cols_in {
310                writeln!(f, "cols_in {col_i}:")?;
311                writeln!(f, "{}:", self.at(row_i, col_i))?;
312            }
313        }
314        Ok(())
315    }
316}
317
318impl<D: DataMut> ZnxZero for MatZnx<D> {
319    fn zero(&mut self) {
320        self.raw_mut().fill(0)
321    }
322
323    fn zero_at(&mut self, i: usize, j: usize) {
324        self.at_mut(i, j).zero();
325    }
326}