feat: add get_field function

2 files changed, 154 insertions, 64 deletions

diff --git a/src/lib.rs b/src/lib.rs
index a8d5e82..dedd1de 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,13 +1,18 @@
 #![deny(clippy::all, clippy::pedantic)]
 use std::alloc::{alloc, dealloc, handle_alloc_error, realloc, Layout};
-use std::any::Any;
+use std::any::{Any, TypeId};
 use std::cmp::max;
 use std::mem::MaybeUninit;
 use std::ptr::NonNull;
 
+use crate::util::MaybeUninitByteSlice;
+
+mod util;
+
 pub struct OwnedAnyPtr
 {
     ptr: *mut dyn Any,
+    drop_in_place: unsafe fn(NonNull<MaybeUninit<u8>>),
 }
 
 impl OwnedAnyPtr
@@ -19,7 +24,12 @@ impl OwnedAnyPtr
 
     pub fn from_boxed<Value: Any>(boxed_value: Box<Value>) -> Self
     {
-        Self { ptr: Box::into_raw(boxed_value) }
+        Self {
+            ptr: Box::into_raw(boxed_value),
+            drop_in_place: |ptr| unsafe {
+                std::ptr::drop_in_place(ptr.cast::<Value>().as_ptr())
+            },
+        }
     }
 
     pub fn as_ptr(&self) -> *const dyn Any
@@ -36,6 +46,11 @@ impl OwnedAnyPtr
     {
         align_of_val(unsafe { &*self.ptr })
     }
+
+    pub fn id(&self) -> TypeId
+    {
+        unsafe { &*self.ptr }.type_id()
+    }
 }
 
 impl Drop for OwnedAnyPtr
@@ -82,7 +97,7 @@ pub struct MultiVec
 {
     ptr: NonNull<MaybeUninit<u8>>,
     field_arr_byte_offsets: Vec<usize>,
-    field_sizes: Vec<usize>,
+    field_metadata: Vec<FieldMetadata>,
     length: usize,
     capacity: usize,
     layout: Option<Layout>,
@@ -119,7 +134,7 @@ impl MultiVec
         Self {
             ptr: NonNull::dangling(),
             field_arr_byte_offsets: Vec::new(),
-            field_sizes: Vec::new(),
+            field_metadata: Vec::new(),
             length: 0,
             capacity: 0,
             layout: None,
@@ -170,7 +185,15 @@ impl MultiVec
             return;
         }
 
-        self.field_sizes = fields.as_ref().iter().map(|field| field.size()).collect();
+        self.field_metadata = fields
+            .as_ref()
+            .iter()
+            .map(|field| FieldMetadata {
+                size: field.size(),
+                type_id: field.id(),
+                drop_in_place: field.drop_in_place,
+            })
+            .collect();
 
         self.do_first_alloc(1, &fields);
 
@@ -205,16 +228,25 @@ impl MultiVec
     //    Some(unsafe { field_ptr.cast().as_ref() })
     //}
 
-    ///// Returns a slice containing the specified field of all items.
-    //#[must_use]
-    //pub fn get_all(&self, field_index: usize) -> &[()]
-    //{
-    //    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) };
-    //
-    //    unsafe { std::slice::from_raw_parts(field_arr_ptr.as_ptr().cast(), self.len()) }
-    //}
+    /// Returns a slice containing the specified field of all items.
+    #[must_use]
+    pub fn get_field(&self, field_index: usize) -> FieldSlice<'_>
+    {
+        let field_arr_byte_offset = self.field_arr_byte_offsets[field_index];
+
+        let field_metadata = &self.field_metadata[field_index];
+
+        let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) };
+
+        let bytes = unsafe {
+            std::slice::from_raw_parts(
+                field_arr_ptr.as_ptr().cast(),
+                self.len() * field_metadata.size,
+            )
+        };
+
+        FieldSlice { bytes, field_metadata }
+    }
 
     /// Returns the number of items stored in this `MultiVec`.
     #[must_use]
@@ -394,11 +426,11 @@ impl Drop for MultiVec
 {
     fn drop(&mut self)
     {
-        assert_eq!(self.field_sizes.len(), self.field_arr_byte_offsets.len());
+        assert_eq!(self.field_metadata.len(), self.field_arr_byte_offsets.len());
 
         for index in 0..self.length {
-            for (field_index, field_size) in self.field_sizes.iter().enumerate() {
-                if *field_size == 0 {
+            for (field_index, field_metadata) in self.field_metadata.iter().enumerate() {
+                if field_metadata.size == 0 {
                     continue;
                 }
 
@@ -406,15 +438,11 @@ impl Drop for MultiVec
 
                 let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) };
 
-                let field_ptr = unsafe { field_arr_ptr.add(field_size * index) };
+                let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) };
 
                 unsafe {
-                    field_ptr.drop_in_place();
+                    (field_metadata.drop_in_place)(field_ptr);
                 }
-
-                //unsafe {
-                //    ItemT::drop_field_inplace(field_index, field_ptr.as_ptr());
-                //}
             }
         }
 
@@ -430,6 +458,30 @@ impl Drop for MultiVec
     }
 }
 
