feat: add remove fn

2 files changed, 244 insertions, 16 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 885f364..8fd5ee4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,7 +5,7 @@ use std::cmp::max;
 use std::mem::MaybeUninit;
 use std::ptr::NonNull;
 
-use crate::util::MaybeUninitByteSlice;
+use crate::util::{AnonUnique, Either, MaybeUninitByteSlice};
 
 mod util;
 
@@ -196,6 +196,70 @@ impl MultiVec
         self.length = 1;
     }
 
+    pub fn remove(&mut self, item_index: usize) -> Option<Vec<Field<'static, '_>>>
+    {
+        if item_index >= self.length {
+            return None;
+        }
+
+        let prev_length = self.length;
+
+        self.length -= 1;
+
+        let mut item_fields = Vec::<Field<'static, '_>>::with_capacity(self.field_cnt());
+
+        for field_index in 0..self.field_cnt() {
+            let field_metadata = &self.get_field_metadata()?.get(field_index)?;
+
+            let field_arr_ptr = unsafe { self.ptr.byte_add(field_metadata.array_offset) };
+
+            let bytes = unsafe {
+                std::slice::from_raw_parts_mut(
+                    field_arr_ptr.as_ptr().cast(),
+                    prev_length * field_metadata.size,
+                )
+            };
+
+            let field_slice = FieldSliceMut {
+                bytes,
+                len: prev_length,
+                field_metadata,
+            };
+
+            let mut item_field_bytes = AnonUnique::new_uninit(
+                Layout::from_size_align(field_metadata.size, field_metadata.alignment)
+                    .unwrap(),
+            );
+
+            item_field_bytes.as_slice_mut().copy_from_slice(
+                field_slice
+                    .get_item(item_index)
+                    .expect("Not possible")
+                    .bytes
+                    .as_a()
+                    .expect("Not possible"),
+            );
+
+            let item_field = Field {
+                bytes: Either::B(item_field_bytes),
+                field_metadata,
+            };
+
+            item_fields.push(item_field);
+
+            let item_field_offset = field_slice
+                .get_item_offset(item_index)
+                .expect("Not possible");
+
+            field_slice.bytes.copy_within(
+                item_field_offset + field_slice.field_metadata.size..,
+                item_field_offset,
+            );
+        }
+
+        Some(item_fields)
+    }
+
     ///// Returns a field of the item with the given index.
     /////
     ///// This function is equivalant to doing `.get_all().get(index)`
@@ -398,6 +462,7 @@ impl MultiVec
 
             uninit_field_metadata.write(FieldMetadata {
                 size: item_field.size(),
+                alignment: item_field.alignment(),
                 type_id: item_field.id(),
                 array_offset: field_arr_byte_offsets[field_index],
                 drop_in_place: item_field.drop_in_place,
@@ -620,7 +685,7 @@ pub struct FieldSliceIter<'mv>
 
 impl<'mv> Iterator for FieldSliceIter<'mv>
 {
-    type Item = Field<'mv>;
+    type Item = Field<'mv, 'mv>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
@@ -637,26 +702,43 @@ impl<'mv> Iterator for FieldSliceIter<'mv>
         self.index += 1;
 
         Some(Field {
-            bytes: field_bytes,
+            bytes: Either::A(field_bytes),
             field_metadata: self.field_metadata,
         })
     }
 }
 
-pub struct Field<'mv>
+pub struct Field<'bytes, 'mv>
 {
-    bytes: &'mv [MaybeUninit<u8>],
+    bytes: Either<&'bytes [MaybeUninit<u8>], AnonUnique>,
     field_metadata: &'mv FieldMetadata,
 }
 
-impl Field<'_>
+impl<'bytes, 'mv> Field<'bytes, 'mv>
 {
-    pub fn cast<T: 'static>(&mut self) -> &T
+    pub fn cast<T: 'static>(&self) -> &T
     {
         assert_eq!(TypeId::of::<T>(), self.field_metadata.type_id);
 
-        unsafe { &*self.bytes.as_ptr().cast::<T>() }
+        // println!("Piss: {:?}", unsafe { self.bytes.as_ref().cast::<u8>() });
+
+        let bytes = match &self.bytes {
+            Either::A(bytes) => bytes,
+            Either::B(bytes) => bytes.as_slice(),
+        };
+
+        let ptr: *const MaybeUninit<u8> = bytes.as_ptr();
+
+        unsafe { &*ptr.cast::<T>() }
     }
+
+    // pub unsafe fn into_with_owned_bytes(self) -> Field<'static, 'mv>
+    // {
+    //     Field {
+    //         bytes: self.bytes.into_owned().into(),
+    //         field_metadata: self.field_metadata,
+    //     }
+    // }
 }
 
 pub struct FieldSliceMut<'mv>
@@ -666,11 +748,31 @@ pub struct FieldSliceMut<'mv>
     field_metadata: &'mv FieldMetadata,
 }
 
