Skip to main content

conjure_cp_core/utils/
view.rs

1use crate::ast::{DomainPtr, GroundDomain, Moo};
2
3#[derive(Debug, Clone, PartialEq, Eq, Hash)]
4pub 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
15impl 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
26impl<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
39pub 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
49impl 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}