1use crate::{
  2    alloc_aligned,
  3    hal::{
  4        api::{DataView, DataViewMut, ZnxInfos, ZnxSliceSize, ZnxView},
  5        layouts::{Data, DataMut, DataRef, ReaderFrom, VecZnx, WriterTo},
  6    },
  7};
  8
  9#[derive(PartialEq, Eq)]
 10pub struct MatZnx<D: Data> {
 11    data: D,
 12    n: usize,
 13    size: usize,
 14    rows: usize,
 15    cols_in: usize,
 16    cols_out: usize,
 17}
 18
 19impl<D: Data> ZnxInfos for MatZnx<D> {
 20    fn cols(&self) -> usize {
 21        self.cols_in
 22    }
 23
 24    fn rows(&self) -> usize {
 25        self.rows
 26    }
 27
 28    fn n(&self) -> usize {
 29        self.n
 30    }
 31
 32    fn size(&self) -> usize {
 33        self.size
 34    }
 35}
 36
 37impl<D: Data> ZnxSliceSize for MatZnx<D> {
 38    fn sl(&self) -> usize {
 39        self.n() * self.cols_out()
 40    }
 41}
 42
 43impl<D: Data> DataView for MatZnx<D> {
 44    type D = D;
 45    fn data(&self) -> &Self::D {
 46        &self.data
 47    }
 48}
 49
 50impl<D: Data> DataViewMut for MatZnx<D> {
 51    fn data_mut(&mut self) -> &mut Self::D {
 52        &mut self.data
 53    }
 54}
 55
 56impl<D: DataRef> ZnxView for MatZnx<D> {
 57    type Scalar = i64;
 58}
 59
 60impl<D: Data> MatZnx<D> {
 61    pub fn cols_in(&self) -> usize {
 62        self.cols_in
 63    }
 64
 65    pub fn cols_out(&self) -> usize {
 66        self.cols_out
 67    }
 68}
 69
 70impl<D: DataRef> MatZnx<D> {
 71    pub fn bytes_of(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 72        rows * cols_in * VecZnx::<Vec<u8>>::alloc_bytes::<i64>(n, cols_out, size)
 73    }
 74}
 75
 76impl<D: DataRef + From<Vec<u8>>> MatZnx<D> {
 77    pub(crate) fn new(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 78        let data: Vec<u8> = alloc_aligned(Self::bytes_of(n, rows, cols_in, cols_out, size));
 79        Self {
 80            data: data.into(),
 81            n,
 82            size,
 83            rows,
 84            cols_in,
 85            cols_out,
 86        }
 87    }
 88
 89    pub(crate) fn new_from_bytes(
 90        n: usize,
 91        rows: usize,
 92        cols_in: usize,
 93        cols_out: usize,
 94        size: usize,
 95        bytes: impl Into<Vec<u8>>,
 96    ) -> Self {
 97        let data: Vec<u8> = bytes.into();
 98        assert!(data.len() == Self::bytes_of(n, rows, cols_in, cols_out, size));
 99        Self {
100            data: data.into(),
101            n,
102            size,
103            rows,
104            cols_in,
105            cols_out,
106        }
107    }
108}
109
110impl<D: DataRef> MatZnx<D> {
111    pub fn at(&self, row: usize, col: usize) -> VecZnx<&[u8]> {
112        #[cfg(debug_assertions)]
113        {
114            assert!(row < self.rows(), "rows: {} >= {}", row, self.rows());
115            assert!(col < self.cols_in(), "cols: {} >= {}", col, self.cols_in());
116        }
117
118        let self_ref: MatZnx<&[u8]> = self.to_ref();
119        let nb_bytes: usize = VecZnx::<Vec<u8>>::alloc_bytes::<i64>(self.n, self.cols_out, self.size);
120        let start: usize = nb_bytes * self.cols() * row + col * nb_bytes;
121        let end: usize = start + nb_bytes;
122
123        VecZnx {
124            data: &self_ref.data[start..end],
125            n: self.n,
126            cols: self.cols_out,
127            size: self.size,
128            max_size: self.size,
129        }
130    }
131}
132
133impl<D: DataMut> MatZnx<D> {
134    pub fn at_mut(&mut self, row: usize, col: usize) -> VecZnx<&mut [u8]> {
135        #[cfg(debug_assertions)]
136        {
137            assert!(row < self.rows(), "rows: {} >= {}", row, self.rows());
138            assert!(col < self.cols_in(), "cols: {} >= {}", col, self.cols_in());
139        }
140
141        let n: usize = self.n();
142        let cols_out: usize = self.cols_out();
143        let cols_in: usize = self.cols_in();
144        let size: usize = self.size();
145
146        let self_ref: MatZnx<&mut [u8]> = self.to_mut();
147        let nb_bytes: usize = VecZnx::<Vec<u8>>::alloc_bytes::<i64>(n, cols_out, size);
148        let start: usize = nb_bytes * cols_in * row + col * nb_bytes;
149        let end: usize = start + nb_bytes;
150
151        VecZnx {
152            data: &mut self_ref.data[start..end],
153            n,
154            cols: cols_out,
155            size,
156            max_size: size,
157        }
158    }
159}
160
161pub type MatZnxOwned = MatZnx<Vec<u8>>;
162pub type MatZnxMut<'a> = MatZnx<&'a mut [u8]>;
163pub type MatZnxRef<'a> = MatZnx<&'a [u8]>;
164
165pub trait MatZnxToRef {
166    fn to_ref(&self) -> MatZnx<&[u8]>;
167}
168
169impl<D: DataRef> MatZnxToRef for MatZnx<D> {
170    fn to_ref(&self) -> MatZnx<&[u8]> {
171        MatZnx {
172            data: self.data.as_ref(),
173            n: self.n,
174            rows: self.rows,
175            cols_in: self.cols_in,
176            cols_out: self.cols_out,
177            size: self.size,
178        }
179    }
180}
181
182pub trait MatZnxToMut {
183    fn to_mut(&mut self) -> MatZnx<&mut [u8]>;
184}
185
186impl<D: DataMut> MatZnxToMut for MatZnx<D> {
187    fn to_mut(&mut self) -> MatZnx<&mut [u8]> {
188        MatZnx {
189            data: self.data.as_mut(),
190            n: self.n,
191            rows: self.rows,
192            cols_in: self.cols_in,
193            cols_out: self.cols_out,
194            size: self.size,
195        }
196    }
197}
198
199impl<D: Data> MatZnx<D> {
200    pub(crate) fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
201        Self {
202            data,
203            n,
204            rows,
205            cols_in,
206            cols_out,
207            size,
208        }
209    }
210}
211
212use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
213
214impl<D: DataMut> ReaderFrom for MatZnx<D> {
215    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
216        self.n = reader.read_u64::<LittleEndian>()? as usize;
217        self.size = reader.read_u64::<LittleEndian>()? as usize;
218        self.rows = reader.read_u64::<LittleEndian>()? as usize;
219        self.cols_in = reader.read_u64::<LittleEndian>()? as usize;
220        self.cols_out = reader.read_u64::<LittleEndian>()? as usize;
221        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
222        let buf: &mut [u8] = self.data.as_mut();
223        if buf.len() != len {
224            return Err(std::io::Error::new(
225                std::io::ErrorKind::UnexpectedEof,
226                format!("self.data.len()={} != read len={}", buf.len(), len),
227            ));
228        }
229        reader.read_exact(&mut buf[..len])?;
230        Ok(())
231    }
232}
233
234impl<D: DataRef> WriterTo for MatZnx<D> {
235    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
236        writer.write_u64::<LittleEndian>(self.n as u64)?;
237        writer.write_u64::<LittleEndian>(self.size as u64)?;
238        writer.write_u64::<LittleEndian>(self.rows as u64)?;
239        writer.write_u64::<LittleEndian>(self.cols_in as u64)?;
240        writer.write_u64::<LittleEndian>(self.cols_out as u64)?;
241        let buf: &[u8] = self.data.as_ref();
242        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
243        writer.write_all(buf)?;
244        Ok(())
245    }
246}