+pub struct FieldSlice<'mv>
+{
+    bytes: &'mv [MaybeUninit<u8>],
+    field_metadata: &'mv FieldMetadata,
+}
+
+impl FieldSlice<'_>
+{
+    pub fn as_slice<Item: 'static>(&self) -> &[Item]
+    {
+        assert_eq!(TypeId::of::<Item>(), self.field_metadata.type_id);
+
+        unsafe { self.bytes.cast::<Item>() }
+    }
+}
+
+#[derive(Debug)]
+struct FieldMetadata
+{
+    size: usize,
+    type_id: TypeId,
+    drop_in_place: unsafe fn(NonNull<MaybeUninit<u8>>),
+}
+
 #[inline]
 const fn array_layout(
     element_size: usize,
@@ -486,7 +538,10 @@ struct CoolLayoutError;
 #[cfg(test)]
 mod tests
 {
-    use crate::{MultiVec, OwnedAnyPtr};
+    use std::any::TypeId;
+    use std::ptr::NonNull;
+
+    use crate::{FieldMetadata, MultiVec, OwnedAnyPtr};
 
     #[test]
     fn single_push_works()
@@ -661,43 +716,47 @@ mod tests
     //    );
     //}
 
-    //unsafe fn cast_slice<SrcItem, DstItem>(slice: &[SrcItem]) -> &[DstItem]
-    //{
-    //    unsafe {
-    //        std::slice::from_raw_parts(slice.as_ptr().cast::<DstItem>(), slice.len())
-    //    }
-    //}
-    //
-    //#[test]
-    //fn get_all_works()
-    //{
-    //    let mut multi_vec = MultiVec::new();
-    //
-    //    struct Data
-    //    {
-    //        _a: [u32; 3],
-    //        _b: [u16; 3],
-    //    }
-    //
-    //    let data = Data {
-    //        _a: [u32::MAX - 3000, 901, 5560000],
-    //        _b: [20210, 7120, 1010],
-    //    };
-    //
-    //    multi_vec.ptr = NonNull::from(&data).cast();
-    //    multi_vec.field_arr_byte_offsets = vec![0, size_of::<u32>() * 3];
-    //    multi_vec.field_sizes = vec![size_of::<u32>(), size_of::<u16>()];
-    //    multi_vec.length = 3;
-    //    multi_vec.capacity = 3;
-    //
-    //    assert_eq!(
-    //        unsafe { cast_slice::<_, u32>(multi_vec.get_all(0)) },
-    //        [u32::MAX - 3000, 901, 5560000]
-    //    );
-    //
-    //    assert_eq!(
-    //        unsafe { cast_slice::<_, u16>(multi_vec.get_all(1)) },
-    //        [20210, 7120, 1010]
-    //    );
-    //}
+    #[test]
+    fn get_field_works()
+    {
+        struct Data
+        {
+            _a: [u32; 3],
+            _b: [u16; 3],
+        }
+
+        let mut data = Data {
+            _a: [u32::MAX - 3000, 901, 5560000],
+            _b: [20210, 7120, 1010],
+        };
+
+        let mut multi_vec = MultiVec::new();
+
+        multi_vec.ptr = NonNull::from(&mut data).cast();
+        multi_vec.field_arr_byte_offsets = vec![0, size_of::<u32>() * 3];
+        multi_vec.field_metadata = vec![
+            FieldMetadata {
+                size: size_of::<u32>(),
+                type_id: TypeId::of::<u32>(),
+                drop_in_place: |_| {},
+            },
+            FieldMetadata {
+                size: size_of::<u16>(),
+                type_id: TypeId::of::<u16>(),
+                drop_in_place: |_| {},
+            },
+        ];
+        multi_vec.length = 3;
+        multi_vec.capacity = 3;
+
+        assert_eq!(
+            multi_vec.get_field(0).as_slice::<u32>(),
+            [u32::MAX - 3000, 901, 5560000]
+        );
+
+        assert_eq!(
+            multi_vec.get_field(1).as_slice::<u16>(),
+            [20210, 7120, 1010]
+        );
+    }
 }
diff --git a/src/util.rs b/src/util.rs
new file mode 100644
index 0000000..6b7180a
--- /dev/null
+++ b/src/util.rs
@@ -0,0 +1,31 @@
+use std::mem::MaybeUninit;
+
+pub trait MaybeUninitByteSlice
+{
+    unsafe fn cast<Item: 'static>(&self) -> &[Item];
+}
+
+impl MaybeUninitByteSlice for &[MaybeUninit<u8>]
+{
+    unsafe fn cast<Item: 'static>(&self) -> &[Item]
+    {
+        assert_eq!(self.len() % size_of::<Item>(), 0, "Invalid item size");
+
+        assert_eq!(
+            self.as_ptr().addr() % align_of::<Item>(),
+            0,
+            "Invalid item alignment"
+        );
+
+        if size_of::<Item>() == 0 {
+            return &[];
+        }
+
+        unsafe {
+            std::slice::from_raw_parts(
+                self.as_ptr().cast::<Item>(),
+                self.len() / size_of::<Item>(),
+            )
+        }
+    }
+}