use crate::error::{Error, Kind};
use crate::hashing::{HashMap, HashMapExt, HashSet, HashSetExt};
use std::{borrow::Cow, fmt, sync::Arc};
pub const MAX_RANK: usize = 128;
pub const MAX_GENERATED_ELEMENTS: usize = 1_000_000_000;
#[inline]
pub fn count(shape: &[usize]) -> usize { match *shape { [] => 1, [n] => n, [m, n] => m * n, _ => shape.iter().product() } }
#[inline]
pub fn element_count(shape: &[usize]) -> Result<usize, Error> {
if shape.len() > MAX_RANK { return Err(Error::new(Kind::Limit, format!("an array has at most {MAX_RANK} axes"))); }
shape.iter().try_fold(1usize, |n, &d| n.checked_mul(d)).or_else(|| shape.contains(&0).then_some(0))
.ok_or_else(|| Error::new(Kind::Limit, "an array of this shape has too many items to count"))
}
#[inline]
pub fn generated_len(shape: &[usize]) -> Result<usize, Error> {
let len = element_count(shape)?;
if len > MAX_GENERATED_ELEMENTS { return Err(Error::new(Kind::Limit, format!("an array has at most {MAX_GENERATED_ELEMENTS} items"))); }
crate::memory::admit(len)?;
Ok(len)
}
#[inline]
pub fn strides(shape: &[usize]) -> Vec<usize> {
let mut strides = vec![1; shape.len()];
for a in (1..shape.len()).rev() { strides[a - 1] = strides[a] * shape[a]; }
strides
}
pub fn exclusive_sums(sizes: impl IntoIterator<Item = usize>) -> Vec<usize> {
let mut at = 0;
sizes.into_iter().map(|n| {
let start = at;
at += n;
start
}).collect()
}
#[inline]
pub fn coordinates(mut flat: usize, shape: &[usize]) -> impl Iterator<Item = (usize, usize)> + '_ {
shape.iter().enumerate().rev().map(move |(axis, &len)| {
let coordinate = flat % len;
flat /= len;
(axis, coordinate)
})
}
#[inline]
pub fn position(n: i64, len: usize) -> Option<usize> { let i = if n < 0 { len as i64 + n } else { n }; usize::try_from(i).ok().filter(|&i| i < len) }
#[derive(Clone, Copy, Debug)]
pub struct Lanes { pub outer: usize, pub len: usize, pub inner: usize }
impl Lanes {
#[inline]
pub fn new(shape: &[usize], axis: usize) -> Result<Self, Error> {
let &len = shape.get(axis).ok_or_else(|| Error::new(Kind::Rank, format!("axis {axis} is outside an array of rank {}", shape.len())))?;
Ok(Self { outer: element_count(&shape[..axis])?, len, inner: element_count(&shape[axis + 1..])? })
}
#[inline]
pub fn offset(&self, i: usize, j: usize, k: usize) -> usize { (i * self.len + j) * self.inner + k }
#[inline]
pub fn dims(&self) -> [usize; 3] { [self.outer, self.len, self.inner] }
}
pub fn agreed<'a, T: PartialEq + Clone + 'a>(values: impl IntoIterator<Item = Option<&'a T>>) -> Option<T> {
let mut found = None;
for value in values.into_iter().flatten() {
if found.is_some_and(|f| f != value) { return None; }
found = Some(value);
}
found.cloned()
}
pub fn repeated_name(names: &[Option<Arc<str>>]) -> Option<&Arc<str>> { let mut seen = HashSet::new(); names.iter().flatten().find(|n| !seen.insert(*n)) }
#[inline]
pub fn hidden(name: &str) -> bool { name.starts_with('_') }
pub const FRESH: char = '·';
pub fn shown(name: &str) -> Cow<'_, str> {
match name.rfind(FRESH) { Some(i) if hidden(name) => format!("{}…", &name[..i + FRESH.len_utf8()]).into(), _ => name.into() }
}
#[inline]
pub fn names_conflict(a: Option<&str>, b: Option<&str>) -> bool { matches!((a, b), (Some(a), Some(b)) if a != b && hidden(a) == hidden(b)) }
#[derive(Clone)]
pub struct Shape(Repr);
#[derive(Clone)]
enum Repr { Small(u8, [usize; 2]), Large(Box<[usize]>) }
impl Shape {
pub const EMPTY: Self = Self(Repr::Small(0, [0; 2]));
fn spliced(parts: [&[usize]; 3]) -> Self {
let len: usize = parts.iter().map(|p| p.len()).sum();
if len > 2 { return Self(Repr::Large(parts.concat().into_boxed_slice())); }
let mut small = [0; 2];
for (slot, &n) in small.iter_mut().zip(parts.into_iter().flatten()) { *slot = n; }
Self(Repr::Small(len as u8, small))
}
}
impl Default for Shape {
#[inline]
fn default() -> Self { Self::EMPTY }
}
impl std::ops::Deref for Shape {
type Target = [usize];
#[inline]
fn deref(&self) -> &[usize] { match &self.0 { Repr::Small(n, lengths) => &lengths[..usize::from(*n)], Repr::Large(lengths) => &lengths[..] } }
}
impl std::ops::DerefMut for Shape {
#[inline]
fn deref_mut(&mut self) -> &mut [usize] {
match &mut self.0 { Repr::Small(n, lengths) => &mut lengths[..usize::from(*n)], Repr::Large(lengths) => &mut lengths[..] }
}
}
impl PartialEq for Shape {
#[inline]
fn eq(&self, other: &Self) -> bool { **self == **other }
}
impl Eq for Shape {}
impl fmt::Debug for Shape { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } }
impl From<&[usize]> for Shape {
#[inline]
fn from(lengths: &[usize]) -> Self {
match *lengths {
[] => Self::EMPTY, [a] => Self(Repr::Small(1, [a, 0])), [a, b] => Self(Repr::Small(2, [a, b])), _ => Self(Repr::Large(lengths.into()))
}
}
}
impl From<Vec<usize>> for Shape {
#[inline]
fn from(lengths: Vec<usize>) -> Self { if lengths.len() > 2 { Self(Repr::Large(lengths.into_boxed_slice())) } else { Self::from(lengths.as_slice()) } }
}
impl<const N: usize> From<[usize; N]> for Shape { fn from(lengths: [usize; N]) -> Self { Self::from(lengths.as_slice()) } }
impl FromIterator<usize> for Shape {
fn from_iter<I: IntoIterator<Item = usize>>(lengths: I) -> Self {
let mut lengths = lengths.into_iter();
let mut small = [0; 2];
for n in 0..2 {
let Some(len) = lengths.next() else { return Self(Repr::Small(n, small)) };
small[usize::from(n)] = len;
}
match lengths.next() {
None => Self(Repr::Small(2, small)), Some(third) => Self(Repr::Large(small.into_iter().chain([third]).chain(lengths).collect()))
}
}
}
#[derive(Clone, Debug, Default)]
pub struct Keys { names: Vec<Option<Arc<str>>>, index: HashMap<Arc<str>, usize> }
impl PartialEq for Keys { fn eq(&self, other: &Self) -> bool { self.names == other.names } }
impl Eq for Keys {}
#[derive(Clone, Debug)]
pub enum Alignment {
Same,
Leading,
Table(Vec<Option<usize>>),
}
impl Keys {
pub fn partial(names: Vec<Option<Arc<str>>>) -> Result<Arc<Self>, Error> {
let mut index = HashMap::with_capacity(names.len());
for (i, name) in names.iter().enumerate() {
let Some(name) = name else { continue };
if index.insert(name.clone(), i).is_some() { return Err(duplicate(name)); }
}
Ok(Arc::new(Self { names, index }))
}
#[inline]
pub fn names(&self) -> &[Option<Arc<str>>] { &self.names }
#[inline]
pub fn len(&self) -> usize { self.names.len() }
#[inline]
pub fn is_empty(&self) -> bool { self.names.is_empty() }
#[inline]
pub fn complete(&self) -> bool { self.index.len() == self.names.len() }
#[inline]
pub fn blank(&self) -> bool { self.index.is_empty() && !self.is_empty() }
#[inline]
pub fn position(&self, name: &str) -> Option<usize> { self.index.get(name).copied() }
pub fn extend(&mut self, names: &[Option<Arc<str>>]) -> Result<(), Error> {
let mut new = HashSet::new();
if let Some(name) = names.iter().flatten().find(|n| self.position(n).is_some() || !new.insert(*n)) { return Err(duplicate(name)); }
for name in names {
if let Some(name) = name { self.index.insert(name.clone(), self.names.len()); }
self.names.push(name.clone());
}
Ok(())
}
pub fn align(&self, wanted: &Self) -> Vec<Option<usize>> {
let mut gaps = self.names.iter().enumerate().filter(|(_, n)| n.is_none()).map(|(i, _)| i);
wanted.names.iter().map(|n| match n { Some(n) => self.position(n), None => gaps.next() }).collect()
}
pub fn positions(&self, wanted: &Self, subset: bool) -> Result<Vec<usize>, Error> {
let kind = if subset { Kind::Index } else { Kind::Length };
if !subset && self.len() != wanted.len() {
return Err(Error::new(kind, format!("keyed axes of lengths {} and {} don't pair", self.len(), wanted.len())));
}
self.align(wanted).into_iter().zip(&wanted.names).map(|(p, name)| {
p.ok_or_else(|| match name {
Some(name) => Error::new(kind, format!("missing key: {}", shown(name))), None => Error::new(kind, "too few positions without keys")
})
}).collect()
}
pub fn select(&self, positions: impl IntoIterator<Item = Option<usize>>) -> Result<Arc<Self>, Error> {
Self::partial(positions.into_iter().map(|i| i.and_then(|i| self.names[i].clone())).collect())
}
pub fn aligned(self: &Arc<Self>, wanted: &Arc<Self>) -> Alignment {
if Arc::ptr_eq(self, wanted) || self == wanted { Alignment::Same }
else if wanted.names.starts_with(&self.names) { Alignment::Leading }
else { Alignment::Table(self.align(wanted)) }
}
pub fn union(self: &Arc<Self>, other: &Arc<Self>) -> Result<(Arc<Self>, [Alignment; 2]), Error> {
if Arc::ptr_eq(self, other) { return Ok((self.clone(), [Alignment::Same, Alignment::Same])); }
let (a, b) = (self.len() - self.index.len(), other.len() - other.index.len());
if a != b { return Err(Error::new(Kind::Length, format!("axes with {a} and {b} positions without keys don't align"))); }
if self == other { return Ok((self.clone(), [Alignment::Same, Alignment::Same])); }
if other.names.starts_with(&self.names) { return Ok((other.clone(), [Alignment::Leading, Alignment::Same])); }
if self.names.starts_with(&other.names) { return Ok((self.clone(), [Alignment::Same, Alignment::Leading])); }
let places = self.align(other);
let mut table = vec![None; self.len()];
for (j, &p) in places.iter().enumerate() { match p { Some(p) => table[p] = Some(j), None => table.push(Some(j)) } }
if table.len() == self.len() { return Ok((self.clone(), [Alignment::Same, Alignment::Table(table)])); }
let keys = Self::partial(self.names.iter().chain(other.names.iter().zip(&places).filter(|(_, p)| p.is_none()).map(|(k, _)| k)).cloned().collect())?;
Ok((keys, [Alignment::Leading, Alignment::Table(table)]))
}
}
fn duplicate(key: &str) -> Error { Error::new(Kind::Domain, format!("duplicate key: {}", shown(key))) }
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Layout { shape: Shape, meta: Option<Box<Meta>> }
#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct Meta { keys: Vec<Option<Arc<Keys>>>, names: Vec<Option<Arc<str>>> }
impl<S: Into<Shape>> From<S> for Layout {
#[inline]
fn from(shape: S) -> Self { Self { shape: shape.into(), meta: None } }
}
impl Layout {
pub const EMPTY: Self = Self { shape: Shape::EMPTY, meta: None };
#[inline]
pub fn shape(&self) -> &[usize] { &self.shape }
#[inline]
pub fn count(&self) -> usize { count(&self.shape) }
#[inline]
pub fn names(&self) -> &[Option<Arc<str>>] { self.meta.as_ref().map_or(&[], |l| &l.names) }
#[inline]
pub fn name(&self, axis: usize) -> Option<&Arc<str>> { self.names().get(axis).and_then(Option::as_ref) }
#[inline]
pub fn axis_named(&self, name: &str) -> Option<usize> { self.names().iter().position(|n| n.as_deref() == Some(name)) }
fn meta_mut(&mut self) -> &mut Meta { self.meta.get_or_insert_with(Default::default) }
fn tidy(mut self) -> Self { if self.meta.as_deref() == Some(&Meta::default()) { self.meta = None; } self }
fn edit(mut self, empty: bool, set: impl FnOnce(&mut Meta)) -> Self {
if empty && self.meta.is_none() { return self; }
set(self.meta_mut());
self.tidy()
}
pub fn with_names(self, names: Vec<Option<Arc<str>>>) -> Result<Self, Error> {
if !names.is_empty() && names.len() != self.shape.len() {
return Err(Error::new(Kind::Rank, format!("{} names for {} axes", names.len(), self.shape.len())));
}
if let Some(name) = repeated_name(&names) { return Err(Error::new(Kind::Domain, format!("the result would have two axes named {}", shown(name)))); }
Ok(self.inherit_names(names))
}
pub fn inherit_names(self, names: Vec<Option<Arc<str>>>) -> Self {
assert!(names.is_empty() || names.len() == self.shape.len());
let empty = names.iter().all(Option::is_none);
self.edit(empty, |m| m.names = if empty { vec![] } else { names })
}
pub fn unique_names(mut self) -> Self {
let Some(meta) = self.meta.as_mut() else { return self };
let mut counts = HashMap::new();
for name in meta.names.iter().flatten() { *counts.entry(name.clone()).or_insert(0) += 1; }
for name in &mut meta.names { if name.as_ref().is_some_and(|n| counts[n] > 1) { *name = None; } }
if meta.names.iter().all(Option::is_none) { meta.names.clear(); }
self.tidy()
}
#[inline]
fn key_list(&self) -> &[Option<Arc<Keys>>] { self.meta.as_ref().map_or(&[], |l| &l.keys) }
#[inline]
pub fn has_keys(&self) -> bool { !self.key_list().is_empty() }
#[inline]
pub fn has_names(&self) -> bool { !self.names().is_empty() }
#[inline]
pub fn plain(&self) -> bool { self.meta.is_none() }
#[inline]
pub fn keys(&self, axis: usize) -> Option<&Arc<Keys>> { self.key_list().get(axis).and_then(Option::as_ref) }
pub fn all_keys(&self) -> Vec<Option<Arc<Keys>>> { (0..self.shape.len()).map(|a| self.keys(a).cloned()).collect() }
pub fn matches(&self, other: &Self) -> bool {
self.shape == other.shape && (0..self.shape.len()).all(|a| {
!names_conflict(self.name(a).map(AsRef::as_ref), other.name(a).map(AsRef::as_ref)) && match (self.keys(a), other.keys(a)) {
(Some(x), Some(y)) => !y.align(x).contains(&None), _ => true
}
})
}
pub fn without_keys(self, axis: usize) -> Self {
if self.keys(axis).is_none() { return self; }
let mut keys = self.all_keys();
keys[axis] = None;
self.with_keys(keys).unwrap()
}
pub fn with_keys(self, keys: Vec<Option<Arc<Keys>>>) -> Result<Self, Error> {
if !keys.is_empty() && keys.len() != self.shape.len() {
return Err(Error::new(Kind::Rank, format!("keys for {} axes of an array with {} axes", keys.len(), self.shape.len())));
}
if let Some((k, n)) = keys.iter().zip(self.shape.iter()).find_map(|(k, &n)| k.as_ref().filter(|k| k.len() != n).map(|k| (k.len(), n))) {
return Err(Error::new(Kind::Length, format!("{k} keys for an axis of length {n}")));
}
let keys: Vec<_> = keys.into_iter().map(|k| k.filter(|k| !k.blank())).collect();
let empty = keys.iter().all(Option::is_none);
Ok(self.edit(empty, |m| m.keys = if empty { vec![] } else { keys }))
}
fn gathered<'a>(sources: impl Iterator<Item = (&'a Self, usize)> + Clone) -> Self {
let shape: Shape = sources.clone().map(|(l, a)| l.shape[a]).collect();
if sources.clone().all(|(l, _)| l.plain()) { return Self::from(shape); }
let keys = sources.clone().map(|(l, a)| l.keys(a).cloned()).collect();
let names = sources.map(|(l, a)| l.name(a).cloned()).collect();
Self::from(shape).with_keys(keys).unwrap().inherit_names(names)
}
pub fn axes(&self, axes: impl Iterator<Item = usize> + Clone) -> Self { Self::gathered(axes.map(|a| (self, a))) }
pub fn concat(&self, other: &Self) -> Self { Self::gathered((0..self.shape.len()).map(|a| (self, a)).chain((0..other.shape.len()).map(|a| (other, a)))) }
pub fn contracted(&self, right: &Self) -> Self { self.axes(0..self.shape.len().saturating_sub(1)).concat(&right.axes(1..right.shape.len())) }
pub fn merged(&self, axes: std::ops::Range<usize>, len: usize) -> Self {
let shape = Shape::spliced([&self.shape[..axes.start], &[len], &self.shape[axes.end..]]);
if self.plain() { return Self::from(shape); }
let source = |a: usize| match a.cmp(&axes.start) {
std::cmp::Ordering::Less => Some(a), std::cmp::Ordering::Equal => None, std::cmp::Ordering::Greater => Some(a - 1 + axes.len())
};
let keys = (0..shape.len()).map(|a| source(a).and_then(|s| self.keys(s).cloned())).collect();
let names = (0..shape.len()).map(|a| source(a).and_then(|s| self.name(s).cloned())).collect();
Self::from(shape).with_keys(keys).unwrap().inherit_names(names)
}
pub fn select(&self, axis: usize, positions: impl ExactSizeIterator<Item = Option<usize>>) -> Result<Self, Error> {
let mut layout = self.clone();
layout.shape[axis] = positions.len();
let Some(keys) = self.keys(axis) else { return Ok(layout) };
let mut all = self.all_keys();
all[axis] = keys.select(positions).ok();
layout.with_keys(all)
}
pub fn append(&mut self, len: usize, names: Option<&[Option<Arc<str>>]>) -> Result<(), Error> {
let (old, rank) = (self.shape[0], self.shape.len());
if let Some(names) = names {
match self.meta.as_mut().and_then(|m| m.keys.first_mut()).and_then(Option::as_mut) {
Some(keys) => Arc::make_mut(keys).extend(names)?,
None => {
let keys = Keys::partial(std::iter::repeat_n(None, old).chain(names.iter().cloned()).collect())?;
let meta = self.meta_mut();
if meta.keys.is_empty() { meta.keys = vec![None; rank]; }
meta.keys[0] = Some(keys);
}
}
}
self.shape[0] += len;
Ok(())
}
}