-impl FieldSliceMut<'_>
+impl<'mv> FieldSliceMut<'mv>
 {
+    pub fn get_item<'this>(&'this self, item_index: usize) -> Option<Field<'mv, 'mv>>
+    where
+        'this: 'mv,
+    {
+        let start_off = self.get_item_offset(item_index)?;
+
+        if item_index >= self.len {
+            return None;
+        }
+
+        let field_bytes = self
+            .bytes
+            .get(start_off..start_off + self.field_metadata.size)?;
+
+        Some(Field {
+            bytes: Either::A(field_bytes),
+            field_metadata: self.field_metadata,
+        })
+    }
+
     pub fn get_item_mut(&mut self, item_index: usize) -> Option<FieldMut<'_>>
     {
-        let start_off = item_index * self.field_metadata.size;
+        let start_off = self.get_item_offset(item_index)?;
 
         if item_index >= self.len {
             return None;
@@ -695,6 +797,15 @@ impl FieldSliceMut<'_>
             field_metadata: self.field_metadata,
         }
     }
+
+    fn get_item_offset(&self, item_index: usize) -> Option<usize>
+    {
+        if item_index >= self.len {
+            return None;
+        }
+
+        Some(item_index * self.field_metadata.size)
+    }
 }
 
 pub struct FieldSliceIterMut<'mv>
