nuvec_core/
shape.rs

1//! Shapes, strides and contiguity.
2
3use smallvec::SmallVec;
4
5use crate::{Error, Result};
6
7/// Dimension sizes. Inline for up to four dimensions.
8pub(crate) type Shape = SmallVec<[usize; 4]>;
9/// Per-dimension steps, in elements. Inline for up to four dimensions.
10pub(crate) type Strides = SmallVec<[isize; 4]>;
11
12/// How an array's elements are laid out in memory.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
14pub struct Flags {
15    /// Elements occupy one dense block, in some dimension order.
16    pub contiguous: bool,
17    /// Dense, last dimension varies fastest (C order).
18    pub row_contiguous: bool,
19    /// Dense, first dimension varies fastest (Fortran order).
20    pub col_contiguous: bool,
21}
22
23/// Element count of `shape`, or an error if it overflows.
24pub(crate) fn checked_size(shape: &[usize]) -> Result<usize> {
25    shape
26        .iter()
27        .try_fold(1usize, |acc, &d| acc.checked_mul(d))
28        .ok_or_else(|| {
29            Error::InvalidArgument(format!("element count of shape {shape:?} overflows"))
30        })
31}
32
33/// Byte size of `shape` with elements of `itemsize` bytes, or an error if it overflows.
34pub(crate) fn checked_nbytes(shape: &[usize], itemsize: usize) -> Result<usize> {
35    checked_size(shape)?
36        .checked_mul(itemsize)
37        .filter(|&n| n <= isize::MAX as usize)
38        .ok_or_else(|| Error::InvalidArgument(format!("byte size of shape {shape:?} overflows")))
39}
40
41/// Strides of a dense row-major array. Requires a shape whose size does not overflow.
42pub(crate) fn row_contiguous_strides(shape: &[usize]) -> Strides {
43    let mut strides: Strides = SmallVec::from_elem(0, shape.len());
44    let mut step = 1isize;
45    for (s, &d) in strides.iter_mut().zip(shape).rev() {
46        *s = step;
47        step = step.saturating_mul(d.max(1) as isize);
48    }
49    strides
50}
51
52/// Contiguity flags for a layout. Dimensions of size 1 do not affect contiguity.
53pub(crate) fn flags_for(shape: &[usize], strides: &[isize]) -> Flags {
54    let dense = |order: &mut dyn Iterator<Item = (&usize, &isize)>| {
55        let mut expected = 1isize;
56        for (&d, &s) in order {
57            if d != 1 && s != expected {
58                return false;
59            }
60            expected = expected.saturating_mul(d as isize);
61        }
62        true
63    };
64    let row = dense(&mut shape.iter().zip(strides).rev());
65    let col = dense(&mut shape.iter().zip(strides));
66    Flags {
67        contiguous: row || col,
68        row_contiguous: row,
69        col_contiguous: col,
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn row_major_strides() {
79        assert_eq!(row_contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
80        assert!(row_contiguous_strides(&[]).is_empty());
81    }
82
83    #[test]
84    fn size_and_overflow() {
85        assert_eq!(checked_size(&[]).unwrap(), 1);
86        assert_eq!(checked_size(&[2, 0, 5]).unwrap(), 0);
87        assert!(matches!(
88            checked_size(&[usize::MAX, 2]),
89            Err(Error::InvalidArgument(_))
90        ));
91    }
92
93    #[test]
94    fn flags() {
95        let f = flags_for(&[2, 3], &[3, 1]);
96        assert!(f.row_contiguous && f.contiguous && !f.col_contiguous);
97        let f = flags_for(&[2, 3], &[1, 2]);
98        assert!(f.col_contiguous && !f.row_contiguous);
99        let f = flags_for(&[1, 4], &[4, 1]);
100        assert!(f.row_contiguous && f.col_contiguous);
101        let scalar = flags_for(&[], &[]);
102        assert!(scalar.row_contiguous && scalar.col_contiguous);
103    }
104}
nuvec_core/shape.rs

nuvec_core/
shape.rs
