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