Do some checked SoA stuff

crates/soa/src/either_index.rs

@@ -0,0 +1,96 @@
+use core::{
+    fmt::{self, Formatter},
+    marker::PhantomData,
+    ptr::NonNull,
+};
+
+use crate::soa_index::Index;
+
+#[derive(PartialEq, Eq)]
+pub struct EitherIndex<T, U> {
+    index: u32,
+    _marker: PhantomData<(T, U)>,
+
+    #[cfg(debug_assertions)]
+    pub(crate) array_start: Result<Option<NonNull<T>>, Option<NonNull<U>>>,
+}
+
+impl<T, U> Clone for EitherIndex<T, U> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T, U> Copy for EitherIndex<T, U> {}
+
+impl<T, U> fmt::Debug for EitherIndex<T, U> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "Index({})", self.index)
+    }
+}
+
+impl<T, U> EitherIndex<T, U> {
+    const MASK: u32 = 1 << 31;
+
+    pub const fn from_left(input: Index<T>) -> Self {
+        assert!(input.index & Self::MASK == 0);
+
+        Self {
+            index: input.index,
+            _marker: PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: Ok(input.array_start),
+        }
+    }
+
+    pub const fn from_right(input: Index<U>) -> Self {
+        assert!(input.index & Self::MASK == 0);
+
+        Self {
+            index: input.index | Self::MASK,
+            _marker: std::marker::PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: Err(input.array_start),
+        }
+    }
+
+    pub const fn split(self) -> Result<Index<T>, Index<U>> {
+        if self.index & Self::MASK == 0 {
+            Ok(Index {
+                index: self.index,
+                _marker: PhantomData,
+
+                #[cfg(debug_assertions)]
+                array_start: {
+                    match self.array_start {
+                        Ok(array) => array,
+                        Err(_) => {
+                            panic!("Tried to split an EitherIndex that was created with from_right, but ended up with a mask that indicates it was created with from_left");
+                        }
+                    }
+                },
+            })
+        } else {
+            Err(Index {
+                index: self.index ^ Self::MASK,
+                _marker: PhantomData,
+
+                #[cfg(debug_assertions)]
+                array_start: {
+                    match self.array_start {
+                        Err(array) => array,
+                        Ok(_) => {
+                            panic!("Tried to split an EitherIndex that was created with from_left, but ended up with a mask that indicates it was created with from_right");
+                        }
+                    }
+                },
+            })
+        }
+    }
+
+    pub fn decrement_index(&mut self) {
+        self.index = self.index.saturating_sub(1);
+    }
+}

crates/soa/src/lib.rs

@@ -1,5 +1,9 @@
+#![cfg_attr(not(any(debug_assertions, test)), no_std)]
+
+mod either_index;
 mod soa_index;
 mod soa_slice;
 
+pub use either_index::*;
 pub use soa_index::*;
 pub use soa_slice::*;

crates/soa/src/soa_index.rs

@@ -1,4 +1,11 @@
-use core::fmt;
+use core::{
+    any,
+    cmp::Ordering,
+    fmt::{self, Formatter},
+    hash::{Hash, Hasher},
+    marker::PhantomData,
+    ptr::NonNull,
+};
 
 use crate::soa_slice::Slice;
 
@@ -7,50 +14,174 @@ use crate::soa_slice::Slice;
 ///
 /// Unlike a Rust pointer, this is a u32 offset
 /// rather than usize.