@@ -757,6 +868,7 @@ impl FieldMut<'_>
 struct FieldMetadata
 {
     size: usize,
+    alignment: usize,
     type_id: TypeId,
     array_offset: usize,
     drop_in_place: unsafe fn(NonNull<MaybeUninit<u8>>),
@@ -819,7 +931,7 @@ struct CoolLayoutError;
 mod tests
 {
     use std::any::TypeId;
-    use std::mem::offset_of;
+    use std::mem::{offset_of, ManuallyDrop};
     use std::ptr::NonNull;
     use std::sync::atomic::{AtomicUsize, Ordering};
 
@@ -862,11 +974,12 @@ mod tests
             assert_eq!(offset_of!(Data, field_cnt), FIELD_COUNT_OFFSET);
             assert_eq!(offset_of!(Data, field_metadata), FIELD_METADATA_ARRAY_OFFSET);
 
-            $data = Data {
+            $data = ManuallyDrop::new(Data {
                 field_cnt: count_rep!($($field_name,)*),
                 field_metadata: [$(
                     FieldMetadata {
                         size: size_of::<$field_type>(),
+                        alignment: align_of::<$field_type>(),
                         type_id: TypeId::of::<$field_type>(),
                         array_offset: offset_of!(Data, $field_name),
                         drop_in_place: |ptr| unsafe {
@@ -875,14 +988,11 @@ mod tests
                     },
                 )*],
                 $($field_name: $field_values.map(|val| val.into()),)*
-            };
+            });
 
             let mut multi_vec = MultiVec::new();
 
             multi_vec.ptr = NonNull::from(&mut $data).cast();
-
-            std::mem::forget($data);
-
             multi_vec.length = $length;
             multi_vec.capacity = multi_vec.length;
 
@@ -1296,6 +1406,54 @@ mod tests
     }
 
     #[test]
+    fn remove_in_middle_works()
+    {
+        let mut data;
+        let mut multi_vec = multi_vec_with_data!(
+            data = &mut data,
+            {
+                _a: u32 = [u32::MAX - 3000, 901, 5560000, 123123],
+                _b: u16 = [20210u16, 7120, 1010, 9009],
+            },
+            length = 4
+        );
+
+        let removed = multi_vec.remove(1).expect("Should be Some");
+
+        assert_eq!(&data._a[0..3], [u32::MAX - 3000, 5560000, 123123]);
+        assert_eq!(&data._b[0..3], [20210u16, 1010, 9009]);
+
+        assert_eq!(removed.len(), 2);
+
+        assert_eq!(*removed[0].cast::<u32>(), 901);
+        assert_eq!(*removed[1].cast::<u16>(), 7120);
+    }
+
+    #[test]
+    fn remove_at_end_works()
+    {
+        let mut data;
+        let mut multi_vec = multi_vec_with_data!(
+            data = &mut data,
+            {
+                _a: u32 = [u32::MAX - 3000, 901, 5560000, 123123],
+                _b: u16 = [20210u16, 7120, 1010, 9009],
+            },
+            length = 4
+        );
+
+        let removed = multi_vec.remove(3).expect("Should be Some");
+
+        assert_eq!(&data._a[0..3], [u32::MAX - 3000, 901, 5560000]);
+        assert_eq!(&data._b[0..3], [20210u16, 7120, 1010]);
+
+        assert_eq!(removed.len(), 2);
+
+        assert_eq!(*removed[0].cast::<u32>(), 123123);
+        assert_eq!(*removed[1].cast::<u16>(), 9009);
+    }
+
+    #[test]
     fn fields_are_dropped()
     {
         static THING_DROPPED_CNT: AtomicUsize = AtomicUsize::new(0);
diff --git a/src/util.rs b/src/util.rs
index 70e114b..8ff79c7 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,4 +1,6 @@
+use std::alloc::{alloc, dealloc, handle_alloc_error, Layout};
 use std::mem::MaybeUninit;
+use std::ptr::NonNull;
 
 pub trait MaybeUninitByteSlice
 {
@@ -26,3 +28,71 @@ impl MaybeUninitByteSlice for &[MaybeUninit<u8>]
         unsafe { std::slice::from_raw_parts(self.as_ptr().cast::<Item>(), new_len) }
     }
 }
+
+#[derive(Debug)]
+pub struct AnonUnique
+{
+    ptr: NonNull<MaybeUninit<u8>>,
+    layout: Layout,
+}
+
+impl AnonUnique
+{
+    pub fn new_uninit(layout: Layout) -> Self
+    {
+        if layout.size() == 0 {
+            return Self { ptr: NonNull::dangling(), layout };
+        }
+
+        let Some(ptr) = NonNull::new(unsafe { alloc(layout) }) else {
+            handle_alloc_error(layout);
+        };
+
+        Self {
+            ptr: ptr.cast::<MaybeUninit<u8>>(),
+            layout,
+        }
+    }
+
+    pub fn as_slice_mut(&mut self) -> &mut [MaybeUninit<u8>]
+    {
+        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.layout.size()) }
+    }
+
+    pub fn as_slice(&self) -> &[MaybeUninit<u8>]
+    {
+        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.layout.size()) }
+    }
+}
+
+impl Drop for AnonUnique
+{
+    fn drop(&mut self)
+    {
+        if self.layout.size() == 0 {
+            return;
+        }
+
+        unsafe {
+            dealloc(self.ptr.as_ptr().cast(), self.layout);
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum Either<A, B>
+{
+    A(A),
+    B(B),
+}
+
+impl<A, B> Either<A, B>
+{
+    pub fn as_a(&self) -> Option<&A>
+    {
+        match self {
+            Self::A(a) => Some(a),
+            Self::B(_) => None,
+        }
+    }
+}