use std::alloc::{alloc, handle_alloc_error, realloc, Layout}; use std::marker::PhantomData; use std::mem::{forget, needs_drop}; use std::ptr::NonNull; /// A list of `ItemT`. This data structure stores a list for every field of `ItemT`, /// reducing memory usage if `ItemT` contains padding and improves memory cache usage if /// only certain fields are needed when iterating. /// /// Inspired by Zig's `MultiArrayList`. /// /// Note: All of the lists are stored in the same allocation. /// /// For example, if you have three of the following struct: /// ``` /// struct Person /// { /// first_name: String, /// age: u8, /// } /// ``` /// /// It would be stored like this in memory: /// ```text /// first_name, first_name, first_name, /// age, age, age, /// ``` #[derive(Debug)] pub struct MultiVec where ItemT: Item, { _pd: PhantomData, ptr: NonNull, field_arr_byte_offsets: Vec, length: usize, layout: Option, } impl MultiVec where ItemT: Item, { pub const fn new() -> Self { Self { _pd: PhantomData, ptr: NonNull::dangling(), field_arr_byte_offsets: Vec::new(), length: 0, layout: None, } } pub fn push(&mut self, item: ItemT) { if self.length != 0 { self.reallocate_for_more(1); self.write_item(self.length - 1, item); return; } self.do_first_alloc(); self.write_item(0, item); } pub fn get( &self, index: usize, ) -> &>::Field where FieldSel: ItemFieldSelection, { if index >= self.length { panic!( "Index {index} is out of bounds in MultiVec with length {}", self.length ); } let field_metadata = FieldSel::metadata(); let field_arr_byte_offset = self.field_arr_byte_offsets[FieldSel::INDEX]; let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) }; let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) }; unsafe { field_ptr.cast().as_ref() } } fn reallocate_for_more(&mut self, item_cnt_inc: usize) { let new_length = self.length + item_cnt_inc; let layout = &self.layout.unwrap(); let (new_layout, new_field_arr_byte_offsets) = Self::create_layout(new_length); let Some(new_ptr) = NonNull::new(unsafe { realloc(self.ptr.as_ptr(), *layout, new_layout.size()) }) else { handle_alloc_error(new_layout); }; for (field_index, field_metadata) in ItemT::iter_field_metadata().enumerate().rev() { let old_field_arr_byte_offset = self.field_arr_byte_offsets[field_index]; let new_field_arr_byte_offset = new_field_arr_byte_offsets[field_index]; let old_field_arr_ptr = unsafe { new_ptr.byte_add(old_field_arr_byte_offset) }; let new_field_arr_ptr = unsafe { new_ptr.byte_add(new_field_arr_byte_offset) }; unsafe { std::ptr::copy( old_field_arr_ptr.as_ptr(), new_field_arr_ptr.as_ptr(), field_metadata.size * self.length, ); } } self.ptr = new_ptr; self.layout = Some(new_layout); self.length = new_length; self.field_arr_byte_offsets = new_field_arr_byte_offsets; } fn do_first_alloc(&mut self) { let (layout, field_arr_byte_offsets) = Self::create_layout(1); let Some(ptr) = NonNull::new(unsafe { alloc(layout) }) else { handle_alloc_error(layout); }; self.ptr = ptr; self.length = 1; self.field_arr_byte_offsets = field_arr_byte_offsets; self.layout = Some(layout); } fn create_layout(length: usize) -> (Layout, Vec) { let mut field_metadata_iter = ItemT::iter_field_metadata(); let first_field_metadata = field_metadata_iter.next().unwrap(); let mut layout = array_layout( first_field_metadata.size, first_field_metadata.alignment, length, ) .unwrap(); let mut field_arr_byte_offsets = Vec::with_capacity(ItemT::FIELD_CNT); field_arr_byte_offsets.push(0); for field_metadata in field_metadata_iter { let (new_layout, array_byte_offset) = layout .extend( array_layout(field_metadata.size, field_metadata.alignment, length) .unwrap(), ) .unwrap(); layout = new_layout; field_arr_byte_offsets.push(array_byte_offset); } (layout, field_arr_byte_offsets) } fn write_item(&mut self, index: usize, item: ItemT) { for (field_index, field_metadata) in ItemT::iter_field_metadata().enumerate() { let field_arr_byte_offset = self.field_arr_byte_offsets[field_index]; let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) }; let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) }; unsafe { std::ptr::copy_nonoverlapping( (&item as *const ItemT).byte_add(field_metadata.offset) as *const u8, field_ptr.as_ptr(), field_metadata.size, ); } } forget(item); } } impl Drop for MultiVec where ItemT: Item, { fn drop(&mut self) { if needs_drop::() { for index in 0..self.length { for (field_index, field_metadata) in ItemT::iter_field_metadata().enumerate() { let field_arr_byte_offset = self.field_arr_byte_offsets[field_index]; let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) }; let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) }; unsafe { ItemT::drop_field_inplace(field_index, field_ptr.as_ptr()); } } } } if let Some(layout) = self.layout { unsafe { std::alloc::dealloc(self.ptr.as_ptr(), layout); } } } } /// Usable as a item of a [`MultiVec`]. /// /// # Safety /// The iterator returned by `iter_field_metadata` must yield [`ItemFieldMetadata`] that /// correctly represents fields of the implementor type. pub unsafe trait Item { type FieldMetadataIter<'a>: Iterator + DoubleEndedIterator + ExactSizeIterator; const FIELD_CNT: usize; fn iter_field_metadata() -> Self::FieldMetadataIter<'static>; unsafe fn drop_field_inplace(field_index: usize, field_ptr: *mut u8); } pub struct ItemFieldMetadata { pub offset: usize, pub size: usize, pub alignment: usize, } /// A field selection for `ItemT`. /// /// # Safety /// The constant `INDEX`, the type `Field` and the `ItemFieldMetadata` returned by the /// `metadata` function must correctly represent a field of `ItemT`; pub unsafe trait ItemFieldSelection where ItemT: Item, { const INDEX: usize; type Field; fn metadata() -> &'static ItemFieldMetadata; } #[inline] const fn array_layout( element_size: usize, align: usize, n: usize, ) -> Result { // We need to check two things about the size: // - That the total size won't overflow a `usize`, and // - That the total size still fits in an `isize`. // By using division we can check them both with a single threshold. // That'd usually be a bad idea, but thankfully here the element size // and alignment are constants, so the compiler will fold all of it. if element_size != 0 && n > max_size_for_align(align) / element_size { return Err(CoolLayoutError); } // SAFETY: We just checked that we won't overflow `usize` when we multiply. // This is a useless hint inside this function, but after inlining this helps // deduplicate checks for whether the overall capacity is zero (e.g., in RawVec's // allocation path) before/after this multiplication. let array_size = unsafe { element_size.unchecked_mul(n) }; // SAFETY: We just checked above that the `array_size` will not // exceed `isize::MAX` even when rounded up to the alignment. // And `Alignment` guarantees it's a power of two. unsafe { Ok(Layout::from_size_align_unchecked(array_size, align)) } } #[inline(always)] const fn max_size_for_align(align: usize) -> usize { // (power-of-two implies align != 0.) // Rounded up size is: // size_rounded_up = (size + align - 1) & !(align - 1); // // We know from above that align != 0. If adding (align - 1) // does not overflow, then rounding up will be fine. // // Conversely, &-masking with !(align - 1) will subtract off // only low-order-bits. Thus if overflow occurs with the sum, // the &-mask cannot subtract enough to undo that overflow. // // Above implies that checking for summation overflow is both // necessary and sufficient. isize::MAX as usize - (align - 1) } #[derive(Debug)] struct CoolLayoutError; #[cfg(test)] mod tests { use std::mem::offset_of; use std::ptr::drop_in_place; use crate::{Item, ItemFieldMetadata, ItemFieldSelection, MultiVec}; struct Foo { num_a: u32, num_b: u16, } struct FooFieldNumA; struct FooFieldNumB; unsafe impl ItemFieldSelection for FooFieldNumA { type Field = u32; const INDEX: usize = 0; fn metadata() -> &'static ItemFieldMetadata { &FOO_FIELD_METADATA[0] } } unsafe impl ItemFieldSelection for FooFieldNumB { type Field = u16; const INDEX: usize = 1; fn metadata() -> &'static ItemFieldMetadata { &FOO_FIELD_METADATA[1] } } static FOO_FIELD_METADATA: [ItemFieldMetadata; 2] = [ ItemFieldMetadata { offset: offset_of!(Foo, num_a), size: size_of::(), alignment: align_of::(), }, ItemFieldMetadata { offset: offset_of!(Foo, num_b), size: size_of::(), alignment: align_of::(), }, ]; unsafe impl Item for Foo { type FieldMetadataIter<'a> = std::slice::Iter<'a, ItemFieldMetadata>; const FIELD_CNT: usize = 2; fn iter_field_metadata() -> Self::FieldMetadataIter<'static> { FOO_FIELD_METADATA.iter() } unsafe fn drop_field_inplace(field_index: usize, field_ptr: *mut u8) { if field_index == 0 { unsafe { drop_in_place::(field_ptr.cast()) } } else if field_index == 1 { unsafe { drop_in_place::(field_ptr.cast()) } } } } #[test] fn works() { let mut multi_vec = MultiVec::::new(); multi_vec.push(Foo { num_a: 123, num_b: 654 }); multi_vec.push(Foo { num_a: 12338, num_b: 191 }); let item = multi_vec.get::(1); println!("yay: {}", *item); } }