1use super::error::*;
4use ndarray::*;
5
6pub use lax::layout::MatrixLayout;
7
8pub trait AllocatedArray {
9 type Elem;
10 fn layout(&self) -> Result<MatrixLayout>;
11 fn square_layout(&self) -> Result<MatrixLayout>;
12 fn ensure_square(&self) -> Result<()>;
14 fn as_allocated(&self) -> Result<&[Self::Elem]>;
15}
16
17pub trait AllocatedArrayMut: AllocatedArray {
18 fn as_allocated_mut(&mut self) -> Result<&mut [Self::Elem]>;
19}
20
21impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
22where
23 S: Data<Elem = A>,
24{
25 type Elem = A;
26
27 fn layout(&self) -> Result<MatrixLayout> {
28 let shape = self.shape();
29 let strides = self.strides();
30 if shape[0] == strides[1] as usize {
31 return Ok(MatrixLayout::F {
32 col: self.ncols() as i32,
33 lda: self.nrows() as i32,
34 });
35 }
36 if shape[1] == strides[0] as usize {
37 return Ok(MatrixLayout::C {
38 row: self.nrows() as i32,
39 lda: self.ncols() as i32,
40 });
41 }
42 Err(LinalgError::InvalidStride {
43 s0: strides[0],
44 s1: strides[1],
45 })
46 }
47
48 fn square_layout(&self) -> Result<MatrixLayout> {
49 let l = self.layout()?;
50 let (n, m) = l.size();
51 if n == m {
52 Ok(l)
53 } else {
54 Err(LinalgError::NotSquare { rows: n, cols: m })
55 }
56 }
57
58 fn ensure_square(&self) -> Result<()> {
59 if self.is_square() {
60 Ok(())
61 } else {
62 Err(LinalgError::NotSquare {
63 rows: self.nrows() as i32,
64 cols: self.ncols() as i32,
65 })
66 }
67 }
68
69 fn as_allocated(&self) -> Result<&[A]> {
70 self.as_slice_memory_order()
71 .ok_or(LinalgError::MemoryNotCont)
72 }
73}
74
75impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix2>
76where
77 S: DataMut<Elem = A>,
78{
79 fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
80 self.as_slice_memory_order_mut()
81 .ok_or(LinalgError::MemoryNotCont)
82 }
83}