use std::fmt;
use smallvec::SmallVec;
use static_assertions::assert_impl_all;
assert_impl_all!(Shape: Send, Sync);
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub struct Shape(SmallVec<[usize; 4]>);
impl Shape {
pub fn new(axes: impl IntoIterator<Item = usize>) -> Self {
Self(axes.into_iter().collect())
}
pub fn scalar() -> Self {
Self(SmallVec::new())
}
pub fn rank(&self) -> usize {
self.0.len()
}
pub fn volume(&self) -> usize {
self.0
.iter()
.try_fold(1usize, |volume, &axis| volume.checked_mul(axis))
.expect("shape volume overflows `usize`")
}
pub fn axes(&self) -> &[usize] {
&self.0
}
pub fn without_axis(&self, axis: usize) -> Shape {
assert!(axis < self.rank(), "axis {axis} is out of rank for {self}");
Shape(
self.0
.iter()
.enumerate()
.filter(|(index, _)| *index != axis)
.map(|(_, &extent)| extent)
.collect(),
)
}
}
impl<const RANK: usize> From<[usize; RANK]> for Shape {
fn from(axes: [usize; RANK]) -> Self {
Shape::new(axes)
}
}
impl From<Vec<usize>> for Shape {
fn from(axes: Vec<usize>) -> Self {
Shape::new(axes)
}
}
impl From<&[usize]> for Shape {
fn from(axes: &[usize]) -> Self {
Shape::new(axes.iter().copied())
}
}
impl From<&Shape> for Shape {
fn from(shape: &Shape) -> Self {
shape.clone()
}
}
impl fmt::Display for Shape {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "[")?;
for (index, axis) in self.0.iter().enumerate() {
if index > 0 {
write!(formatter, ", ")?;
}
write!(formatter, "{axis}")?;
}
write!(formatter, "]")
}
}
#[cfg(test)]
#[path = "tests/shape_tests.rs"]
mod tests;