Codedump!

crates/bevy_ecs/macros/src/lib.rs

@@ -133,11 +133,10 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 component_ids.push(components.init_component::<#field_type>(storages));
             });
             field_get_components.push(quote! {
-                func((&mut self.#field as *mut #field_type).cast::<u8>());
-                std::mem::forget(self.#field);
+                #ecs_path::ptr::OwningPtr::make(self.#field, func);
             });
             field_from_components.push(quote! {
-                #field: func().cast::<#field_type>().read(),
+                #field: func().inner().cast::<#field_type>().read(),
             });
         }
     }
@@ -159,14 +158,17 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             }
 
             #[allow(unused_variables, unused_mut, non_snake_case)]
-            unsafe fn from_components(mut func: impl FnMut() -> *mut u8) -> Self {
+            unsafe fn from_components<'a, F>(func: F) -> Self
+            where
+                F: FnMut() -> #ecs_path::ptr::OwningPtr<'a> + 'a
+            {
                 Self {
                     #(#field_from_components)*
                 }
             }
 
             #[allow(unused_variables, unused_mut, forget_copy, forget_ref)]
-            fn get_components(mut self, mut func: impl FnMut(*mut u8)) {
+            fn get_components(self, mut func: impl FnMut(#ecs_path::ptr::OwningPtr<'_>)) {
                 #(#field_get_components)*
             }
         }

crates/bevy_ecs/src/bundle.rs

@@ -8,6 +8,7 @@ use crate::{
     archetype::{AddBundle, Archetype, ArchetypeId, Archetypes, ComponentStatus},
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
+    ptr::OwningPtr,
     storage::{SparseSetIndex, SparseSets, Storages, Table},
 };
 use bevy_ecs_macros::all_tuples;
@@ -85,14 +86,15 @@ pub unsafe trait Bundle: Send + Sync + 'static {
     /// # Safety
     /// Caller must return data for each component in the bundle, in the order of this bundle's
     /// Components
-    unsafe fn from_components(func: impl FnMut() -> *mut u8) -> Self
+    unsafe fn from_components<'a, F>(func: F) -> Self
     where
+        F: FnMut() -> OwningPtr<'a> + 'a,
         Self: Sized;
 
     /// Calls `func` on each value, in the order of this bundle's Components. This will
     /// "mem::forget" the bundle fields, so callers are responsible for dropping the fields if
     /// that is desirable.
-    fn get_components(self, func: impl FnMut(*mut u8));
+    fn get_components(self, func: impl FnMut(OwningPtr<'_>));
 }
 
 macro_rules! tuple_impl {
@@ -106,21 +108,23 @@ macro_rules! tuple_impl {
 
             #[allow(unused_variables, unused_mut)]
             #[allow(clippy::unused_unit)]
-            unsafe fn from_components(mut func: impl FnMut() -> *mut u8) -> Self {
+            unsafe fn from_components<'a, F>(func: F) -> Self
+            where
+                F: FnMut() -> OwningPtr<'a> + 'a
+            {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = (
-                    $(func().cast::<$name>(),)*
+                    $(func().inner().cast::<$name>(),)*
                 );
                 ($($name.read(),)*)
             }
 
             #[allow(unused_variables, unused_mut)]
-            fn get_components(self, mut func: impl FnMut(*mut u8)) {
+            fn get_components(self, mut func: impl FnMut(OwningPtr<'_>)) {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = self;
                 $(
-                    func((&mut $name as *mut $name).cast::<u8>());
-                    std::mem::forget($name);
+                    OwningPtr::make($name, func);
                 )*
             }
         }

crates/bevy_ecs/src/component.rs

@@ -1,6 +1,7 @@
 //! Types for declaring and storing [`Component`]s.
 
 use crate::{
+    ptr::OwningPtr,
     storage::{SparseSetIndex, Storages},
     system::Resource,
 };
@@ -118,7 +119,7 @@ impl ComponentInfo {
     }
 
     #[inline]
-    pub fn drop(&self) -> unsafe fn(*mut u8) {
+    pub fn drop(&self) -> unsafe fn(OwningPtr<'_>) {
         self.descriptor.drop
     }
 
@@ -163,7 +164,6 @@ impl SparseSetIndex for ComponentId {
     }
 }
 
-#[derive(Debug)]
 pub struct ComponentDescriptor {
     name: String,
     // SAFETY: This must remain private. It must match the statically known StorageType of the
@@ -174,13 +174,25 @@ pub struct ComponentDescriptor {
     is_send_and_sync: bool,
     type_id: Option<TypeId>,
     layout: Layout,
-    drop: unsafe fn(*mut u8),
+    drop: for<'a> unsafe fn(OwningPtr<'a>),
+}
+
+impl std::fmt::Debug for ComponentDescriptor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ComponentDescriptor")
+            .field("name", &self.name)
+            .field("storage_type", &self.storage_type)
+            .field("is_send_and_sync", &self.is_send_and_sync)
+            .field("type_id", &self.type_id)
+            .field("layout", &self.layout)
+            .finish()
+    }
 }
 
 impl ComponentDescriptor {
     // SAFETY: The pointer points to a valid value of type `T` and it is safe to drop this value.
-    unsafe fn drop_ptr<T>(x: *mut u8) {
-        x.cast::<T>().drop_in_place()
+    unsafe fn drop_ptr<T>(x: OwningPtr<'_>) {
+        x.inner().cast::<T>().drop_in_place()
     }
 
     pub fn new<T: Component>() -> Self {

crates/bevy_ecs/src/lib.rs

@@ -4,6 +4,7 @@ pub mod change_detection;
 pub mod component;
 pub mod entity;
 pub mod event;
+pub mod ptr;
 pub mod query;
 #[cfg(feature = "bevy_reflect")]
 pub mod reflect;

crates/bevy_ecs/src/ptr.rs

@@ -0,0 +1,83 @@
+use std::{marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
+
+#[derive(Copy, Clone)]
+pub struct Ptr<'a>(NonNull<u8>, PhantomData<&'a u8>);
+#[derive(Copy, Clone)]
+pub struct PtrMut<'a>(NonNull<u8>, PhantomData<&'a mut u8>);
+
+pub struct OwningPtr<'a>(NonNull<u8>, PhantomData<&'a mut u8>);
+
+macro_rules! impl_ptr {
+    ($ptr:ident) => {
+        impl $ptr<'_> {
+            /// Safety: the offset cannot make the existing ptr null, or take it out of bounds for it's allocation.
+            pub unsafe fn offset(self, count: isize) -> Self {
+                Self(
+                    NonNull::new_unchecked(self.0.as_ptr().offset(count)),
+                    PhantomData,
+                )
+            }
+
+            /// Safety: the offset cannot make the existing ptr null, or take it out of bounds for it's allocation.
+            pub unsafe fn add(self, count: usize) -> Self {
+                Self(
+                    NonNull::new_unchecked(self.0.as_ptr().add(count)),
+                    PhantomData,
+                )
+            }
+
+            pub unsafe fn new(inner: NonNull<u8>) -> Self {
+                Self(inner, PhantomData)
+            }
+
+            pub unsafe fn inner_nonnull(self) -> NonNull<u8> {
+                self.0
+            }
+        }
+    };
+}
+
+impl_ptr!(Ptr);
+impl<'a> Ptr<'a> {
+    pub unsafe fn inner(self) -> *const u8 {
+        self.0.as_ptr() as *const _
+    }
+
+    pub unsafe fn deref<T>(self) -> &'a T {
+        &*self.inner().cast()
+    }
+}
+impl_ptr!(PtrMut);
+impl<'a> PtrMut<'a> {
+    pub unsafe fn inner(self) -> *mut u8 {
+        self.0.as_ptr()
+    }
+
+    pub unsafe fn promote(self) -> OwningPtr<'a> {
+        OwningPtr(self.0, PhantomData)
+    }
+
+    pub unsafe fn deref_mut<T>(self) -> &'a mut T {
+        &mut *self.inner().cast()
+    }
+}
+impl_ptr!(OwningPtr);
+impl<'a> OwningPtr<'a> {
+    pub unsafe fn inner(self) -> *mut u8 {
+        self.0.as_ptr()
+    }
+
+    pub fn make<T, F: FnOnce(OwningPtr<'_>) -> R, R>(val: T, f: F) -> R {
+        let temp = MaybeUninit::new(val);
+        let ptr = unsafe { NonNull::new_unchecked(temp.as_mut_ptr().cast::<u8>()) };
+        f(Self(ptr, PhantomData))
+    }
+
+    pub unsafe fn read<T>(self) -> T {
+        self.inner().cast::<T>().read()
+    }
+
+    pub unsafe fn deref_mut<T>(self) -> &'a mut T {
+        &mut *self.inner().cast()
+    }
+}

crates/bevy_ecs/src/query/fetch.rs

@@ -336,7 +336,7 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
                 let column = tables[archetype.table_id()]
                     .get_column(state.component_id)
                     .unwrap();
-                self.table_components = column.get_data_ptr().cast::<T>();
+                self.table_components = column.get_data_ptr().inner_nonnull().cast::<T>();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
         }
