multi-array-list 0.2.1

//! A `MultiArrayList` stores a list of a struct.
//!
//! **Experimental**: Only a small subset of the array list API is implemented.
//!
//! ---
//!
//! > Instead of storing a single list of items, `MultiArrayList` stores separate lists for each field of the struct.
//! > This allows for memory savings if the struct has padding,
//! > and also improves cache usage if only some fields are needed for a computation.
//!
//! The primary API for accessing fields is the [`items::<name>`][`MultiArrayList::items()`] function.
//!
//! ---
//! _inspired by [Zig's `MultiArrayList`](https://ziglang.org/documentation/master/std/#std.MultiArrayList)._
//!
//! # Example
//! ```rust
//! #![allow(incomplete_features)]
//! #![feature(generic_const_exprs)]
//! # use multi_array_list::MultiArrayList;
//! struct Pizza {
//!     radius: u32,
//!     toppings: Vec<Topping>,
//! }
//!
//! enum Topping {
//!     Tomato,
//!     Mozzarella,
//!     Anchovies,
//! }
//!
//! let mut order = MultiArrayList::<Pizza>::new();
//!
//! let margherita = Pizza {
//!     radius: 12,
//!     toppings: vec![Topping::Tomato],
//! };
//! order.push(margherita);
//!
//! let napoli = Pizza {
//!     radius: 12,
//!     toppings: vec![Topping::Tomato, Topping::Anchovies],
//! };
//! order.push(napoli);
//!
//! for topping in order.items_mut::<"toppings", Vec<Topping>>() {
//!     topping.push(Topping::Mozzarella);
//! }
//! ```
#![allow(incomplete_features)]
#![feature(adt_const_params)]
#![feature(generic_const_exprs)]
#![feature(type_info)]
#![feature(unsized_const_params)]
#![feature(const_cmp)]
#![feature(const_trait_impl)]
#![warn(missing_docs)]

use std::alloc::{self, Layout};
use std::any::TypeId;
use std::marker::PhantomData;
use std::mem::type_info::{Field, Type, TypeKind};
use std::mem::{self, ManuallyDrop, MaybeUninit};
use std::ptr;
use std::slice;

/// Alignment of various types, imprecise for anything but numbers, chars and bools.
const fn align_of(kind: TypeKind) -> usize {
    use TypeKind::*;
    match kind {
        Int(int) => (int.bits / 8) as usize,
        Float(float) => (float.bits / 8) as usize,
        Char(_) => 4,
        Bool(_) => 1,
        _ => mem::size_of::<usize>(),
    }
}

/// Number of fields in `T`
pub const fn field_len<T>() -> usize {
    assert_valid_fields::<T>();
    let typeinfo = Type::of::<T>();

    typeinfo.size.unwrap();
    let TypeKind::Struct(struct_info) = typeinfo.kind else {
        panic!("MultiArrayList only works for structs");
    };
    struct_info.fields.len()
}

const fn assert_valid_fields<T>() {
    let typeinfo = Type::of::<T>();

    typeinfo.size.unwrap();
    let TypeKind::Struct(struct_info) = typeinfo.kind else {
        panic!("MultiArrayList only works for structs");
    };

    assert!(struct_info.generics.is_empty());

    let mut i = 0;
    while i < struct_info.fields.len() {
        let field = &struct_info.fields[i];
        let field_info = field.ty.info();

        assert!(field_info.size.is_some(), "known size required");
        assert!(matches!(
            field_info.kind,
            TypeKind::Bool(_)
                | TypeKind::Int(_)
                | TypeKind::Char(_)
                | TypeKind::Float(_)
                | TypeKind::Tuple(_)
                | TypeKind::Struct(_)
                | TypeKind::Enum(_)
        ));

        i += 1;
    }
}

const fn field_sizes<T>() -> [usize; field_len::<T>()] {
    let typeinfo = Type::of::<T>();

    let TypeKind::Struct(struct_info) = typeinfo.kind else {
        panic!("MultiArrayList only works for structs");
    };

    let fields = struct_info.fields;
    let mut new_fields = [0; field_len::<T>()];

    let mut i = 0;
    while i < fields.len() {
        let field = &fields[i];
        let field_info = field.ty.info();
        let element_size = field_info.size.unwrap();
        new_fields[i] = element_size;
        i += 1;
    }

    new_fields
}

