[Merged by Bors] - Use lifetimed, type erased pointers in bevy_ecs

crates/bevy_ecs/macros/src/lib.rs

@@ -124,18 +124,17 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 self.#field.get_components(&mut func);
             });
             field_from_components.push(quote! {
-                #field: <#field_type as #ecs_path::bundle::Bundle>::from_components(&mut func),
+                #field: <#field_type as #ecs_path::bundle::Bundle>::from_components(ctx, &mut func),
             });
         } else {
             field_component_ids.push(quote! {
                 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, &mut func);
             });
             field_from_components.push(quote! {
-                #field: func().cast::<#field_type>().read(),
+                #field: func(ctx).inner().as_ptr().cast::<#field_type>().read(),
             });
         }
     }
@@ -157,14 +156,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<T, F>(ctx: &mut T, mut func: F) -> Self
+            where
+                F: FnMut(&mut T) -> #ecs_path::ptr::OwningPtr<'_>
+            {
                 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
     /// [`Component`]s
-    unsafe fn from_components(func: impl FnMut() -> *mut u8) -> Self
+    unsafe fn from_components<T, F>(ctx: &mut T, func: F) -> Self
     where
+        F: FnMut(&mut T) -> OwningPtr<'_>,
         Self: Sized;
 
     /// Calls `func` on each value, in the order of this bundle's [`Component`]s. This will
     /// [`std::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 {
@@ -105,21 +107,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<T, F>(ctx: &mut T, mut func: F) -> Self
+            where
+                F: FnMut(&mut T) -> OwningPtr<'_>
+            {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = (
-                    $(func().cast::<$name>(),)*
+                    $(func(ctx).inner().cast::<$name>(),)*
                 );
-                ($($name.read(),)*)
+                ($($name.as_ptr().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, &mut 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,
 };
@@ -117,7 +118,7 @@ impl ComponentInfo {
     }
 
     #[inline]
-    pub fn drop(&self) -> unsafe fn(*mut u8) {
+    pub fn drop(&self) -> unsafe fn(OwningPtr<'_>) {
         self.descriptor.drop
     }
 
@@ -162,7 +163,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
@@ -173,13 +173,26 @@ 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>),
+}
+
+// We need to ignore the `drop` field in our `Debug` impl
+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>().as_ptr().drop_in_place()
     }
 
     pub fn new<T: Component>() -> Self {
@@ -338,7 +351,7 @@ impl Components {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Copy, Clone, Debug)]
 pub struct ComponentTicks {
     pub(crate) added: u32,
     pub(crate) changed: u32,

crates/bevy_ecs/src/lib.rs

@@ -6,6 +6,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;
@@ -46,6 +47,7 @@ pub use bevy_ecs_macros::all_tuples;
 #[cfg(test)]
 mod tests {
     use crate as bevy_ecs;
+    use crate::query::QueryFetch;
     use crate::{
         bundle::Bundle,
         component::{Component, ComponentId},
@@ -901,7 +903,7 @@ mod tests {
 
         fn get_filtered<F: WorldQuery>(world: &mut World) -> Vec<Entity>
         where
-            F::Fetch: FilterFetch,
+            for<'x> QueryFetch<'x, F>: FilterFetch<'x>,
         {
             world
                 .query_filtered::<Entity, F>()

crates/bevy_ecs/src/ptr.rs

@@ -0,0 +1,149 @@
+use std::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
+
+/// Type-erased borrow of some unknown type chosen when constructing this type.
+///
+/// This type tries to act "borrow-like" which means that:
+/// - It should be considered immutable: its target must not be changed while this pointer is alive.
+/// - It must always points to a valid value of whatever the pointee type is.
+/// - The lifetime `'a` accurately represents how long the pointer is valid for.
+///
+/// It may be helpful to think of this type as similar to `&'a dyn Any` but without
+/// the metadata and able to point to data that does not correspond to a Rust type.
+#[derive(Copy, Clone)]
+pub struct Ptr<'a>(NonNull<u8>, PhantomData<&'a u8>);
+
+/// Type-erased mutable borrow of some unknown type chosen when constructing this type.
+///
+/// This type tries to act "borrow-like" which means that:
+/// - Pointer is considered exclusive and mutable. It cannot be cloned as this would lead to
+///   aliased mutability.
+/// - It must always points to a valid value of whatever the pointee type is.
+/// - The lifetime `'a` accurately represents how long the pointer is valid for.
+///
+/// It may be helpful to think of this type as similar to `&'a mut dyn Any` but without
+/// the metadata and able to point to data that does not correspond to a Rust type.
+pub struct PtrMut<'a>(NonNull<u8>, PhantomData<&'a mut u8>);
+
+/// Type-erased Box-like pointer to some unknown type chosen when constructing this type.
+/// Conceptually represents ownership of whatever data is being pointed to and so is
+/// responsible for calling its `Drop` impl. This pointer is _not_ responsible for freeing
+/// the memory pointed to by this pointer as it may be pointing to an element in a `Vec` or
+/// to a local in a function etc.
+///
+/// This type tries to act "borrow-like" like which means that:
+/// - Pointer should be considered exclusive and mutable. It cannot be cloned as this would lead
+///   to aliased mutability and potentially use after free bugs.
+/// - It must always points to a valid value of whatever the pointee type is.
+/// - The lifetime `'a` accurately represents how long the pointer is valid for.
+///
+/// It may be helpful to think of this type as similar to `&'a mut ManuallyDrop<dyn Any>` but
+/// without the metadata and able to point to data that does not correspond to a Rust type.
+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 its 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 its allocation.
+            pub unsafe fn add(self, count: usize) -> Self {
+                Self(
+                    NonNull::new_unchecked(self.0.as_ptr().add(count)),
+                    PhantomData,
+                )
+            }
+
+            /// # Safety
+            ///
+            /// The lifetime for the returned item must not exceed the lifetime `inner` is valid for
+            pub unsafe fn new(inner: NonNull<u8>) -> Self {
+                Self(inner, PhantomData)
+            }
+
+            pub fn inner(&self) -> NonNull<u8> {
+                self.0
+            }
+        }
+    };
+}
+
+impl_ptr!(Ptr);
+impl<'a> Ptr<'a> {
+    /// # Safety
+    ///
+    /// Another [`PtrMut`] for the same [`Ptr`] must not be created until the first is dropped.
+    pub unsafe fn assert_unique(self) -> PtrMut<'a> {
+        PtrMut(self.0, PhantomData)
+    }
+
+    /// # Safety
+    /// Must point to a valid `T`
+    pub unsafe fn deref<T>(self) -> &'a T {
+        &*self.0.as_ptr().cast()
+    }
+}
+impl_ptr!(PtrMut);
+impl<'a> PtrMut<'a> {
+    /// Transforms this [`PtrMut`] into an [`OwningPtr`]
+    ///
+    /// # Safety
+    /// Must have right to drop or move out of [`PtrMut`], and current [`PtrMut`] should not be accessed again unless it's written to again.
+    pub unsafe fn promote(self) -> OwningPtr<'a> {
+        OwningPtr(self.0, PhantomData)
+    }
+
+    /// Transforms this [`PtrMut<T>`] into a `&mut T` with the same lifetime
+    ///
+    /// # Safety
+    /// Must point to a valid `T`
+    pub unsafe fn deref_mut<T>(self) -> &'a mut T {
+        &mut *self.inner().as_ptr().cast()
+    }
+}
+impl_ptr!(OwningPtr);
+impl<'a> OwningPtr<'a> {
+    pub fn make<T, F: FnOnce(OwningPtr<'_>) -> R, R>(val: T, f: F) -> R {
+        let mut temp = MaybeUninit::new(val);
+        let ptr = unsafe { NonNull::new_unchecked(temp.as_mut_ptr().cast::<u8>()) };
+        f(Self(ptr, PhantomData))
+    }
+
+    //// Consumes the [`OwningPtr`] to obtain ownership of the underlying data of type `T`.
+    ///
+    /// # Safety
+    /// Must point to a valid `T`.
+    pub unsafe fn read<T>(self) -> T {
+        self.inner().as_ptr().cast::<T>().read()
+    }
+}
+
+pub(crate) trait UnsafeCellDeref<'a, T> {
+    unsafe fn deref_mut(self) -> &'a mut T;
+    unsafe fn deref(self) -> &'a T;
+    unsafe fn read(self) -> T
+    where
+        T: Copy;
+}
+impl<'a, T> UnsafeCellDeref<'a, T> for &'a UnsafeCell<T> {
+    unsafe fn deref_mut(self) -> &'a mut T {
+        &mut *self.get()
+    }
+    unsafe fn deref(self) -> &'a T {
+        &*self.get()
+    }
+
+    unsafe fn read(self) -> T
+    where
+        T: Copy,
+    {
+        self.get().read()
+    }
+}