use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ops::{Index, IndexMut};
use tinyvec::ArrayVec;
use crate::element_op::{ComposableElementOp, ElementOp, ElementOpApply, Identity};
use crate::{AxisVec, Result, StrideVec, StridedError, INLINE_RANK};
#[inline(always)]
fn arrayvec_from_slice<T: Copy + Default>(src: &[T]) -> ArrayVec<[T; INLINE_RANK]> {
debug_assert!(src.len() <= INLINE_RANK);
let n = src.len();
let mut buf = [T::default(); INLINE_RANK];
match n {
0 => {}
1 => buf[0] = src[0],
2 => {
buf[0] = src[0];
buf[1] = src[1];
}
3 => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
}
4 => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
buf[3] = src[3];
}
5 => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
buf[3] = src[3];
buf[4] = src[4];
}
6 => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
buf[3] = src[3];
buf[4] = src[4];
buf[5] = src[5];
}
7 => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
buf[3] = src[3];
buf[4] = src[4];
buf[5] = src[5];
buf[6] = src[6];
}
_ => {
buf[0] = src[0];
buf[1] = src[1];
buf[2] = src[2];
buf[3] = src[3];
buf[4] = src[4];
buf[5] = src[5];
buf[6] = src[6];
buf[7] = src[7];
}
}
ArrayVec::from_array_len(buf, n)
}
#[derive(Clone, Debug)]
struct SpilledMeta {
dims: Box<[usize]>,
strides: Box<[isize]>,
}
#[derive(Debug)]
struct ViewMeta {
inline_dims: ArrayVec<[usize; INLINE_RANK]>,
inline_strides: ArrayVec<[isize; INLINE_RANK]>,
spilled: Option<Box<SpilledMeta>>,
}
impl Clone for ViewMeta {
#[inline]
fn clone(&self) -> Self {
Self {
inline_dims: self.inline_dims,
inline_strides: self.inline_strides,
spilled: self.spilled.clone(),
}
}
}
impl ViewMeta {
#[inline(always)]
fn from_slices(dims: &[usize], strides: &[isize]) -> Self {
debug_assert_eq!(dims.len(), strides.len());
if dims.len() <= INLINE_RANK {
Self {
inline_dims: arrayvec_from_slice(dims),
inline_strides: arrayvec_from_slice(strides),
spilled: None,
}
} else {
Self {
inline_dims: ArrayVec::new(),
inline_strides: ArrayVec::new(),
spilled: Some(Box::new(SpilledMeta {
dims: dims.into(),
strides: strides.into(),
})),
}
}
}
#[inline]
fn dims(&self) -> &[usize] {
match &self.spilled {
Some(spilled) => &spilled.dims,
None => self.inline_dims.as_slice(),
}
}
#[inline]
fn strides(&self) -> &[isize] {
match &self.spilled {
Some(spilled) => &spilled.strides,
None => self.inline_strides.as_slice(),
}
}
}
#[inline]
fn empty_layout(dims: &[usize]) -> bool {
dims.contains(&0)
}
#[inline]
fn empty_const_ptr<T>() -> *const T {
core::ptr::NonNull::<T>::dangling().as_ptr()
}
#[inline]
fn empty_mut_ptr<T>() -> *mut T {
core::ptr::NonNull::<T>::dangling().as_ptr()
}
#[inline]
pub(crate) fn validate_bounds(
len: usize,
dims: &[usize],
strides: &[isize],
offset: isize,
) -> Result<()> {
if dims.len() != strides.len() {
return Err(StridedError::StrideLengthMismatch);
}
if dims.contains(&0) {
return Ok(());
}
let mut min_offset = offset;
let mut max_offset = offset;
for (&dim, &stride) in dims.iter().zip(strides.iter()) {
if dim > 1 {
let last = dim
.checked_sub(1)
.and_then(|value| isize::try_from(value).ok())
.ok_or(StridedError::OffsetOverflow)?;
let end = stride
.checked_mul(last)
.ok_or(StridedError::OffsetOverflow)?;
if end >= 0 {
max_offset = max_offset
.checked_add(end)
.ok_or(StridedError::OffsetOverflow)?;
} else {
min_offset = min_offset
.checked_add(end)
.ok_or(StridedError::OffsetOverflow)?;
}
}
}
if min_offset < 0 || max_offset < 0 {
return Err(StridedError::OffsetOverflow);
}
if max_offset as usize >= len {
return Err(StridedError::OffsetOverflow);
}
Ok(())
}
#[inline]
fn indexed_offset(dims: &[usize], strides: &[isize], indices: &[usize], base: isize) -> isize {
assert_eq!(indices.len(), dims.len(), "wrong number of indices");
let mut offset = base;
for (i, &index) in indices.iter().enumerate() {
assert!(
index < dims[i],
"index {} out of bounds for dim {}",
index,
dims[i]
);
offset = isize::try_from(index)
.ok()
.and_then(|index| index.checked_mul(strides[i]))
.and_then(|step| offset.checked_add(step))
.expect("index offset overflow");
}
offset
}
#[inline]
fn checked_element_count(dims: &[usize]) -> Result<usize> {
if dims.contains(&0) {
return Ok(0);
}
dims.iter()
.try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
.ok_or(StridedError::OffsetOverflow)
}
#[inline]
fn element_count(dims: &[usize]) -> usize {
checked_element_count(dims).expect("element count overflows usize")
}
fn checked_dense_strides(dims: &[usize], order: impl Iterator<Item = usize>) -> Result<StrideVec> {
let empty = dims.contains(&0);
let mut strides = StrideVec::new();
strides.resize(dims.len(), 0isize);
let mut step = 1isize;
for axis in order {
strides[axis] = step;
let dim = isize::try_from(dims[axis]).ok();
step = match dim.and_then(|dim| step.checked_mul(dim)) {
Some(next) => next,
None if empty => isize::MAX,
None => return Err(StridedError::OffsetOverflow),
};
}
Ok(strides)
}
fn checked_col_major_strides(dims: &[usize]) -> Result<StrideVec> {
checked_dense_strides(dims, 0..dims.len())
}
fn checked_row_major_strides(dims: &[usize]) -> Result<StrideVec> {
checked_dense_strides(dims, (0..dims.len()).rev())
}
#[inline]
pub fn col_major_strides(dims: &[usize]) -> StrideVec {
checked_col_major_strides(dims).expect("column-major strides overflow isize")
}
#[inline]
pub fn row_major_strides(dims: &[usize]) -> StrideVec {
checked_row_major_strides(dims).expect("row-major strides overflow isize")
}
pub struct StridedView<'a, T, Op = Identity> {
ptr: *const T,
data: &'a [T],
meta: ViewMeta,
offset: isize,
_op: PhantomData<Op>,
}
unsafe impl<T: Sync, Op: Send> Send for StridedView<'_, T, Op> {}
unsafe impl<T: Sync, Op: Sync> Sync for StridedView<'_, T, Op> {}
impl<T, Op> Clone for StridedView<'_, T, Op> {
#[inline]
fn clone(&self) -> Self {
Self {
ptr: self.ptr,
data: self.data,
meta: self.meta.clone(),
offset: self.offset,
_op: PhantomData,
}
}
}
impl<T: std::fmt::Debug, Op> std::fmt::Debug for StridedView<'_, T, Op> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StridedView")
.field("dims", &self.meta.dims())
.field("strides", &self.meta.strides())
.field("offset", &self.offset)
.finish()
}
}
#[inline]
fn base_ptr<T>(data: &[T], dims: &[usize], offset: isize) -> *const T {
if empty_layout(dims) {
empty_const_ptr()
} else {
unsafe { data.as_ptr().offset(offset) }
}
}
#[inline]
fn from_meta<'a, T, Op>(
data: &'a [T],
dims: &[usize],
strides: &[isize],
offset: isize,
) -> StridedView<'a, T, Op> {
let meta = ViewMeta::from_slices(dims, strides);
let ptr = base_ptr(data, meta.dims(), offset);
StridedView {
ptr,
data,
meta,
offset,
_op: PhantomData,
}
}
impl<'a, T, Op> StridedView<'a, T, Op> {
#[inline]
pub fn new(data: &'a [T], dims: &[usize], strides: &[isize], offset: isize) -> Result<Self> {
validate_bounds(data.len(), dims, strides, offset)?;
Ok(from_meta(data, dims, strides, offset))
}
#[inline]
pub unsafe fn new_unchecked(
data: &'a [T],
dims: &[usize],
strides: &[isize],
offset: isize,
) -> Self {
from_meta(data, dims, strides, offset)
}
#[inline]
pub fn dims(&self) -> &[usize] {
self.meta.dims()
}
#[inline]
pub fn strides(&self) -> &[isize] {
self.meta.strides()
}
#[inline]
pub fn offset(&self) -> isize {
self.offset
}
#[inline]
pub fn ndim(&self) -> usize {
self.dims().len()
}
#[inline]
pub fn len(&self) -> usize {
element_count(self.dims())
}
#[inline]
pub fn is_empty(&self) -> bool {
self.dims().contains(&0)
}
#[inline]
pub fn data(&self) -> &'a [T] {
self.data
}
#[inline]
pub fn ptr(&self) -> *const T {
self.ptr
}
pub fn permute(&self, perm: &[usize]) -> Result<StridedView<'a, T, Op>> {
let rank = self.dims().len();
if perm.len() != rank {
return Err(StridedError::RankMismatch(perm.len(), rank));
}
let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
seen.resize(rank, false);
for &p in perm {
if p >= rank {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
if seen[p] {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
seen[p] = true;
}
let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
Ok(from_meta(self.data, &new_dims, &new_strides, self.offset))
}
pub fn diagonal_view(&self, axis_pairs: &[(usize, usize)]) -> Result<StridedView<'a, T, Op>> {
let ndim = self.ndim();
let mut root: tinyvec::TinyVec<[usize; INLINE_RANK]> = (0..ndim).collect();
fn find(root: &mut [usize], mut ax: usize) -> usize {
while root[ax] != ax {
root[ax] = root[root[ax]];
ax = root[ax];
}
ax
}
for &(a, b) in axis_pairs {
let hi = a.max(b);
if hi >= ndim {
return Err(StridedError::InvalidAxis {
axis: hi,
rank: ndim,
});
}
if a == b {
return Err(StridedError::InvalidAxis {
axis: a,
rank: ndim,
});
}
let (ra, rb) = (find(&mut root, a), find(&mut root, b));
root[ra.max(rb)] = ra.min(rb);
}
let mut dims: AxisVec = AxisVec::new();
let mut strides: StrideVec = StrideVec::new();
let mut position: tinyvec::TinyVec<[usize; INLINE_RANK]> =
tinyvec::TinyVec::with_capacity(ndim);
for ax in 0..ndim {
let r = find(&mut root, ax);
if r == ax {
position.push(dims.len());
dims.push(self.dims()[ax]);
strides.push(self.strides()[ax]);
} else {
position.push(usize::MAX);
let p = position[r];
strides[p] = strides[p]
.checked_add(self.strides()[ax])
.ok_or(StridedError::OffsetOverflow)?;
dims[p] = dims[p].min(self.dims()[ax]);
}
}
validate_bounds(self.data().len(), &dims, &strides, self.offset())?;
Ok(from_meta(self.data(), &dims, &strides, self.offset()))
}
pub fn broadcast(&self, target_dims: &[usize]) -> Result<StridedView<'a, T, Op>> {
if self.dims().len() != target_dims.len() {
return Err(StridedError::RankMismatch(
self.dims().len(),
target_dims.len(),
));
}
let mut new_strides = StrideVec::with_capacity(self.dims().len());
for i in 0..self.dims().len() {
if self.dims()[i] == target_dims[i] {
new_strides.push(self.strides()[i]);
} else if self.dims()[i] == 1 {
new_strides.push(0);
} else {
return Err(StridedError::ShapeMismatch(
self.dims().to_vec(),
target_dims.to_vec(),
));
}
}
Ok(from_meta(self.data, target_dims, &new_strides, self.offset))
}
}
impl<'a, T: Copy + ElementOpApply, Op: ComposableElementOp<T>> StridedView<'a, T, Op> {
pub fn transpose_2d(&self) -> Result<StridedView<'a, T, Op::ComposeTranspose>> {
if self.dims().len() != 2 {
return Err(StridedError::RankMismatch(self.dims().len(), 2));
}
Ok(from_meta(
self.data,
&[self.dims()[1], self.dims()[0]],
&[self.strides()[1], self.strides()[0]],
self.offset,
))
}
pub fn adjoint_2d(&self) -> Result<StridedView<'a, T, Op::ComposeAdjoint>> {
if self.dims().len() != 2 {
return Err(StridedError::RankMismatch(self.dims().len(), 2));
}
Ok(from_meta(
self.data,
&[self.dims()[1], self.dims()[0]],
&[self.strides()[1], self.strides()[0]],
self.offset,
))
}
#[inline]
pub fn conj(&self) -> StridedView<'a, T, Op::ComposeConj> {
StridedView {
ptr: self.ptr,
data: self.data,
meta: self.meta.clone(),
offset: self.offset,
_op: PhantomData,
}
}
}
impl<'a, T: Copy, Op: ElementOp<T>> StridedView<'a, T, Op> {
pub fn get(&self, indices: &[usize]) -> T {
let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
Op::apply(unsafe { *self.ptr.offset(offset) })
}
#[inline]
pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
let mut idx = 0isize;
for (i, &index) in indices.iter().enumerate() {
idx += index as isize * self.strides()[i];
}
Op::apply(*self.ptr.offset(idx))
}
}
pub struct StridedViewMut<'a, T> {
ptr: *mut T,
base: *mut T,
src_len: usize,
meta: ViewMeta,
offset: isize,
_marker: PhantomData<&'a mut [T]>,
}
unsafe impl<T: Send> Send for StridedViewMut<'_, T> {}
impl<T: std::fmt::Debug> std::fmt::Debug for StridedViewMut<'_, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StridedViewMut")
.field("dims", &self.meta.dims())
.field("strides", &self.meta.strides())
.field("offset", &self.offset)
.finish()
}
}
impl<'a, T> StridedViewMut<'a, T> {
#[inline]
fn from_meta(data: &'a mut [T], dims: &[usize], strides: &[isize], offset: isize) -> Self {
let meta = ViewMeta::from_slices(dims, strides);
let src_len = data.len();
let base = data.as_mut_ptr();
let ptr = if empty_layout(meta.dims()) {
empty_mut_ptr()
} else {
unsafe { base.offset(offset) }
};
Self {
ptr,
base,
src_len,
meta,
offset,
_marker: PhantomData,
}
}
#[inline]
pub fn new(
data: &'a mut [T],
dims: &[usize],
strides: &[isize],
offset: isize,
) -> Result<Self> {
validate_bounds(data.len(), dims, strides, offset)?;
Ok(Self::from_meta(data, dims, strides, offset))
}
#[inline]
pub unsafe fn new_unchecked(
data: &'a mut [T],
dims: &[usize],
strides: &[isize],
offset: isize,
) -> Self {
Self::from_meta(data, dims, strides, offset)
}
#[inline]
pub fn dims(&self) -> &[usize] {
self.meta.dims()
}
#[inline]
pub fn strides(&self) -> &[isize] {
self.meta.strides()
}
#[inline]
pub fn offset(&self) -> isize {
self.offset
}
#[inline]
pub fn ndim(&self) -> usize {
self.dims().len()
}
#[inline]
pub fn len(&self) -> usize {
element_count(self.dims())
}
#[inline]
pub fn is_empty(&self) -> bool {
self.dims().contains(&0)
}
#[inline]
pub fn ptr(&self) -> *const T {
self.ptr as *const T
}
#[inline]
pub fn as_mut_ptr(&self) -> *mut T {
self.ptr
}
#[inline]
pub fn data(&self) -> &[T] {
unsafe { std::slice::from_raw_parts(self.base, self.src_len) }
}
#[inline]
pub fn data_mut(&mut self) -> &mut [T] {
unsafe { std::slice::from_raw_parts_mut(self.base, self.src_len) }
}
pub fn permute(self, perm: &[usize]) -> Result<StridedViewMut<'a, T>> {
let rank = self.dims().len();
if perm.len() != rank {
return Err(StridedError::RankMismatch(perm.len(), rank));
}
let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
seen.resize(rank, false);
for &p in perm {
if p >= rank {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
if seen[p] {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
seen[p] = true;
}
let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
Ok(Self {
meta: ViewMeta::from_slices(&new_dims, &new_strides),
..self
})
}
#[inline]
pub fn as_view(&self) -> StridedView<'_, T, Identity> {
unsafe { StridedView::new_unchecked(self.data(), self.dims(), self.strides(), self.offset) }
}
}
impl<'a, T: Copy> StridedViewMut<'a, T> {
pub fn get(&self, indices: &[usize]) -> T {
let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
unsafe { *self.ptr.offset(offset) }
}
pub fn set(&mut self, indices: &[usize], value: T) {
let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
unsafe {
*self.ptr.offset(offset) = value;
}
}
}
pub struct StridedArray<T> {
data: Vec<T>,
meta: ViewMeta,
offset: isize,
}
impl<T: std::fmt::Debug> std::fmt::Debug for StridedArray<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StridedArray")
.field("dims", &self.meta.dims())
.field("strides", &self.meta.strides())
.field("offset", &self.offset)
.finish()
}
}
impl<T: Clone> Clone for StridedArray<T> {
fn clone(&self) -> Self {
Self {
data: self.data.clone(),
meta: self.meta.clone(),
offset: self.offset,
}
}
}
impl<T: Clone + Default> StridedArray<T> {
pub fn col_major(dims: &[usize]) -> Self {
Self::try_col_major(dims).expect("col_major: shape overflows isize")
}
pub fn try_col_major(dims: &[usize]) -> Result<Self> {
let strides = checked_col_major_strides(dims)?;
let total = checked_element_count(dims)?;
Ok(Self {
data: vec![T::default(); total],
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
pub fn row_major(dims: &[usize]) -> Self {
Self::try_row_major(dims).expect("row_major: shape overflows isize")
}
pub fn try_row_major(dims: &[usize]) -> Result<Self> {
let strides = checked_row_major_strides(dims)?;
let total = checked_element_count(dims)?;
Ok(Self {
data: vec![T::default(); total],
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
pub fn from_fn_col_major(dims: &[usize], f: impl FnMut(&[usize]) -> T) -> Self {
Self::try_from_fn_col_major(dims, f).expect("from_fn_col_major: shape overflows isize")
}
pub fn try_from_fn_col_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Result<Self> {
let strides = checked_col_major_strides(dims)?;
let total = checked_element_count(dims)?;
let rank = dims.len();
let mut data = Vec::with_capacity(total);
let mut idx: AxisVec = AxisVec::new();
idx.resize(rank, 0usize);
let idx: &mut [usize] = &mut idx[..];
for _ in 0..total {
data.push(f(idx));
for d in 0..rank {
idx[d] += 1;
if idx[d] < dims[d] {
break;
}
idx[d] = 0;
}
}
Ok(Self {
data,
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
pub fn from_fn_row_major(dims: &[usize], f: impl FnMut(&[usize]) -> T) -> Self {
Self::try_from_fn_row_major(dims, f).expect("from_fn_row_major: shape overflows isize")
}
pub fn try_from_fn_row_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Result<Self> {
let strides = checked_row_major_strides(dims)?;
let total = checked_element_count(dims)?;
let rank = dims.len();
let mut data = Vec::with_capacity(total);
let mut idx: AxisVec = AxisVec::new();
idx.resize(rank, 0usize);
let idx: &mut [usize] = &mut idx[..];
for _ in 0..total {
data.push(f(idx));
for d in (0..rank).rev() {
idx[d] += 1;
if idx[d] < dims[d] {
break;
}
idx[d] = 0;
}
}
Ok(Self {
data,
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
}
impl<T> StridedArray<T> {
pub fn from_parts(
data: Vec<T>,
dims: &[usize],
strides: &[isize],
offset: isize,
) -> Result<Self> {
validate_bounds(data.len(), dims, strides, offset)?;
Ok(Self {
data,
meta: ViewMeta::from_slices(dims, strides),
offset,
})
}
#[inline]
pub fn dims(&self) -> &[usize] {
self.meta.dims()
}
#[inline]
pub fn strides(&self) -> &[isize] {
self.meta.strides()
}
#[inline]
pub fn ndim(&self) -> usize {
self.dims().len()
}
#[inline]
pub fn len(&self) -> usize {
element_count(self.dims())
}
#[inline]
pub fn is_empty(&self) -> bool {
self.dims().contains(&0)
}
#[inline]
pub fn data(&self) -> &[T] {
&self.data
}
#[inline]
pub fn data_mut(&mut self) -> &mut [T] {
&mut self.data
}
#[inline]
pub fn view(&self) -> StridedView<'_, T> {
let ptr = if self.is_empty() {
empty_const_ptr()
} else {
unsafe { self.data.as_ptr().offset(self.offset) }
};
StridedView {
ptr,
data: &self.data,
meta: self.meta.clone(),
offset: self.offset,
_op: PhantomData,
}
}
#[inline]
pub fn view_mut(&mut self) -> StridedViewMut<'_, T> {
let src_len = self.data.len();
let base = self.data.as_mut_ptr();
let ptr = if self.is_empty() {
empty_mut_ptr()
} else {
unsafe { base.offset(self.offset) }
};
StridedViewMut {
ptr,
base,
src_len,
meta: self.meta.clone(),
offset: self.offset,
_marker: PhantomData,
}
}
pub fn permuted(self, perm: &[usize]) -> Result<Self> {
let rank = self.dims().len();
if perm.len() != rank {
return Err(StridedError::RankMismatch(perm.len(), rank));
}
let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
seen.resize(rank, false);
for &p in perm {
if p >= rank {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
if seen[p] {
return Err(StridedError::InvalidAxis { axis: p, rank });
}
seen[p] = true;
}
let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
Ok(Self {
data: self.data,
meta: ViewMeta::from_slices(&new_dims, &new_strides),
offset: self.offset,
})
}
pub fn into_data(self) -> Vec<T> {
self.data
}
pub fn iter(&self) -> std::slice::Iter<'_, T> {
self.data.iter()
}
pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
self.data.iter_mut()
}
}
impl<T: Default> StridedArray<T> {
pub fn col_major_from_buffer(buf: Vec<T>, dims: &[usize]) -> Self {
Self::try_col_major_from_buffer(buf, dims)
.expect("col_major_from_buffer: shape overflows isize")
}
pub fn try_col_major_from_buffer(mut buf: Vec<T>, dims: &[usize]) -> Result<Self> {
let strides = checked_col_major_strides(dims)?;
let total = checked_element_count(dims)?;
if buf.len() >= total {
buf.truncate(total);
} else {
buf.resize_with(total, T::default);
}
for v in buf.iter_mut() {
*v = T::default();
}
Ok(Self {
data: buf,
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
}
impl<T: Copy> StridedArray<T> {
pub unsafe fn col_major_uninit(dims: &[usize]) -> StridedArray<MaybeUninit<T>> {
Self::try_col_major_uninit(dims).expect("col_major_uninit: shape overflows isize")
}
pub unsafe fn try_col_major_uninit(dims: &[usize]) -> Result<StridedArray<MaybeUninit<T>>> {
let strides = checked_col_major_strides(dims)?;
let total = checked_element_count(dims)?;
Ok(StridedArray {
data: vec![MaybeUninit::uninit(); total],
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
pub unsafe fn col_major_from_buffer_uninit(
buf: Vec<MaybeUninit<T>>,
dims: &[usize],
) -> StridedArray<MaybeUninit<T>> {
Self::try_col_major_from_buffer_uninit(buf, dims)
.expect("col_major_from_buffer_uninit: shape overflows isize")
}
pub unsafe fn try_col_major_from_buffer_uninit(
mut buf: Vec<MaybeUninit<T>>,
dims: &[usize],
) -> Result<StridedArray<MaybeUninit<T>>> {
let strides = checked_col_major_strides(dims)?;
let total = checked_element_count(dims)?;
buf.resize(total, MaybeUninit::uninit());
Ok(StridedArray {
data: buf,
meta: ViewMeta::from_slices(dims, &strides),
offset: 0,
})
}
}
impl<T: Copy> StridedArray<T> {
pub fn get(&self, indices: &[usize]) -> T {
self.view().get(indices)
}
pub fn set(&mut self, indices: &[usize], value: T) {
let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
self.data[offset as usize] = value;
}
}
impl<T: Copy> Index<&[usize]> for StridedArray<T> {
type Output = T;
fn index(&self, indices: &[usize]) -> &T {
let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
&self.data[offset as usize]
}
}
impl<T: Copy> IndexMut<&[usize]> for StridedArray<T> {
fn index_mut(&mut self, indices: &[usize]) -> &mut T {
let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
&mut self.data[offset as usize]
}
}
#[cfg(test)]
mod tests;