-pub struct Index<Array, Elem> {
-    pub index: u32,
-    pub(crate) _marker: core::marker::PhantomData<(Array, Elem)>,
+pub struct Index<T> {
+    pub(crate) index: u32,
+    pub(crate) _marker: PhantomData<T>,
+
+    #[cfg(debug_assertions)]
+    pub(crate) array_start: Option<NonNull<T>>,
 }
 
-impl<Array, Elem> fmt::Debug for Index<Array, Elem> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(
-            f,
-            "Index<{}, {}>({})",
-            core::any::type_name::<Array>(),
-            core::any::type_name::<Elem>(),
-            self.index
-        )
+impl<T> PartialEq for Index<T> {
+    fn eq(&self, other: &Self) -> bool {
+        #[cfg(debug_assertions)]
+        self.verify_array(other.array_start.map(NonNull::as_ptr), ".eq()");
+
+        self.index == other.index
+    }
+}
+
+impl<T> Eq for Index<T> {}
+
+impl<T> PartialOrd for Index<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        #[cfg(debug_assertions)]
+        self.verify_array(other.array_start.map(NonNull::as_ptr), ".partial_cmp()");
+
+        Some(self.cmp(other))
+    }
+}
+
+impl<T> Ord for Index<T> {
+    fn cmp(&self, other: &Self) -> Ordering {
+        #[cfg(debug_assertions)]
+        self.verify_array(other.array_start.map(NonNull::as_ptr), ".cmp()");
+
+        self.index.cmp(&other.index)
+    }
+}
+
+impl<T> fmt::Debug for Index<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "Index<{}>({})", any::type_name::<T>(), self.index)
     }
 }
 
 // derive of copy and clone does not play well with PhantomData
 
-impl<Array, Elem> Copy for Index<Array, Elem> {}
+impl<T> Copy for Index<T> {}
 