@@ -345,9 +345,10 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
     #[inline]
     unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
         self.table_components = table
-            .get_column(state.component_id)
+            .get_column_mut(state.component_id)
             .unwrap()
-            .get_data_ptr()
+            .get_data_ptr_mut()
+            .inner_nonnull()
             .cast::<T>();
     }
 
@@ -360,7 +361,12 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
             }
             StorageType::SparseSet => {
                 let entity = *self.entities.add(archetype_index);
-                &*(*self.sparse_set).get(entity).unwrap().cast::<T>()
+                &*(*self.sparse_set)
+                    .get(entity)
+                    .unwrap()
+                    .inner_nonnull()
+                    .cast::<T>()
+                    .as_ptr()
             }
         }
     }
@@ -496,7 +502,7 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
                 let column = tables[archetype.table_id()]
                     .get_column(state.component_id)
                     .unwrap();
-                self.table_components = column.get_data_ptr().cast::<T>();
+                self.table_components = column.get_data_ptr().inner_nonnull().cast::<T>();
                 self.table_ticks = column.get_ticks_ptr();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
@@ -506,7 +512,7 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
     #[inline]
     unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
         let column = table.get_column(state.component_id).unwrap();
-        self.table_components = column.get_data_ptr().cast::<T>();
+        self.table_components = column.get_data_ptr().inner_nonnull().cast::<T>();
         self.table_ticks = column.get_ticks_ptr();
     }
 
@@ -529,7 +535,7 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
                 let (component, component_ticks) =
                     (*self.sparse_set).get_with_ticks(entity).unwrap();
                 Mut {
-                    value: &mut *component.cast::<T>(),
+                    value: &mut *component.inner_nonnull().cast::<T>().as_ptr(),
                     ticks: Ticks {
                         component_ticks: &mut *component_ticks,
                         change_tick: self.change_tick,