use crate::dense::array::Array;
use crate::dense::layout::{check_multi_index_in_bounds, convert_1d_nd_from_shape};
use crate::traits::{
accessors::{
UnsafeDisjointRandom1DAccessMut, UnsafeRandom1DAccessByRef, UnsafeRandom1DAccessByValue,
UnsafeRandom1DAccessMut, UnsafeRandomAccessByRef, UnsafeRandomAccessByValue,
UnsafeRandomAccessMut,
},
base_operations::{BaseItem, Shape, Stride},
};
use crate::{ContainerType, Unknown};
pub struct ArraySubView<ArrayImpl, const NDIM: usize> {
arr: Array<ArrayImpl, NDIM>,
offset: [usize; NDIM],
shape: [usize; NDIM],
}
impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> ArraySubView<ArrayImpl, NDIM> {
pub fn new(arr: Array<ArrayImpl, NDIM>, offset: [usize; NDIM], shape: [usize; NDIM]) -> Self {
let arr_shape = arr.shape();
for index in 0..NDIM {
assert!(
offset[index] + shape[index] <= arr_shape[index],
"View out of bounds for dimension {}. {} > {}",
index,
offset[index] + shape[index],
arr_shape[index]
)
}
Self { arr, offset, shape }
}
}
impl<ArrayImpl, const NDIM: usize> ContainerType for ArraySubView<ArrayImpl, NDIM>
where
ArrayImpl: ContainerType,
{
type Type = Unknown;
}
impl<ArrayImpl, const NDIM: usize> BaseItem for ArraySubView<ArrayImpl, NDIM>
where
ArrayImpl: BaseItem,
{
type Item = ArrayImpl::Item;
}
impl<ArrayImpl, const NDIM: usize> Shape<NDIM> for ArraySubView<ArrayImpl, NDIM> {
#[inline(always)]
fn shape(&self) -> [usize; NDIM] {
self.shape
}
}
impl<ArrayImpl: Stride<NDIM>, const NDIM: usize> Stride<NDIM> for ArraySubView<ArrayImpl, NDIM> {
#[inline(always)]
fn stride(&self) -> [usize; NDIM] {
self.arr.stride()
}
}
impl<ArrayImpl: UnsafeRandomAccessByValue<NDIM>, const NDIM: usize> UnsafeRandomAccessByValue<NDIM>
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_value_unchecked(&self, multi_index: [usize; NDIM]) -> Self::Item {
debug_assert!(check_multi_index_in_bounds(multi_index, self.shape()));
unsafe {
self.arr
.get_value_unchecked(offset_multi_index(multi_index, self.offset))
}
}
}
impl<ArrayImpl: UnsafeRandomAccessByRef<NDIM>, const NDIM: usize> UnsafeRandomAccessByRef<NDIM>
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_unchecked(&self, multi_index: [usize; NDIM]) -> &Self::Item {
debug_assert!(check_multi_index_in_bounds(multi_index, self.shape()));
unsafe {
self.arr
.get_unchecked(offset_multi_index(multi_index, self.offset))
}
}
}
impl<ArrayImpl: UnsafeRandomAccessByValue<NDIM>, const NDIM: usize> UnsafeRandom1DAccessByValue
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_value_1d_unchecked(&self, index: usize) -> Self::Item {
unsafe { self.get_value_unchecked(convert_1d_nd_from_shape(index, self.shape)) }
}
}
impl<ArrayImpl: UnsafeRandomAccessByRef<NDIM>, const NDIM: usize> UnsafeRandom1DAccessByRef
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_1d_unchecked(&self, index: usize) -> &Self::Item {
unsafe { self.get_unchecked(convert_1d_nd_from_shape(index, self.shape)) }
}
}
impl<ArrayImpl: UnsafeRandomAccessMut<NDIM>, const NDIM: usize> UnsafeRandom1DAccessMut
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_1d_unchecked_mut(&mut self, index: usize) -> &mut Self::Item {
unsafe { self.get_unchecked_mut(convert_1d_nd_from_shape(index, self.shape)) }
}
}
unsafe impl<ArrayImpl, const NDIM: usize> UnsafeDisjointRandom1DAccessMut
for ArraySubView<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + UnsafeRandomAccessMut<NDIM>,
{
}
impl<ArrayImpl: UnsafeRandomAccessMut<NDIM>, const NDIM: usize> UnsafeRandomAccessMut<NDIM>
for ArraySubView<ArrayImpl, NDIM>
{
#[inline(always)]
unsafe fn get_unchecked_mut(&mut self, multi_index: [usize; NDIM]) -> &mut Self::Item {
debug_assert!(crate::dense::layout::check_multi_index_in_bounds(
multi_index,
self.shape()
));
unsafe {
self.arr
.get_unchecked_mut(offset_multi_index(multi_index, self.offset))
}
}
}
fn offset_multi_index<const NDIM: usize>(
multi_index: [usize; NDIM],
offset: [usize; NDIM],
) -> [usize; NDIM] {
let mut output = [0; NDIM];
for (ind, elem) in output.iter_mut().enumerate() {
*elem = multi_index[ind] + offset[ind]
}
output
}