use crate::error::MattenError;
use crate::limits::MAX_NDIM;
pub(crate) fn validate_shape(
shape: &[usize],
operation: &'static str,
) -> Result<usize, MattenError> {
if shape.len() > MAX_NDIM {
return Err(MattenError::Shape {
operation,
message: format!(
"rank {} exceeds the maximum supported rank of {MAX_NDIM} (shape {shape:?})",
shape.len()
),
});
}
checked_shape_len(shape, operation)
}
pub(crate) fn checked_shape_len(
shape: &[usize],
operation: &'static str,
) -> Result<usize, MattenError> {
let mut len: usize = 1;
for &dim in shape {
if dim == 0 {
return Err(MattenError::Shape {
operation,
message: format!(
"zero-sized dimensions are not supported in the current matten shape model (shape {shape:?})"
),
});
}
len = len.checked_mul(dim).ok_or_else(|| MattenError::Allocation {
requested_elements: usize::MAX,
message: format!("shape {shape:?} overflows usize when computing the element count in {operation}"),
})?;
}
Ok(len)
}
#[allow(dead_code)]
pub(crate) fn strides_for_shape(shape: &[usize]) -> Vec<usize> {
let mut strides = vec![0usize; shape.len()];
let mut acc: usize = 1;
for j in (0..shape.len()).rev() {
strides[j] = acc;
acc *= shape[j];
}
strides
}
#[allow(dead_code)]
pub(crate) fn coord_to_flat(coord: &[usize], shape: &[usize]) -> Option<usize> {
if coord.len() != shape.len() {
return None;
}
let strides = strides_for_shape(shape);
let mut flat = 0usize;
for (i, (&c, &dim)) in coord.iter().zip(shape).enumerate() {
if c >= dim {
return None;
}
flat += c * strides[i];
}
Some(flat)
}
#[allow(dead_code)]
pub(crate) fn flat_to_coord(flat: usize, shape: &[usize]) -> Vec<usize> {
let strides = strides_for_shape(shape);
let mut coord = vec![0usize; shape.len()];
let mut rem = flat;
for (i, &stride) in strides.iter().enumerate() {
coord[i] = rem / stride;
rem %= stride;
}
coord
}