-impl<Array, Elem> Clone for Index<Array, Elem> {
+impl<T> Clone for Index<T> {
     fn clone(&self) -> Self {
         *self
     }
 }
 
-impl<Array, Elem> Index<Array, Elem> {
-    pub const fn new(start: u32) -> Self {
+impl<T> Hash for Index<T> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.index.hash(state);
+    }
+}
+
+impl<T> Index<T> {
+    pub const fn new(
+        start: u32,
+        _array: &[T], // only used in debug builds
+    ) -> Self {
         Self {
             index: start,
-            _marker: core::marker::PhantomData,
+            _marker: PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: Some(
+                // Safety: this is safe beccause references are never to null pointers.
+                // Once NonNull::new is const on stable Rust, use that instead. Docs to check for const-ness:
+                // https://doc.rust-lang.org/std/ptr/struct.NonNull.html#method.new
+                unsafe { NonNull::new_unchecked(_array.as_ptr() as *mut T) },
+            ),
         }
     }
 
-    pub fn push_new(vector: &mut Vec<Elem>, value: Elem) -> Self {
-        let index = Self::new(vector.len() as _);
+    /// Create a new index that isn't associated with any particular array.
+    /// This is marked as unsafe because it omits the runtime checks (in debug builds)
+    /// which verify in debug builds that indices are dereferened into the original array.
+    pub const unsafe fn new_unchecked(start: u32) -> Self {
+        Self {
+            index: start,
+            _marker: PhantomData,
 
-        vector.push(value);
-
-        index
+            #[cfg(debug_assertions)]
+            array_start: None,
+        }
     }
 
-    pub const fn as_slice(self) -> Slice<Array, Elem> {
-        Slice::new(self.index, 1)
+    pub const fn as_slice(self) -> Slice<T> {
+        Slice {
+            start: self.index,
+            length: 1,
+            _marker: PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: self.array_start,
+        }
+    }
+
+    /// Given the original slice this indexes into, return the equivalent of slice.get_unchecked(index)
+    /// (except that in debug build it actually does do a bounds check, and panics if it's out of bounds,
+    /// but that should never happen since this is an Index into a known slice - basically, a reference).
+    pub fn get_in(self, slice: &[T]) -> &T {
+        #[cfg(debug_assertions)]
+        {
+            self.verify_array(Some(slice.as_ptr() as *mut T), "get_in");
+            slice.get(self.index as usize).expect("Accessing an Index in its original slice went out of bounds. This should never happen!")
+        }
+
+        #[cfg(not(debug_assertions))]
+        unsafe {
+            slice.get_unchecked(self.index as usize)
+        }
+    }
+
+    /// Given the original slice this indexes into, return the equivalent of slice.get_unchecked_mut(index)
+    /// (except that in debug build it actually does do a bounds check, and panics if it's out of bounds,
+    /// but that should never happen since this is an Index into a known slice - basically, a reference).
+    pub fn get_in_mut(self, slice: &mut [T]) -> &mut T {
+        #[cfg(debug_assertions)]
+        {
+            self.verify_array(Some(slice.as_ptr() as *mut T), "get_in");
+            slice.get_mut(self.index as usize).expect("Accessing an Index in its original slice went out of bounds. This should never happen!")
+        }
+
+        #[cfg(not(debug_assertions))]
+        unsafe {
+            slice.get_unchecked_mut(self.index as usize)
+        }
+    }
+
+    /// This is unsafe because it doesn't verify that the index being returned is being used with the original
+    /// slice it was created with. Self::get_in is the safe alternative to this.
+    pub const unsafe fn index(self) -> usize {
+        self.index as usize
+    }
+
+    #[cfg(debug_assertions)]
+    fn verify_array(&self, other: Option<*mut T>, operation: &'static str) {
+        match (self.array_start, other) {
+            (Some(array_start), Some(other_array_start))
+                if other_array_start != array_start.as_ptr() =>
+            {
+                panic!(
+                    "Tried to {} an index from one array to the index of a different array!",
+                    operation
+                );
+            }
+            _ => {
+                // Either one of them was None, meaning this is unchecked, or they point to the same memory - as expected!
+            }
+        }
+    }
+}
+
+impl<'a, T> core::ops::Index<Index<T>> for [T] {
+    type Output = T;
+
+    fn index(&self, index: Index<T>) -> &Self::Output {
+        index.get_in(self)
+    }
+}
+
+impl<'a, T> core::ops::IndexMut<Index<T>> for [T] {
+    fn index_mut(&mut self, index: Index<T>) -> &mut Self::Output {
+        index.get_in_mut(self)
     }
 }

crates/soa/src/soa_slice.rs

@@ -1,5 +1,4 @@
-use core::fmt;
-use core::iter::Map;
+use core::{fmt, marker::PhantomData, ops::Range, ptr::NonNull};
 
 use crate::soa_index::Index;
 
@@ -8,19 +7,22 @@ use crate::soa_index::Index;
 ///
 /// Unlike a Rust slice, this is a u32 offset
 /// rather than a pointer, and the length is u16.
-pub struct Slice<Array, Elem> {
-    pub start: u32,
-    pub length: u16,
-    _marker: core::marker::PhantomData<(Array, Elem)>,
+#[derive(PartialEq, Eq, PartialOrd, Ord)]
+pub struct Slice<T> {
+    pub(crate) start: u32,
+    pub(crate) length: u16,
+    pub(crate) _marker: core::marker::PhantomData<T>,
+
+    #[cfg(debug_assertions)]
+    pub(crate) array_start: Option<NonNull<T>>,
 }
 
-impl<T, U> fmt::Debug for Slice<T, U> {
+impl<T> fmt::Debug for Slice<T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "Slice<{}, {}> {{ start: {}, length: {} }}",
+            "Slice<{}> {{ start: {}, length: {} }}",
             core::any::type_name::<T>(),
-            core::any::type_name::<U>(),
             self.start,
             self.length
         )
@@ -29,42 +31,89 @@ impl<T, U> fmt::Debug for Slice<T, U> {
 
 // derive of copy and clone does not play well with PhantomData
 
-impl<T, U> Copy for Slice<T, U> {}
+impl<T> Copy for Slice<T> {}
 
-impl<T, U> Clone for Slice<T, U> {
+impl<T> Clone for Slice<T> {
     fn clone(&self) -> Self {
         *self
     }
 }
 
-impl<T, U> Default for Slice<T, U> {
-    fn default() -> Self {
-        Self::empty()
-    }
-}
-
-impl<Array, Elem> Slice<Array, Elem> {
-    pub fn empty() -> Self {
+impl<T> Slice<T> {
+    pub fn empty(_array: &[T]) -> Self {
         Self {
             start: 0,
             length: 0,
             _marker: Default::default(),
+
+            #[cfg(debug_assertions)]
+            array_start: NonNull::new(_array.as_ptr() as *mut T),
         }
     }
 
-    pub fn get_slice<'a>(&self, slice: &'a [Elem]) -> &'a [Elem] {
-        &slice[self.indices()]
+    /// Create an empty slice that isn't associated with any particular array.
+    /// This is marked as unsafe because it omits the runtime checks (in debug builds)
+    /// which verify that indices made from this slice are compared with other
+    /// indices into the original array.
+    pub unsafe fn empty_unchecked() -> Self {
+        Self {
+            start: 0,
+            length: 0,
+            _marker: Default::default(),
+
+            #[cfg(debug_assertions)]
+            array_start: None,
+        }
     }
 
-    pub fn get_slice_mut<'a>(&self, slice: &'a mut [Elem]) -> &'a mut [Elem] {
-        &mut slice[self.indices()]
+    /// This is unsafe because it doesn't verify that the start index being returned is being used with the original
+    /// slice it was created with. Self::get_in is the safe alternative to this.
+    pub const unsafe fn start(self) -> usize {
+        self.start as usize
     }
 
+    pub fn get_slice<'a>(&self, slice: &'a [T]) -> &'a [T] {
+        &slice[self.as_range()]
+    }
+
+    pub fn get_slice_mut<'a>(&self, slice: &'a mut [T]) -> &'a mut [T] {
+        &mut slice[self.as_range()]
+    }
+
+    /// This is unsafe because it doesn't verify that the start index being returned is being used with the original
+    /// slice it was created with. Self::get_in is the safe alternative to this.
     #[inline(always)]
-    pub fn indices(&self) -> core::ops::Range<usize> {
+    pub unsafe fn indices(&self) -> Range<usize> {
+        self.as_range()
+    }
+
+    /// Given the original &[T] that this Slice<T> was based on,
+    /// return the &[T] representation of this Slice<T>.
+    #[inline(always)]
+    pub fn slice_from<'a>(&self, original: &'a [T]) -> &'a [T] {
+        #[cfg(debug_assertions)]
+        self.verify_array(Some(original.as_ptr() as *mut T), "slice_from");
+
+        &original[self.as_range()]
+    }
+
+    /// The pub version of this is indices() and it's marked as unsafe. This one isn't marked as unsafe,
+    /// for internal convenience.
+    #[inline(always)]
+    fn as_range(&self) -> Range<usize> {
         self.start as usize..(self.start as usize + self.length as usize)
     }
 