const fn field_aligns<T>() -> [usize; field_len::<T>()] {
    let typeinfo = Type::of::<T>();

    let TypeKind::Struct(struct_info) = typeinfo.kind else {
        panic!("MultiArrayList only works for structs");
    };

    let fields = struct_info.fields;
    let mut new_fields = [0; field_len::<T>()];

    let mut i = 0;
    while i < fields.len() {
        let field = &fields[i];
        let align = align_of(field.ty.info().kind);
        new_fields[i] = align;
        i += 1;
    }

    new_fields
}

/// A `MultiArrayList` stores a list of a struct.
///
/// > Instead of storing a single list of items, `MultiArrayList` stores separate lists for each field of the struct.
/// > This allows for memory savings if the struct has padding,
/// > and also improves cache usage if only some fields are needed for a computation.
///
/// The primary API for accessing fields is the [`items(name)`][`MultiArrayList::items()`] function.
#[derive(Debug)]
pub struct MultiArrayList<T>
where
    T: 'static,
    [(); field_len::<T>()]:,
{
    elems: [*mut u8; field_len::<T>()],
    cap: usize,
    len: usize,

    _t: PhantomData<T>,
}

impl<T> MultiArrayList<T>
where
    T: 'static,
    [(); field_len::<T>()]:,
{
    /// Constructs a new, empty `MultiArrayList<T>`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// # use multi_array_list::MultiArrayList;
    /// struct Point {
    ///     x: i32,
    ///     y: i32
    /// }
    /// let mut list: MultiArrayList<Point> = MultiArrayList::new();
    /// ```
    ///
    /// # Type can't be not a struct
    ///
    /// ```compile_fail
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// use multi_array_list::MultiArrayList;
    /// #[repr(u8)]
    /// enum Random {
    ///     Four,
    /// }
    /// let m = MultiArrayList::<Random>::new();
    /// ```
    ///
    /// # Type can't be empty
    ///
    /// ```compile_fail
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// use multi_array_list::MultiArrayList;
    /// let m = MultiArrayList::<()>::new();
    /// ```
    pub fn new() -> MultiArrayList<T> {
        let elems = [ptr::null_mut(); field_len::<T>()];

        MultiArrayList {
            elems,
            cap: 0,
            len: 0,
            _t: PhantomData,
        }
    }

    /// Constructs a new, empty `MultiArrayList<T>` with at least the specified capacity.
    ///
    /// # Examples
    ///
    /// ```rust
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// # use multi_array_list::MultiArrayList;
    /// struct Point {
    ///     x: i32,
    ///     y: i32
    /// }
    /// let mut list: MultiArrayList<Point> = MultiArrayList::with_capacity(10);
    /// ```
    pub fn with_capacity(capacity: usize) -> MultiArrayList<T> {
        let mut list = MultiArrayList::<T>::new();
        list.grow(capacity);
        list
    }

    fn elem_ptrs(&self) -> &[*mut u8] {
        &self.elems[..]
    }

    const fn fields() -> &'static [Field] {
        let typeinfo = Type::of::<T>();

        let TypeKind::Struct(struct_info) = typeinfo.kind else {
            panic!("MultiArrayList only works for structs");
        };

        struct_info.fields
    }

    /// Returns the total number of elements the vector can hold without reallocating.
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Returns the number of elements in the `MultiArrayList`, also referred to as its 'length'.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Appends an element to the back of a collection.
    pub fn push(&mut self, value: T) {
        let len = self.len;

        // We will consume this value, don't drop it.
        let value = ManuallyDrop::new(value);

        if len == self.capacity() {
            self.grow(1);
        }

        let sizes = const { field_sizes::<T>() };

        unsafe {
            let value: *const T = &*value as *const _;
            let value_ptr: *const u8 = value.cast();

            let elem_ptrs = self.elem_ptrs();
            let fields = const { Self::fields() };

            for (idx, (field, elem)) in fields.iter().zip(elem_ptrs).enumerate() {
                let offset = field.offset;
                let size = sizes[idx];

                if size > 0 {
                    let field_elem = elem.offset((size * self.len()) as isize);
                    let src = value_ptr.offset(offset as isize);
                    ptr::copy_nonoverlapping(src, field_elem, size);
                }
            }
        }

        self.len += 1;
    }

    /// Removes the last element from the `MultiArrayList` and returns it, or `None` if it is empty.
    pub fn pop(&mut self) -> Option<Box<T>> {
        if self.len() == 0 {
            return None;
        }
        let idx = self.len() - 1;
        let wip: Box<MaybeUninit<T>> = Box::new(MaybeUninit::uninit());

        let wip = unsafe {
            let elem_ptrs = self.elem_ptrs();
            let wip = Box::into_raw(wip);
            let wip_ptr: *mut u8 = wip.cast();

            let sizes = const { field_sizes::<T>() };

            for (i, (field, elem)) in Self::fields().iter().zip(elem_ptrs).enumerate() {
                let offset = field.offset;
                let size = sizes[i];
                if size > 0 {
                    let field_elem = elem.offset((size * idx) as isize);
                    let dst_ptr = wip_ptr.offset(offset as isize);
                    ptr::copy_nonoverlapping(field_elem, dst_ptr, size);
                }
            }

            Box::from_raw(wip)
        };

        self.len -= 1;

        // SAFETY: We initialized all fields at their correct offset and size
        // based on values we previously copied into our element array.
        unsafe { Some(wip.assume_init()) }
    }

    /// Returns an iterator over the `MultiArrayList`.
    ///
    /// The iterator yields all items from start to end.
    ///
    /// ```
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// use multi_array_list::MultiArrayList;
    ///
    /// struct Point {
    ///     x: i32,
    ///     y: i32,
    /// }
    ///
    /// let mut points = MultiArrayList::<Point>::new();
    /// points.push(Point { x: 1, y: 0 });
    /// for point in points.iter() {
    ///     assert_eq!(1, *point.get::<"x", _>());
    ///     assert_eq!(0, *point.get::<"y", _>());
    /// }
    /// ```
    pub fn iter<'a>(&'a self) -> Iter<'a, T> {
        Iter {
            array: self,
            pos: 0,
        }
    }

    const fn get_field_by_name<const NAME: &'static str, V: 'static>() -> (usize, usize) {
        let mut idx = 0;
        while idx < Self::fields().len() {
            let field = &Self::fields()[idx];
            if field.name.eq_ignore_ascii_case(NAME) {
                break;
            }
            idx += 1
        }

        if idx == Self::fields().len() {
            panic!("unknown item name");
        }

        assert!(idx < Self::fields().len());
        let field = &Self::fields()[idx];
        assert!(TypeId::of::<V>() == field.ty, "types don't match");

        let elem_size = mem::size_of::<V>();
        (idx, elem_size)
    }

    /// Get an iterator of values for a specified field.
    ///
    /// # Compile errors
    ///
    /// * Fails to compile if the requested field does not exist.
    /// * Fails to compile if the requested type does not match the found field (by size).
    ///
    /// ```compile_fail
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// use multi_array_list::MultiArrayList;
    /// struct HereBeDragons {
    ///     x: u32,
    ///     y: (u8, u16),
    /// }
    /// let mut list = MultiArrayList::<HereBeDragons>::new();
    /// _ = list.items::<"y", (u16, u8)>(); // will not compile
    /// ```
    pub fn items<'a, const NAME: &'static str, V: 'static>(&'a self) -> Slice<'a, V> {
        let (idx, elem_size) = const { Self::get_field_by_name::<NAME, V>() };

        unsafe {
            let elem_ptrs = self.elem_ptrs();
            let ptr = elem_ptrs[idx];
            let end = ptr.offset((self.len * elem_size) as isize);

            return Slice {
                ptr,
                end,
                typ: PhantomData,
            };
        }
    }

    /// Get an iterator of mutable values for a specified field.
    ///
    /// # Compile errors
    ///
    /// * Fails to compile if the requested field does not exist.
    /// * Fails to compile if the requested type does not match the found field (by size).
    pub fn items_mut<'a, const NAME: &'static str, V: 'static>(&'a mut self) -> SliceMut<'a, V> {
        let (idx, elem_size) = const { Self::get_field_by_name::<NAME, V>() };

        unsafe {
            let elem_ptrs = self.elem_ptrs();
            let ptr = elem_ptrs[idx];
            let end = ptr.offset((self.len * elem_size) as isize);

            return SliceMut {
                ptr,
                end,
                typ: PhantomData,
            };
        }
    }

    fn grow(&mut self, additional: usize) {
        let old_cap = self.capacity();
        let new_cap = old_cap + additional;

        let sizes = const { field_sizes::<T>() };
        let aligns = const { field_aligns::<T>() };

        for idx in 0..Self::fields().len() {
            let element_size = sizes[idx];
            let align = aligns[idx];
            let old_array_size = element_size * old_cap;
            let new_array_size = element_size * new_cap;

            if new_array_size == 0 {
                continue;
            }

            unsafe {
                let ptr = &mut self.elems[idx];
                if old_cap == 0 {
                    let layout = Layout::from_size_align_unchecked(new_array_size, align);
                    *ptr = alloc::alloc(layout);
                } else {
                    let layout = Layout::from_size_align_unchecked(old_array_size, align);
                    *ptr = alloc::realloc(*ptr, layout, new_array_size);
                }
            }
        }

        self.cap = new_cap;
    }

    /// Shrinks the capacity of the vector as much as possible.
    pub fn shrink_to_fit(&mut self) {
        let cur_cap = self.capacity();
        let len = self.len();
        if len == cur_cap {
            return;
        }

        let sizes = const { field_sizes::<T>() };
        let aligns = const { field_aligns::<T>() };

        unsafe {
            for idx in 0..Self::fields().len() {
                let element_size = sizes[idx];
                let align = aligns[idx];
                let old_array_size = element_size * cur_cap;
                let new_array_size = element_size * len;

                if element_size == 0 {
                    continue;
                }

                let ptr = &mut self.elems[idx];
                let layout = Layout::from_size_align_unchecked(old_array_size, align);
                if new_array_size == 0 {
                    alloc::dealloc(*ptr, layout);
                    *ptr = ptr::null_mut();
                } else {
                    *ptr = alloc::realloc(*ptr, layout, new_array_size);
                }
            }
        }

        self.cap = len;
    }
}

