use crate::arr::core::{
adaptative_result_shape, dim_expansion_for_shape, next_index, plain_index, shape_size,
};
use std::vec;
#[test]
fn shape_size_works() {
let cases = [
(vec![2, 6, 8], 2 * 6 * 8),
(vec![1], 1),
(vec![3, 2, 6, 8], 3 * 2 * 6 * 8),
];
for (shape, expected_size) in cases.iter() {
assert_eq!(*expected_size, shape_size(shape));
}
}
#[test]
fn plain_index_works() {
let cases = [
(vec![2, 6, 1], vec![2, 6, 8], vec![0, 1, 3], 1 * 1),
(
vec![2, 1, 8],
vec![2, 6, 8],
vec![1, 3, 5],
1 * (1 * 8) + 0 * 8 + 5,
),
];
for (shape, index_shape, index, expected_plain_index) in cases.iter() {
assert_eq!(
*expected_plain_index,
plain_index(shape, index_shape, index)
);
}
}
#[test]
fn next_index_works() {
let shape = vec![2, 1, 2];
let mut index = vec![0, 0, 0];
let end = next_index(&shape, &mut index);
assert_eq!(vec![0, 0, 1], index);
assert!(!end);
let end = next_index(&shape, &mut index);
assert_eq!(vec![1, 0, 0], index);
assert!(!end);
let end = next_index(&shape, &mut index);
assert_eq!(vec![1, 0, 1], index);
assert!(!end);
let end = next_index(&shape, &mut index);
assert_eq!(vec![0, 0, 0], index);
assert!(end);
}
#[test]
fn commutative_result_shape_works() {
let cases = [
(vec![2, 6, 8], vec![1], vec![2, 6, 8]),
(vec![1], vec![3, 5, 8], vec![3, 5, 8]),
(vec![1, 1], vec![3, 5, 8], vec![3, 5, 8]),
(vec![1, 8], vec![3, 5, 8], vec![3, 5, 8]),
(vec![10, 10], vec![10], vec![10, 10]),
];
for (shape_a, shape_b, expected_result_shape) in cases.iter() {
assert_eq!(
*expected_result_shape,
adaptative_result_shape(shape_a, shape_b)
);
}
}
#[test]
fn dim_expansion_for_shape_works() {
assert_eq!(vec![1, 1, 5], dim_expansion_for_shape(&vec![5], 3));
assert_eq!(vec![1, 1, 3, 5], dim_expansion_for_shape(&vec![3, 5], 4));
assert_eq!(
vec![8, 1, 3, 5],
dim_expansion_for_shape(&vec![8, 1, 3, 5], 4)
);
}