//! Shape / stride helpers. Strides are in **elements** (not bytes), row-major.

use smallvec::SmallVec;

/// Inline storage for the common ≤4D case; spills to the heap beyond that.
pub type Dims = SmallVec<[usize; 4]>;

/// Row-major (C order) contiguous strides for a shape.
/// Example: `[2, 3, 4]` => `[12, 4, 1]`.
pub fn contiguous_strides(shape: &[usize]) -> Dims {
    let mut strides: Dims = SmallVec::with_capacity(shape.len());
    strides.resize(shape.len(), 0);
    if shape.is_empty() {
        return strides;
    }
    strides[shape.len() - 1] = 1;
    for i in (0..shape.len() - 1).rev() {
        strides[i] = strides[i + 1] * shape[i + 1];
    }
    strides
}

/// True if `strides` describe a row-major contiguous layout for `shape`.
/// A mismatched stride on a size-1 dimension is fine (it is never stepped).
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
    let ndim = shape.len();
    let mut expected = 1usize;
    for d in (0..ndim).rev() {
        if shape[d] != 1 && strides[d] != expected {
            return false;
        }
        expected *= shape[d];
    }
    true
}

/// Total element count.
pub fn num_elements(shape: &[usize]) -> usize {
    shape.iter().product()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn contiguous_strides_basic() {
        assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
        assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
    }

    #[test]
    fn is_contiguous_detects_transpose() {
        assert!(is_contiguous(&[2, 3], &[3, 1]));
        assert!(!is_contiguous(&[3, 2], &[1, 3]));
    }
}