impl<T> Drop for MultiArrayList<T>
where
    T: 'static,
    [(); field_len::<T>()]:,
{
    fn drop(&mut self) {
        while let Some(elem) = self.pop() {
            drop(elem);
        }

        self.shrink_to_fit();
    }
}

/// An iterator over all values of a specified field.
pub struct Slice<'a, V> {
    ptr: *const u8,
    end: *const u8,
    typ: PhantomData<&'a V>,
}

impl<'a, V> Iterator for Slice<'a, V> {
    type Item = &'a V;

    fn next(&mut self) -> Option<Self::Item> {
        if self.ptr >= self.end {
            return None;
        }
        let elem_size = mem::size_of::<V>();

        unsafe {
            let slice: &[u8] = slice::from_raw_parts(self.ptr, elem_size);
            let val = mem::transmute_copy::<&[u8], &V>(&slice);
            self.ptr = self.ptr.offset(elem_size as isize);
            return Some(val);
        }
    }
}

/// An iterator over mutable values of a specified field.
pub struct SliceMut<'a, V> {
    ptr: *mut u8,
    end: *mut u8,
    typ: PhantomData<&'a mut V>,
}

impl<'a, V> Iterator for SliceMut<'a, V> {
    type Item = &'a mut V;

    fn next(&mut self) -> Option<Self::Item> {
        if self.ptr >= self.end {
            return None;
        }
        let elem_size = mem::size_of::<V>();

        unsafe {
            let mut slice: &mut [u8] = slice::from_raw_parts_mut(self.ptr, elem_size);
            let val = mem::transmute_copy::<&mut [u8], &mut V>(&mut slice);
            self.ptr = self.ptr.offset(elem_size as isize);
            return Some(val);
        }
    }
}

