Skip to main content

strided_view/
auxiliary.rs

1//! Auxiliary routines ported from StridedViews.jl/src/auxiliary.jl
2
3/// Compute the relative order of strides.
4///
5/// Returns a vector where `result[i]` is the rank of `strides[i]` among all non-zero strides.
6/// Zero strides have order 1.
7///
8/// Returns an inline [`crate::AxisVec`]: the result does not allocate at rank
9/// `<=` [`crate::INLINE_RANK`] and spills above it.
10///
11/// # Examples
12///
13/// ```
14/// use strided_view::auxiliary::index_order;
15///
16/// let order = index_order(&[4, 1, 2]);
17/// assert_eq!(order.as_slice(), &[3, 1, 2]);
18/// ```
19///
20/// # Julia equivalent
21/// ```julia
22/// function indexorder(strides::NTuple{N,Int}) where {N}
23///     return ntuple(Val(N)) do i
24///         si = abs(strides[i])
25///         si == 0 && return 1
26///         k = 1
27///         for s in strides
28///             if s != 0 && abs(s) < si
29///                 k += 1
30///             end
31///         end
32///         return k
33///     end
34/// end
35/// ```
36pub fn index_order(strides: &[isize]) -> crate::AxisVec {
37    let n = strides.len();
38    let mut result = crate::AxisVec::new();
39    result.resize(n, 1usize);
40
41    for i in 0..n {
42        let si = strides[i].unsigned_abs();
43        if si == 0 {
44            result[i] = 1;
45            continue;
46        }
47        let mut k = 1usize;
48        for &s in strides {
49            if s != 0 && s.unsigned_abs() < si {
50                k += 1;
51            }
52        }
53        result[i] = k;
54    }
55
56    result
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    #[test]
64    fn test_index_order() {
65        // strides [4, 1, 2]: order is [3, 1, 2] (4 is largest, 1 is smallest, 2 is middle)
66        let order = index_order(&[4, 1, 2]);
67        assert_eq!(order.as_slice(), &[3, 1, 2]);
68    }
69
70    #[test]
71    fn test_index_order_with_zero() {
72        // Zero strides have order 1
73        let order = index_order(&[4, 0, 2]);
74        assert_eq!(order.as_slice(), &[2, 1, 1]);
75    }
76
77    #[test]
78    fn test_index_order_negative_strides() {
79        let order = index_order(&[-4, 1, -2]);
80        assert_eq!(order.as_slice(), &[3, 1, 2]);
81    }
82
83    #[test]
84    fn test_index_order_tied_strides() {
85        let order = index_order(&[2, 2, 1]);
86        assert_eq!(order.as_slice(), &[2, 2, 1]);
87    }
88
89    #[test]
90    fn test_index_order_all_same() {
91        let order = index_order(&[3, 3, 3]);
92        assert_eq!(order.as_slice(), &[1, 1, 1]);
93    }
94}