nd_slice/
addressing.rs

1//! Indexing into n-dimensional arrays with different underlying storage mechanisms.
2//!
3//! <https://en.wikipedia.org/wiki/Row-_and_column-major_order#Address_calculation_in_general>
4
5#[derive(Debug, Clone)]
6/// Enum to indicate the underlying memory order of the `NdSlice` or `NdSliceMut`.
7///
8/// In row-major order, the last dimension is contiguous.
9/// In column-major order, the first dimension is contiguous. 
10/// Read more: <https://en.wikipedia.org/wiki/Row-_and_column-major_order>
11pub enum Order {
12    /// In row-major order, the last dimension is contiguous. Read more: <https://en.wikipedia.org/wiki/Row-_and_column-major_order>
13    RowMajor,
14    /// In column-major order, the first dimension is contiguous. Read more: <https://en.wikipedia.org/wiki/Row-_and_column-major_order>
15    ColumnMajor,
16}
17
18/// Given the shape and index calculate the "address" for a given memory order.
19pub fn address(order: &Order, shape: &[usize], index: &[usize]) -> usize {
20    match order {
21        Order::RowMajor => row_major_address(shape, index),
22        Order::ColumnMajor => col_major_address(shape, index),
23    }
24
25}
26
27fn row_major_address(shape: &[usize], index: &[usize]) -> usize {
28    let d = {
29        assert_eq!(shape.len(), index.len());
30        shape.len()
31    };
32
33    let mut res = index[0];
34    for i in 1..d {
35        res = res * shape[i] + index[i];
36    }
37
38    res
39}
40
41fn col_major_address(shape: &[usize], index: &[usize]) -> usize {
42    let d = {
43        assert_eq!(shape.len(), index.len());
44        shape.len()
45    };
46
47    let mut res = index[d - 1];
48    for i in (0..(d - 1)).rev() {
49        res = res * shape[i] + index[i];
50    }
51
52    res
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[test]
60    fn row_major_test_2d() {
61        assert_eq!(row_major_address(&[2, 3], &[0, 0]), 0);
62        assert_eq!(row_major_address(&[2, 3], &[0, 1]), 1);
63        assert_eq!(row_major_address(&[2, 3], &[0, 2]), 2);
64        assert_eq!(row_major_address(&[2, 3], &[1, 0]), 3);
65        assert_eq!(row_major_address(&[2, 3], &[1, 1]), 4);
66        assert_eq!(row_major_address(&[2, 3], &[1, 2]), 5);
67    }
68
69    #[test]
70    fn col_major_test_2d() {
71        assert_eq!(col_major_address(&[2, 3], &[0, 0]), 0);
72        assert_eq!(col_major_address(&[2, 3], &[1, 0]), 1);
73        assert_eq!(col_major_address(&[2, 3], &[0, 1]), 2);
74        assert_eq!(col_major_address(&[2, 3], &[1, 1]), 3);
75        assert_eq!(col_major_address(&[2, 3], &[0, 2]), 4);
76        assert_eq!(col_major_address(&[2, 3], &[1, 2]), 5);
77    }
78}
79
nd_slice/addressing.rs

nd_slice/
addressing.rs
