use std::fmt;
use std::ops::{Index, IndexMut};
#[cfg(feature = "afserde")]
use serde::{Deserialize, Serialize};
#[derive(Copy, Clone, PartialEq, Debug)]
#[cfg_attr(feature = "afserde", derive(Serialize, Deserialize))]
pub struct Dim4 {
dims: [u64; 4],
}
impl Default for Dim4 {
fn default() -> Self {
Self { dims: [1, 1, 1, 1] }
}
}
impl Index<usize> for Dim4 {
type Output = u64;
fn index(&self, _index: usize) -> &u64 {
&self.dims[_index]
}
}
impl IndexMut<usize> for Dim4 {
fn index_mut(&mut self, _index: usize) -> &mut Self::Output {
&mut self.dims[_index]
}
}
impl fmt::Display for Dim4 {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"[{} {} {} {}]",
self.dims[0], self.dims[1], self.dims[2], self.dims[3]
)
}
}
impl Dim4 {
pub fn new(dims: &[u64; 4]) -> Self {
Self {
dims: [dims[0], dims[1], dims[2], dims[3]],
}
}
pub fn elements(&self) -> u64 {
self.dims[0] * self.dims[1] * self.dims[2] * self.dims[3]
}
pub fn ndims(&self) -> usize {
let nelems = self.elements();
match nelems {
0 => 0,
1 => 1,
_ => {
if self.dims[3] != 1 {
4
} else if self.dims[2] != 1 {
3
} else if self.dims[1] != 1 {
2
} else {
1
}
}
}
}
pub fn get(&self) -> &[u64; 4] {
&self.dims
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "afserde")]
mod serde_tests {
use super::super::Dim4;
use crate::dim4;
#[test]
fn dim4_serde() {
let dims = dim4!(4, 4);
let serd = match serde_json::to_string(&dims) {
Ok(serialized_str) => serialized_str,
Err(e) => e.to_string(),
};
assert_eq!(serd, "{\"dims\":[4,4,1,1]}");
let deserd: Dim4 = serde_json::from_str(&serd).unwrap();
assert_eq!(deserd, dims);
}
}
}