/// An iterator over all elements of the `MultiArrayList`.
pub struct Iter<'a, T>
where
    T: 'static,
    [(); field_len::<T>()]:,
{
    array: &'a MultiArrayList<T>,
    pos: usize,
}

impl<'a, T> Iterator for Iter<'a, T>
where
    T: 'static,
    [(); field_len::<T>()]:,
{
    type Item = IterRef<'a, T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos >= self.array.len() {
            return None;
        }

        let elems_slice = &self.array.elems[..];

        self.pos += 1;
        Some(IterRef {
            ptr: elems_slice,
            idx: self.pos - 1,
            elem: PhantomData,
        })
    }
}

/// A reference to a single element of `T`.
///
/// Use [`get(field)`][`IterRef::get`] to access a single field.
pub struct IterRef<'a, T>
where
    T: 'static,
{
    ptr: &'a [*mut u8],
    idx: usize,
    elem: PhantomData<&'a T>,
}

impl<'a, T> IterRef<'a, T>
where
    [(); field_len::<T>()]:,
{
    /// Get a specified field of `T`.
    ///
    /// # Compile errors
    ///
    /// * Fails to compile if the requested field does not exist.
    /// * Fails to compile if the requested type does not match the found field (by size).
    ///
    /// ## It needs the correct type or it fails to compile
    ///
    /// ```compile_fail
    /// #![allow(incomplete_features)]
    /// #![feature(generic_const_exprs)]
    /// use multi_array_list::MultiArrayList;
    ///
    /// struct Point {
    ///     x: i32,
    ///     y: i32,
    /// }
    ///
    /// let mut points = MultiArrayList::<Point>::new();
    /// points.push(Point { x: 1, y: 0 });
    /// for point in points.iter() {
    ///     point.get::<"x", bool>();
    /// }
    /// ```
    pub fn get<const NAME: &'static str, V: 'static>(&self) -> &V {
        let (idx, elem_size) = const { MultiArrayList::<T>::get_field_by_name::<NAME, V>() };

        unsafe {
            let ptr = self.ptr[idx].offset((self.idx * elem_size) as isize);
            let slice: &[u8] = slice::from_raw_parts(ptr, elem_size);
            let val = mem::transmute_copy::<&[u8], &V>(&slice);
            return val;
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::mem;

    #[derive(Debug, PartialEq, Eq)]
    struct Point {
        x: i32,
        y: i32,
    }

    #[repr(u8)]
    #[allow(dead_code)]
    enum Random {
        Four,
    }

    #[test]
    fn size() {
        assert_eq!(32, mem::size_of::<MultiArrayList<Point>>());
    }

    #[test]
    fn empty() {
        let list = MultiArrayList::<Point>::new();
        assert_eq!(0, list.len());
        assert_eq!(0, list.capacity());
    }

    #[test]
    fn push() {
        let mut list = MultiArrayList::<Point>::new();
        assert_eq!(0, list.len());
        assert_eq!(0, list.capacity());

        let p1 = Point { x: 2, y: 8 };
        list.push(p1);
        assert_eq!(1, list.len());
        assert!(0 < list.capacity());
    }

    #[test]
    fn pop() {
        let mut list = MultiArrayList::<Point>::new();
        assert_eq!(0, list.len());
        assert_eq!(0, list.capacity());

        let p1 = Point { x: 2, y: 8 };
        list.push(p1);
        assert_eq!(1, list.len());
        assert!(0 < list.capacity());

        let p2 = list.pop().unwrap();
        assert_eq!(0, list.len());
        assert_eq!(&Point { x: 2, y: 8 }, &*p2);
    }

    #[test]
    fn capacity() {
        let mut list = MultiArrayList::<Point>::with_capacity(10);
        assert_eq!(0, list.len());
        assert_eq!(10, list.capacity());

        let p1 = Point { x: 2, y: 8 };
        list.push(p1);
        assert_eq!(1, list.len());
        assert_eq!(10, list.capacity());
    }

    #[test]
    fn zst() {
        #[allow(dead_code)]
        #[derive(Debug, PartialEq, Eq)]
        struct Empty {
            x: (),
        }

        let mut list = MultiArrayList::<Empty>::new();
        assert_eq!(0, list.len());
        assert_eq!(0, list.capacity());

        list.push(Empty { x: () });

        let elem = list.pop().unwrap();
        assert_eq!(Empty { x: () }, *elem);
    }
}

#[doc = include_str!("../README.md")]
#[cfg(doctest)]
pub struct ReadmeDoctests;