Make specialize_types lower to a new representation

crates/soa/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "soa"
+description = "Struct-of-Array helpers"
+
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+version.workspace = true
+
+[dependencies]
+
+[dev-dependencies]

crates/soa/src/either_index.rs

@@ -0,0 +1,66 @@
+use core::{
+    fmt::{self, Formatter},
+    marker::PhantomData,
+};
+
+use crate::soa_index::Index;
+
+#[derive(PartialEq, Eq)]
+pub struct EitherIndex<T, U> {
+    index: u32,
+    _marker: PhantomData<(T, 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, "EitherIndex({})", 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,
+        }
+    }
+
+    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,
+        }
+    }
+
+    pub const fn split(self) -> Result<Index<T>, Index<U>> {
+        if self.index & Self::MASK == 0 {
+            Ok(Index {
+                index: self.index,
+                _marker: PhantomData,
+            })
+        } else {
+            Err(Index {
+                index: self.index ^ Self::MASK,
+                _marker: PhantomData,
+            })
+        }
+    }
+
+    pub fn decrement_index(&mut self) {
+        self.index = self.index.saturating_sub(1);
+    }
+}

crates/soa/src/lib.rs

@@ -0,0 +1,9 @@
+mod either_index;
+mod soa_index;
+mod soa_slice;
+mod soa_slice2;
+
+pub use either_index::*;
+pub use soa_index::*;
+pub use soa_slice::*;
+pub use soa_slice2::*;

crates/soa/src/soa_index.rs

