rstango 0.1.0

//! Generic multi-dimensional arrays.

use std::mem::ManuallyDrop;
use std::ptr::{self, NonNull};
use std::marker::PhantomData;
use std::{fmt, mem, slice};


/// An owned multi-dimensional array of values.
/// 
/// This array is internally backend by a single array, the **row-major** order is used.
pub struct MdArray<T, const DIM: usize> {
    ptr: NonNull<T>,
    dim: [usize; DIM],
    _phantom: PhantomData<Box<[T]>>,
}

unsafe impl<T: Send, const DIM: usize> Send for MdArray<T, DIM> { }
unsafe impl<T: Sync, const DIM: usize> Sync for MdArray<T, DIM> { }

impl<T, const DIM: usize> MdArray<T, DIM> {

    /// The caller has to ensure that the given dimensions are correct for the given 
    /// slice, or else the later use of this array will be unsound.
    #[inline]
    pub unsafe fn from_boxed_slice_unchecked(slice: Box<[T]>, dim: [usize; DIM]) -> Self {
        // SAFETY: Box::into_raw doc says "the pointer will be aligned and non-null".
        Self {
            ptr: unsafe { NonNull::new_unchecked(Box::into_raw(slice) as *mut T) },
            dim,
            _phantom: PhantomData
        }
    }

    pub fn try_from_boxed_slice<S: Into<Box<[T]>>>(slice: S, dim: [usize; DIM]) -> Option<Self> {
        
        let len = array_len(&dim)?;
        let slice = slice.into();
        if slice.len() != len {
            return None;
        }

        // SAFETY: We checked that the array len matches the given slice length.
        Some(unsafe { Self::from_boxed_slice_unchecked(slice, dim) })

    }

    #[inline]
    pub fn into_boxed_slice(self: Self) -> Box<[T]> {
        // SAFETY: We created that pointer from Box::<[T]>::into_raw, so we construct it 
        // back in order to destruct it! We have to reconstruct the original *mut [T] 
        // pointer by computing the length from all dimensions.
        // NOTE: Unwrap because we already validated at construction time!!
        unsafe { 
            let this = ManuallyDrop::new(self);
            let len = array_len(&this.dim).unwrap_unchecked();
            let ptr = ptr::slice_from_raw_parts_mut(this.ptr.as_ptr(), len);
            Box::from_raw(ptr)
        }
    }

    #[inline]
    pub fn as_view(&self) -> MdView<'_, T, DIM> {
        MdView {
            ptr: self.ptr,
            dim: self.dim,
            _phantom: PhantomData,
        }
    }

    #[inline]
    pub fn is_empty(self) -> bool {
        self.as_view().is_empty()
    }

    #[inline]
    pub fn raw_len(&self) -> usize {
        self.as_view().raw_len()
    }

    #[inline]
    pub fn as_raw_slice(&self) -> &[T] {
        self.as_view().as_raw_slice()
    }

    #[inline]
    pub fn as_raw_slice_mut(&mut self) -> &mut [T] {
        let len = self.raw_len();
        // SAFETY: We compute the flat array length from the known dimensions, using that
        // we can create a *mut [T] pointer that can be transformed into &mut [T].
        unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), len) }
    }
    
    #[inline]
    pub fn len(&self, dim_index: usize) -> usize {
        self.dim[dim_index]
    }

    #[inline]
    pub fn get(&self, index: [usize; DIM]) -> Option<&T> {
        self.as_view().get(index)
    }

    #[inline]
    pub fn get_mut(&mut self, index: [usize; DIM]) -> Option<&mut T> {
        let index = array_index(&self.dim, &index)?;
        // SAFETY: We checked that the index is correct for the dimensions.
        unsafe { Some(&mut *self.ptr.add(index).as_ptr()) }
    }

    /// Map each element by value and transform this array into another one with all 
    /// results from the transform function.
    #[inline]
    pub fn map<R, F: FnMut(T) -> R>(self, func: F) -> MdArray<R, DIM> {
        
        let dim = self.dim;
        let new_slice = Self::into_boxed_slice(self)
            .into_iter()
            .map(func)
            .collect::<Box<[R]>>();

        // SAFETY: We are just mapping each element, the sliced box has the same size
        // and the dimension is unchanged: the precondition is guaranteed.
        unsafe { MdArray::from_boxed_slice_unchecked(new_slice, dim) }
        
    }

    /// Map each element through its reference to a new owned value, that is returned as
    /// a new owned array.
    #[inline]
    pub fn map_to_owned<'a, U, F: FnMut(&'a T) -> U>(&'a self, func: F) -> MdArray<U, DIM> {
        self.as_view().map_to_owned(func)
    }

    #[inline]
    pub fn as_shape<const NEW_DIM: usize>(&self, new_dim: [usize; NEW_DIM]) -> Option<MdView<'_, T, NEW_DIM>> {
        self.as_view().as_shape(new_dim)
    }

    /// Re-interpret this array as another shape!
    #[inline]
    pub fn into_shape<const NEW_DIM: usize>(self, new_dim: [usize; NEW_DIM]) -> Result<MdArray<T, NEW_DIM>, Self> {

        if array_len(&new_dim) != Some(self.raw_len()) {
            return Err(self);
        }

        Ok(MdArray {
            ptr: self.ptr,
            dim: new_dim,
            _phantom: PhantomData,
        })

    }

}

