xtk_core/

array.rs

#![expect(
    unsafe_op_in_unsafe_fn,
    reason = "Array requires unsafe code in some places"
)]

//! Variable length heap array

use core::ops::{Deref, DerefMut, Index, IndexMut};
use std::alloc;

/// Variable length array.
///
/// The purpose of this data structure is to provide an alternative to [Vec] with a static length.
/// Use a normal array if you know the length at compile time.
///
/// A possible use for [Array] is where you need to read from a [std::io::Read]er to a [Vec] using
/// the [std::io::Read::read] function, this wouldn't work because some elements of a [Vec] may be
/// uninitialized whereas [Array] functions just like a static compile-time known array.
///
/// # Examples
///
/// ```
/// use xtk_core::array::Array;
///
/// let mut array = Array::new(5);
///
/// array[0] = 1;
/// array[1] = 2;
/// array[2] = 3;
/// array[3] = 4;
/// array[4] = 5;
/// ```
pub struct Array<T> {
    ptr: *mut T,
    size: usize,
    layout: alloc::Layout,
}

impl<T> Array<T> {
    /// Allocate an [Array]
    ///
    ///
    /// # Examples
    /// ```
    /// use xtk_core::array::Array;
    ///
    /// let mut array = Array::new(10);
    ///
    /// array[0] = 1;
    /// ```
    #[must_use]
    pub fn new(size: usize) -> Self {
        unsafe {
            let layout = alloc::Layout::array::<T>(size).expect("layout creation failed");
            let ptr = alloc::alloc(layout) as *mut T;

            assert!(!ptr.is_null(), "array allocation failed");

            Self { ptr, size, layout }
        }
    }

    /// Get the size of the array
    #[must_use]
    #[doc(alias = "len")]
    pub const fn size(&self) -> usize {
        self.size
    }

    /// Return the array as an immutable pointer
    #[must_use]
    pub const fn as_ptr(&self) -> *const T {
        self.ptr
    }

    /// Return the array as a mutable pointer
    #[must_use]
    pub const fn as_mut_ptr(&self) -> *mut T {
        self.ptr
    }

    /// Set value at `index` to `value`
    ///
    /// # Panics
    ///
    /// Panics when `index` is out of bounds
    pub fn set(&mut self, index: usize, value: T) {
        if index >= self.size {
            panic!("index out of bounds");
        }

        // SAFETY: it is ensured that whatever this pointer points to is inside of the array by
        // the previous statement
        unsafe { *(self.ptr.add(index)) = value }
    }

    /// Get value at index `index`
    #[must_use]
    pub const fn get(&self, index: usize) -> Option<&T> {
        if index >= self.size {
            return None;
        }

        // SAFETY: it is ensured that whatever this pointer points to is inside of the `Array` by
        // the previous statement
        unsafe { Some(&(*(self.ptr.add(index)))) }
    }

    /// Get value at index `index`
    #[must_use]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        if index >= self.size {
            return None;
        }

        // SAFETY: it is ensured that whatever this pointer points to is inside of the `Array` by
        // the previous statement
        unsafe { Some(&mut (*(self.ptr.add(index)))) }
    }

    /// Returns an immutable pointer to value at `index`.
    ///
    /// # Safety
    ///
    /// The pointer is obtained using unchecked pointer arithmetic.
    #[must_use]
    pub const unsafe fn get_ptr(&self, index: usize) -> *const T {
        self.ptr.add(index) as *const T
    }

    /// Returns a mutable pointer to value at `index`.
    ///
    /// # Safety
    ///
    /// The pointer is obtained using unchecked pointer arithmetic.
    #[must_use]
    pub const unsafe fn get_ptr_mut(&mut self, index: usize) -> *mut T {
        self.ptr.add(index)
    }

    /// Return an iterator over the array
    #[must_use]
    pub fn iter(&self) -> iter::Iter<'_, T> {
        iter::Iter::new(self)
    }

    /// Return a mutable iterator over the array
    #[must_use]
    pub fn iter_mut(&mut self) -> iter::IterMut<'_, T> {
        iter::IterMut::new(self)
    }
}

impl<T: Clone> Array<T> {
    /// Fill array with elements of `slice` by cloning the elements
    pub fn fill(&mut self, slice: &[T]) {
        for (i, t) in slice.iter().enumerate() {
            if i >= self.size {
                break;
            }
            self[i] = t.clone();
        }
    }

    /// Allocates a new array with `n` extra length, clones all elements of the array into the new
    /// one, hence `T` must implement [Clone].
    ///
    /// All elements of the array starting from [Array::size] + `n` will be uninitialized.
    ///
    /// # Panics
    ///
    /// This function will panic if the reallocation of the underlying data fails.
    pub fn grow(&mut self, n: usize) {
        let new_size = self.size() + n;

        unsafe {
            self.ptr = alloc::realloc(self.ptr as _, self.layout, new_size) as _;

            if self.ptr.is_null() {
                panic!("Reallocation failed");
            }
        }

        self.size = new_size;
    }

    /// Create a new [Array] of length `n`
    ///
    /// `n` should not be less than the length of the array.
    pub fn resized(&self, n: usize) -> Array<T> {
        let new_size = self.size() + n;

        let mut a = Array::new(new_size);

        for i in 0..new_size {
            a[i] = self.get(i).cloned().unwrap();
        }

        a
    }
}

impl<T> Drop for Array<T> {
    fn drop(&mut self) {
        unsafe {
            self.ptr.drop_in_place();
            alloc::dealloc(self.ptr as *mut u8, self.layout);
        }
    }
}