+    /// Given the original &mut [T] that this Slice<T> was based on,
+    /// return the &mut [T] representation of this Slice<T>.
+    #[inline(always)]
+    pub fn slice_from_mut<'a>(&self, original: &'a mut [T]) -> &'a mut [T] {
+        #[cfg(debug_assertions)]
+        self.verify_array(Some(original.as_ptr() as *mut T), "slice_from");
+
+        &mut original[self.start as usize..(self.start as usize + self.length as usize)]
+    }
+
     pub const fn len(&self) -> usize {
         self.length as usize
     }
@@ -73,43 +122,130 @@ impl<Array, Elem> Slice<Array, Elem> {
         self.len() == 0
     }
 
-    pub const fn new(start: u32, length: u16) -> Self {
-        Self {
-            start,
-            length,
-            _marker: std::marker::PhantomData,
+    pub fn at_start(&self, array: &[T]) -> Index<T> {
+        self.at(0, array)
+    }
+
+    pub fn at(&self, i: usize, _array: &[T]) -> Index<T> {
+        #[cfg(debug_assertions)]
+        {
+            if !self.matches_stored_array(_array) {
+                panic!("Tried to call .at() on a slice passing an array that was different from the array it was created with!");
+            }
+        }
+
+        Index {
+            index: self.start + i as u32,
+            _marker: PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: self.array_start,
         }
     }
 
-    pub fn extend_new(vec: &mut Vec<Elem>, it: impl IntoIterator<Item = Elem>) -> Self {
-        let start = vec.len();
+    #[cfg(debug_assertions)]
+    fn matches_stored_array(&self, other_array: &[T]) -> bool {
+        match self.array_start {
+            Some(self_array) => other_array.as_ptr() == self_array.as_ptr(),
+            None => true,
+        }
+    }
 
-        vec.extend(it);
+    pub const fn new(start: u32, length: u16, _array: &[T]) -> Self {
+        Self {
+            start,
+            length,
+            _marker: PhantomData,
 
-        let end = vec.len();
+            #[cfg(debug_assertions)]
+            array_start: Some(
+                // Safety: this is safe beccause references are never to null pointers.
+                // Once NonNull::new is const on stable Rust, use that instead. Docs to check for const-ness:
+                // https://doc.rust-lang.org/std/ptr/struct.NonNull.html#method.new
+                unsafe { NonNull::new_unchecked(_array.as_ptr() as *mut T) },
+            ),
+        }
+    }
 
-        Self::new(start as u32, (end - start) as u16)
+    /// Create a new slice that isn't associated with any particular array.
+    /// This is marked as unsafe because it omits the runtime checks (in debug builds)
+    /// which verify in debug builds that indices made from this slice are compared with other
+    /// indices into the original array.
+    pub const unsafe fn new_unchecked(start: u32, length: u16) -> Self {
+        Self {
+            start,
+            length,
+            _marker: PhantomData,
+
+            #[cfg(debug_assertions)]
+            array_start: None,
+        }
+    }
+
+    #[cfg(debug_assertions)]
+    fn verify_array(&self, other: Option<*mut T>, operation: &'static str) {
+        match (self.array_start, other) {
+            (Some(array_start), Some(other_array_start))
+                if other_array_start != array_start.as_ptr() =>
+            {
+                panic!(
+                    "Tried to {} an index from one array to the index of a different array!",
+                    operation
+                );
+            }
+            _ => {
+                // Either one of them was None, meaning this is unchecked, or they point to the same memory - as expected!
+            }
+        }
     }
 }
 
-impl<Array, Elem> IntoIterator for Slice<Array, Elem> {
-    type Item = Index<Array, Elem>;
-
-    #[allow(clippy::type_complexity)]
-    type IntoIter = Map<core::ops::Range<u32>, fn(u32) -> Self::Item>;
+impl<T> IntoIterator for Slice<T> {
+    type Item = Index<T>;
+    type IntoIter = SliceIterator<T>;
 
     fn into_iter(self) -> Self::IntoIter {
-        (self.start..(self.start + self.length as u32)).map(u32_to_index)
+        SliceIterator {
+            slice: self,
+            current: self.start,
+        }
     }
 }
 
-fn u32_to_index<T, U>(i: u32) -> Index<T, U> {
-    Index {
-        index: i,
-        _marker: core::marker::PhantomData,
+pub struct SliceIterator<T> {
+    slice: Slice<T>,
+    current: u32,
+}
+
+impl<T> Iterator for SliceIterator<T> {
+    type Item = Index<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current < self.slice.start + self.slice.length as u32 {
+            let index = Index {
+                index: self.current,
+                _marker: PhantomData,
+
+                #[cfg(debug_assertions)]
+                array_start: self.slice.array_start,
+            };
+
+            self.current += 1;
+
+            Some(index)
+        } else {
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = (self.slice.start + self.slice.length as u32 - self.current) as usize;
+        (remaining, Some(remaining))
     }
 }
 
-pub trait GetSlice<Array, Elem> {
-    fn get_slice(&self, slice: Slice<Array, Elem>) -> &[Elem];
+impl<T> ExactSizeIterator for SliceIterator<T> {}
+
+pub trait GetSlice<T> {
+    fn get_slice(&self, slice: Slice<T>) -> &[T];
 }