@@ -0,0 +1,98 @@
+use core::{
+    any,
+    cmp::Ordering,
+    fmt::{self, Formatter},
+    hash::{Hash, Hasher},
+    marker::PhantomData,
+};
+
+use crate::soa_slice::Slice;
+
+pub type Id<T> = Index<T>;
+
+/// An index into an array of values, based
+/// on an offset into the array rather than a pointer.
+///
+/// Unlike a Rust pointer, this is a u32 offset
+/// rather than usize.
+pub struct Index<T> {
+    pub index: u32,
+    pub _marker: PhantomData<T>,
+}
+
+impl<T> PartialEq for Index<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.index == other.index
+    }
+}
+
+impl<T> Eq for Index<T> {}
+
+impl<T> PartialOrd for Index<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<T> Ord for Index<T> {
+    fn cmp(&self, other: &Self) -> Ordering {
+        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<T> Copy for Index<T> {}
+
+impl<T> Clone for Index<T> {
+    fn clone(&self) -> Self {
+        *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) -> Self {
+        Self {
+            index: start,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn as_slice(self) -> Slice<T> {
+        Slice {
+            start: self.index,
+            length: 1,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn index(self) -> usize {
+        self.index as usize
+    }
+}
+
+impl<T> core::ops::Index<Index<T>> for [T] {
+    type Output = T;
+
+    fn index(&self, index: Index<T>) -> &Self::Output {
+        &self[index.index()]
+    }
+}
+
+impl<T> core::ops::IndexMut<Index<T>> for [T] {
+    fn index_mut(&mut self, index: Index<T>) -> &mut Self::Output {
+        &mut self[index.index()]
+    }
+}

crates/soa/src/soa_slice.rs

@@ -0,0 +1,151 @@
+use core::{fmt, marker::PhantomData, ops::Range};
+
+use crate::soa_index::Index;
+
+/// A slice into an array of values, based
+/// on an offset into the array rather than a pointer.
+///
+/// Unlike a Rust slice, this is a u32 offset
+/// rather than a pointer, and the length is u16.
+#[derive(PartialEq, Eq, PartialOrd, Ord)]
+pub struct Slice<T> {
+    pub start: u32,
+    pub length: u16,
+    pub _marker: core::marker::PhantomData<T>,
+}
+
+impl<T> fmt::Debug for Slice<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Slice<{}> {{ start: {}, length: {} }}",
+            core::any::type_name::<T>(),
+            self.start,
+            self.length
+        )
+    }
+}
+
+// derive of copy and clone does not play well with PhantomData
+
+impl<T> Copy for Slice<T> {}
+
+impl<T> Clone for Slice<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T> Default for Slice<T> {
+    fn default() -> Self {
+        Self::empty()
+    }
+}
+
+impl<T> Slice<T> {
+    pub const fn empty() -> Self {
+        Self {
+            start: 0,
+            length: 0,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn start(self) -> u32 {
+        self.start
+    }
+
+    pub fn advance(&mut self, amount: u32) {
+        self.start += amount
+    }
+
+    pub fn get_slice<'a>(&self, elems: &'a [T]) -> &'a [T] {
+        &elems[self.indices()]
+    }
+
+    pub fn get_slice_mut<'a>(&self, elems: &'a mut [T]) -> &'a mut [T] {
+        &mut elems[self.indices()]
+    }
+
+    #[inline(always)]
+    pub const fn indices(&self) -> Range<usize> {
+        self.start as usize..(self.start as usize + self.length as usize)
+    }
+
+    pub const fn len(&self) -> usize {
+        self.length as usize
+    }
+
+    pub const fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    pub fn at_start(&self) -> Index<T> {
+        Index {
+            index: self.start,
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn at(&self, i: usize) -> Index<T> {
+        Index {
+            index: self.start + i as u32,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn new(start: u32, length: u16) -> Self {
+        Self {
+            start,
+            length,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T> IntoIterator for Slice<T> {
+    type Item = Index<T>;
+    type IntoIter = SliceIterator<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SliceIterator {
+            slice: self,
+            current: self.start,
+        }
+    }
+}
+
+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,
+            };
+
+            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))
+    }
+}
+
+impl<T> ExactSizeIterator for SliceIterator<T> {}
+
+pub trait GetSlice<T> {
+    fn get_slice(&self, slice: Slice<T>) -> &[T];
+}

crates/soa/src/soa_slice2.rs

@@ -0,0 +1,139 @@
+use core::{fmt, marker::PhantomData, ops::Range};
+
+use crate::{soa_index::Index, soa_slice::Slice};
+
+/// Two slices of the same length, each based on a different
+/// offset into the same array (rather than a pointer).
+///
+/// Unlike a Rust slice, this is a u32 offset
+/// rather than a pointer, and the length is u16.
+#[derive(PartialEq, Eq, PartialOrd, Ord)]
+pub struct Slice2<T, U> {
+    pub start1: u32,
+    pub start2: u32,
+    pub length: u16,
+    pub _marker: core::marker::PhantomData<(T, U)>,
+}
+
+impl<T, U> fmt::Debug for Slice2<T, U> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Slice2<{}, {}> {{ start1: {}, start2: {}, length: {} }}",
+            core::any::type_name::<T>(),
+            core::any::type_name::<U>(),
+            self.start1,
+            self.start2,
+            self.length
+        )
+    }
+}
+
+// derive of copy and clone does not play well with PhantomData
+
+impl<T, U> Copy for Slice2<T, U> {}
+
+impl<T, U> Clone for Slice2<T, U> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T, U> Default for Slice2<T, U> {
+    fn default() -> Self {
+        Self::empty()
+    }
+}
+
+impl<T, U> Slice2<T, U> {
+    pub const fn empty() -> Self {
+        Self {
+            start1: 0,
+            start2: 0,
+            length: 0,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn slice_first(self) -> Slice<T> {
+        Slice {
+            start: self.start1,
+            length: self.length,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn slice_second(self) -> Slice<U> {
+        Slice {
+            start: self.start2,
+            length: self.length,
+            _marker: PhantomData,
+        }
+    }
+
+    pub const fn len(&self) -> usize {
+        self.length as usize
+    }
+
+    pub const fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    pub const fn new(start1: u32, start2: u32, length: u16) -> Self {
+        Self {
+            start1,
+            start2,
+            length,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T, U> IntoIterator for Slice2<T, U> {
+    type Item = (Index<T>, Index<U>);
+    type IntoIter = SliceIterator<T, U>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SliceIterator {
+            slice: self,
+            offset: 0,
+        }
+    }
+}
+
+pub struct SliceIterator<T, U> {
+    slice: Slice2<T, U>,
+    offset: u32,
+}
+
+impl<T, U> Iterator for SliceIterator<T, U> {
+    type Item = (Index<T>, Index<U>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let offset = self.offset;
+
+        if offset < self.slice.length as u32 {
+            let index1 = Index {
+                index: self.slice.start1 + offset,
+                _marker: PhantomData,
+            };
+            let index2 = Index {
+                index: self.slice.start2 + offset,
+                _marker: PhantomData,
+            };
+
+            self.offset += 1;
+
+            Some((index1, index2))
+        } else {
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = (self.slice.length as u32 - self.offset) as usize;
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T, U> ExactSizeIterator for SliceIterator<T, U> {}