1
use crate::ast::{DomainPtr, GroundDomain, Moo};
2

            
3
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
4
pub struct MatrixShape<T> {
5
    // Total count of elements
6
    pub size: usize,
7
    // Sizes along each dimension
8
    pub dims: Vec<usize>,
9
    // Strides for each dimension
10
    pub strides: Vec<usize>,
11
    // Index domains for each dimension
12
    pub idx_doms: Vec<T>,
13
}
14

            
15
impl From<MatrixShape<Moo<GroundDomain>>> for MatrixShape<DomainPtr> {
16
    fn from(value: MatrixShape<Moo<GroundDomain>>) -> Self {
17
        MatrixShape {
18
            size: value.size,
19
            strides: value.strides,
20
            dims: value.dims,
21
            idx_doms: value.idx_doms.into_iter().map(Into::into).collect(),
22
        }
23
    }
24
}
25

            
26
impl<T> From<MatrixShape<T>> for View {
27
    fn from(value: MatrixShape<T>) -> Self {
28
        Self {
29
            offset: 0,
30
            dims: value.dims,
31
            strides: value.strides,
32
        }
33
    }
34
}
35

            
36
#[allow(dead_code)] // will be used in following commits
37
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
38
/// A view into a 1D slice that we can manipulate
39
pub struct View {
40
    // Number of elements to skip at the beginning
41
    pub offset: usize,
42
    // Sizes along each dimension
43
    pub dims: Vec<usize>,
44
    // Strides for each dimension
45
    pub strides: Vec<usize>,
46
}
47

            
48
#[allow(dead_code)] // will be used in following commits
49
impl View {
50
    pub fn new(offset: usize, dims: Vec<usize>, strides: Vec<usize>) -> Self {
51
        Self {
52
            offset,
53
            dims,
54
            strides,
55
        }
56
    }
57

            
58
    fn apply_impl<'a, T>(
59
        elems: &'a [T],
60
        offset: usize,
61
        dims: &[usize],
62
        strides: &[usize],
63
    ) -> Vec<&'a T> {
64
        assert_eq!(dims.len(), strides.len());
65

            
66
        if dims.is_empty() {
67
            return vec![&elems[offset]];
68
        }
69

            
70
        let mut ans: Vec<&'a T> = Vec::new();
71
        for i in 0..dims[0] {
72
            let new_off = offset + i * strides[0];
73
            ans.extend(View::apply_impl(elems, new_off, &dims[1..], &strides[1..]));
74
        }
75
        ans
76
    }
77

            
78
    /// Get this view into `elems`
79
    pub fn apply<'a, T>(&self, elems: &'a [T]) -> Vec<&'a T> {
80
        View::apply_impl(elems, self.offset, &self.dims, &self.strides)
81
    }
82

            
83
    /// Reorder dimensions according to the given permutation.
84
    ///
85
    /// `perm` must be a permutation of `0..self.dims.len()`. The returned view
86
    /// iterates through the same backing store but with reordered dimensions:
87
    /// dimension `i` of the new view corresponds to dimension `perm[i]` of `self`.
88
    pub fn permute(&self, perm: &[usize]) -> View {
89
        debug_assert_eq!(perm.len(), self.dims.len());
90
        View::new(
91
            self.offset,
92
            perm.iter().map(|&i| self.dims[i]).collect(),
93
            perm.iter().map(|&i| self.strides[i]).collect(),
94
        )
95
    }
96

            
97
    /// Compute standard row-major strides for the given dimension sizes.
98
    ///
99
    /// For dims `[d0, d1, .., dN]` the strides are
100
    /// `[d1*d2*..*dN, d2*..*dN, .., 1]`.
101
    pub fn row_major_strides(dims: &[usize]) -> Vec<usize> {
102
        let mut strides = vec![1usize; dims.len()];
103
        for i in (0..dims.len().saturating_sub(1)).rev() {
104
            strides[i] = strides[i + 1] * dims[i + 1];
105
        }
106
        strides
107
    }
108
}