impl<T> Deref for Array<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        // SAFETY: the length is known in the struct so this wont result in UB
        unsafe { core::slice::from_raw_parts(self.ptr, self.size) }
    }
}

impl<T> DerefMut for Array<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: the length is known in the struct so this wont result in UB
        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.size) }
    }
}

impl<T> AsRef<[T]> for Array<T> {
    fn as_ref(&self) -> &[T] {
        self
    }
}

impl<T> AsMut<[T]> for Array<T> {
    fn as_mut(&mut self) -> &mut [T] {
        self
    }
}

impl<T> From<Vec<T>> for Array<T> {
    fn from(mut value: Vec<T>) -> Self {
        let mut v = Self::new(value.len());
        v.swap_with_slice(&mut value);
        v
    }
}

/// [std::io::Write] is implemented for any [`Array<u8>`]
impl std::io::Write for Array<u8> {
    // TODO: is this implementation even good?
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let start = self.ptr;

        unsafe {
            for i in 0..self.size {
                if i > buf.len() {
                    break;
                }
                *self.ptr = buf[i];
                self.ptr = self.ptr.add(1);
            }
        }
        self.ptr = start;

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

impl<T> Index<usize> for Array<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("index out of bounds")
    }
}

impl<T> IndexMut<usize> for Array<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.get_mut(index).expect("index out of bounds")
    }
}

impl<T: Clone> Clone for Array<T> {
    fn clone(&self) -> Self {
        let mut array = Array::new(self.size);

        for (i, v) in self.iter().enumerate() {
            array[i] = v.clone();
        }

        array
    }
}

impl<T> IntoIterator for Array<T> {
    type Item = T;
    type IntoIter = iter::IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        iter::IntoIter::new(self)
    }
}

impl<'a, T> IntoIterator for &'a Array<T> {
    type Item = &'a T;
    type IntoIter = iter::Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a, T> IntoIterator for &'a mut Array<T> {
    type Item = &'a mut T;
    type IntoIter = iter::IterMut<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}

#[cfg(feature = "serde")]
impl<T: serde::Serialize> serde::Serialize for Array<T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        (&self).serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<'de, T: serde::Deserialize<'de>> serde::Deserialize<'de> for Array<T> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let v = Vec::<T>::deserialize(deserializer)?;
        Ok(Array::<T>::from(v))
    }
}

/// Iterators for [Array]
pub mod iter {
    use super::Array;
    use core::marker::PhantomData;

    /// Immutable [Array] iterator.
    pub struct Iter<'a, T> {
        _marker: PhantomData<&'a T>,
        ptr: *const T,
        end: *const T,
    }

    impl<'a, T> Iter<'a, T> {
        pub(crate) fn new(array: &'a Array<T>) -> Self {
            let ptr = array.ptr;
            Self {
                _marker: PhantomData,
                ptr,
                end: unsafe { ptr.add(array.size) },
            }
        }
    }

    impl<'a, T> Iterator for Iter<'a, T> {
        type Item = &'a T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(&*ptr)
                }
            }
        }
    }

    /// Mutable [Array] iterator
    pub struct IterMut<'a, T> {
        _marker: PhantomData<&'a T>,
        ptr: *mut T,
        end: *mut T,
    }

    impl<'a, T> IterMut<'a, T> {
        pub(crate) fn new(array: &'a Array<T>) -> Self {
            let ptr = array.ptr;
            Self {
                _marker: PhantomData,
                ptr,
                end: unsafe { ptr.add(array.size) },
            }
        }
    }

    impl<'a, T> Iterator for IterMut<'a, T> {
        type Item = &'a mut T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(&mut *ptr)
                }
            }
        }
    }

    /// Owned [Array] Iterator
    pub struct IntoIter<T> {
        _array: Array<T>,
        ptr: *const T,
        end: *const T,
    }

    impl<T> IntoIter<T> {
        pub(crate) fn new(array: Array<T>) -> Self {
            unsafe {
                let ptr = array.ptr.cast_const();
                let end = ptr.add(array.size);
                Self {
                    _array: array,
                    ptr,
                    end,
                }
            }
        }
    }

    impl<T> Iterator for IntoIter<T> {
        type Item = T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(ptr.read())
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::Read;

    use super::*;

    #[test]
    fn index_and_iter() {
        let mut array = Array::new(5);

        array[0] = 1;
        array[1] = 2;
        array[2] = 3;
        array[3] = 4;
        array[4] = 5;

        for (i, v) in array.iter().enumerate() {
            assert_eq!(*v, i + 1);
        }
    }

    #[test]
    fn len() {
        let array: Array<()> = Array::new(15);

        assert_eq!(array.size(), 15);
    }

    #[test]
    fn grow() {
        let mut array: Array<i32> = Array::new(5);

        array[0] = 1;
        array[1] = 2;
        array[2] = 3;

        assert_eq!(array.size(), 5);
        array.grow(30);
        assert_eq!(array.size(), 35);

        assert_eq!(array[1], 2);
    }

    #[test]
    fn writer() {
        let array = {
            let mut a = Array::<u8>::new(12);
            a.fill(b"Hello world!");
            a
        };

        let mut cursor = std::io::Cursor::new(array);

        let mut buf = [0u8; 5];
        let read = cursor.read(&mut buf).unwrap();
        assert!(read == 5);
        let s = str::from_utf8(&buf).unwrap();
        println!("{}", s);
    }
}