impl<T, const DIM: usize> Drop for MdArray<T, DIM> {
    fn drop(&mut self) {
        // Take the data to transform it into a boxed slice that we immediately drop.
        // Read safety comment inside `into_boxed_slice`.
        drop(mem::replace(self, Self {
            ptr: NonNull::dangling(),
            dim: [0; DIM],
            _phantom: PhantomData,
        }).into_boxed_slice())
    }
}

impl<T: Clone, const DIM: usize> Clone for MdArray<T, DIM> {
    fn clone(&self) -> Self {
        self.as_view().to_owned()
    }
}

impl<T: fmt::Debug, const DIM: usize> fmt::Debug for MdArray<T, DIM> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_fmt(format_args!("MdArray"))?;
        fmt::Debug::fmt(&self.as_view(), f)
    }
}

impl<T> From<T> for MdArray<T, 0> {
    #[inline]
    fn from(value: T) -> Self {
        // SAFETY: We have no dimension.
        unsafe {
            Self::from_boxed_slice_unchecked(Box::new([value]), [])
        }
    }
}

impl<T> From<Box<[T]>> for MdArray<T, 1> {
    #[inline]
    fn from(value: Box<[T]>) -> Self {
        // SAFETY: We have only one dimension, and we give the slice's length, which is
        // equal to... the slice length, which validate the precondition.
        unsafe {
            let len = value.len();
            Self::from_boxed_slice_unchecked(value, [len])
        }
    }
}

impl<T> From<Vec<T>> for MdArray<T, 1> {
    #[inline]
    fn from(value: Vec<T>) -> Self {
        value.into_boxed_slice().into()
    }
}

impl<T, const N: usize> From<[T; N]> for MdArray<T, 1> {
    #[inline]
    fn from(value: [T; N]) -> Self {
        Box::<[T]>::from(value).into()
    }
}

/// A view to a full [`MdArray`].
pub struct MdView<'a, T: 'a, const DIM: usize> {
    ptr: NonNull<T>,
    dim: [usize; DIM],
    _phantom: PhantomData<&'a [T]>,
}

unsafe impl<'a, T: Send, const DIM: usize> Send for MdView<'a, T, DIM> { }
unsafe impl<'a, T: Sync, const DIM: usize> Sync for MdView<'a, T, DIM> { }

impl<'a, T, const DIM: usize> MdView<'a, T, DIM> {

    /// The caller has to ensure that the given dimensions are correct for the given 
    /// slice, or else the later use of this array will be unsound.
    #[inline]
    pub unsafe fn from_slice_unchecked(slice: &[T], dim: [usize; DIM]) -> Self {
        Self {
            ptr: NonNull::from_ref(&slice[0]),
            dim,
            _phantom: PhantomData
        }
    }

    pub fn try_from_slice<S: AsRef<[T]>>(slice: S, dim: [usize; DIM]) -> Option<Self> {

        let len = array_len(&dim)?;
        let slice = slice.as_ref();
        if slice.len() != len {
            return None;
        }

        // SAFETY: We checked that the array len matches the given slice length.
        Some(unsafe { Self::from_slice_unchecked(slice, dim) })

    }

    #[inline]
    pub fn is_empty(self) -> bool {
        self.raw_len() == 0
    }

    #[inline]
    pub fn raw_len(self) -> usize {
        // SAFETY: Unwrap because we already validated at construction time!!
        unsafe { array_len(&self.dim).unwrap_unchecked() }
    }

    #[inline]
    pub fn as_raw_slice(self) -> &'a [T] {
        let len = self.raw_len();
        // SAFETY: We compute the flat array length from the known dimensions, using that
        // we can create a *const [T] pointer that can be transformed into &[T].
        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), len) }
    }

    #[inline]
    pub fn to_owned(self) -> MdArray<T, DIM>
    where 
        T: Clone,
    {
        // SAFETY: We copy the slice into a new box, so we know that we have the same 
        // dimensions and length as the current view.
        let slice = Box::<[T]>::from(self.as_raw_slice());
        unsafe {
            MdArray::from_boxed_slice_unchecked(slice, self.dim)
        }
    }

    #[inline]
    pub fn len(self, dim_index: usize) -> usize {
        self.dim[dim_index]
    }

    #[inline]
    pub fn get(self, index: [usize; DIM]) -> Option<&'a T> {
        // SAFETY: We checked that the index is correct for the dimensions.
        let index = array_index(&self.dim, &index)?;
        unsafe {
            Some(&*self.ptr.add(index).as_ptr())
        }
    }

    /// Map each element through its reference to a new owned value, that is returned as
    /// a new owned array.
    #[inline]
    pub fn map_to_owned<U, F: FnMut(&'a T) -> U>(self, func: F) -> MdArray<U, DIM> {
        
        let dim = self.dim;
        let new_slice = self.as_raw_slice()
            .into_iter()
            .map(func)
            .collect::<Box<[U]>>();

        // SAFETY: We are just mapping each element, the sliced box has the same size
        // and the dimension is unchanged: the precondition is guaranteed.
        unsafe { MdArray::from_boxed_slice_unchecked(new_slice, dim) }
        
    }

    /// Re-interpret this array view as another shape!
    #[inline]
    pub fn as_shape<const NEW_DIM: usize>(self, new_dim: [usize; NEW_DIM]) -> Option<MdView<'a, T, NEW_DIM>> {

        if array_len(&new_dim) != Some(self.raw_len()) {
            return None;
        }

        Some(MdView {
            ptr: self.ptr,
            dim: new_dim,
            _phantom: PhantomData,
        })

    }

}

impl<'a, T, const DIM: usize> Clone for MdView<'a, T, DIM> {
    #[inline]
    fn clone(&self) -> Self {
        Self { 
            ptr: self.ptr, 
            dim: self.dim, 
            _phantom: PhantomData,
        }
    }
}

impl<'a, T, const DIM: usize> Copy for MdView<'a, T, DIM> { }

impl<'a, T: fmt::Debug, const DIM: usize> fmt::Debug for MdView<'a, T, DIM> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        
        f.write_fmt(format_args!("<{DIM}>"))?;
        
        let mut counter = [0; DIM];
        for val in self.as_raw_slice() {

            let mut open_count = 0;
            while open_count < DIM && counter[open_count] == 0 {
                open_count += 1;
            }
            
            if counter != [0; DIM] {
                f.write_str(", ")?;
                if open_count != 0 && f.alternate() {
                    f.write_str("\n")?;
                }
            }
            
            for i in (0..open_count).rev() {
                f.write_str("[")?;
                if i != 0 && f.alternate() {
                    f.write_str("\n")?;
                }
            }

            f.write_fmt(format_args!("{val:?}"))?;
            
            counter[0] += 1;
            for i in 0..DIM {
                if counter[i] == self.dim[i] {
                    f.write_str("]")?;
                    if let Some(next_count) = counter.get_mut(i + 1) {
                        *next_count += 1;
                    }
                    counter[i] = 0;
                }
            }

        }

        Ok(())

    }
}

impl<'a, T> From<&'a T> for MdView<'a, T, 0> {
    #[inline]
    fn from(value: &'a T) -> Self {
        // SAFETY: We have no dimension.
        unsafe {
            Self::from_slice_unchecked(slice::from_ref(value), [])
        }
    }
}

impl<'a, T> From<&'a [T]> for MdView<'a, T, 1> {
    #[inline]
    fn from(value: &'a [T]) -> Self {
        // SAFETY: We have only one dimension, and we give the slice's length, which is
        // equal to... the slice length, which validate the precondition.
        unsafe {
            let len = value.len();
            Self::from_slice_unchecked(value, [len])
        }
    }
}

fn array_len<const DIM: usize>(dim: &[usize; DIM]) -> Option<usize> {
    let mut flat = 1usize;
    for &val in dim {
        flat = flat.checked_mul(val)?;
    }
    Some(flat)
}

fn array_index<const DIM: usize>(dim: &[usize; DIM], index: &[usize; DIM]) -> Option<usize> {
    let mut flat = 0;
    for i in (0..DIM).rev() {
        if index[i] >= dim[i] {
            return None;
        }
        flat *= dim[i];
        flat += index[i];
    }
    Some(flat)
}