improve type safety for generic hashed labels

Author: joseph-gio

crates/bevy_ecs/examples/derive_label.rs

@@ -45,6 +45,14 @@ fn main() {
         // created from a value of type `ComplexLabel`.
         unreachable!();
     }
+
+    // Generic heap-allocated labels;
+    let id = ComplexerLabel(1_i128).as_label();
+    assert_eq!(format!("{id:?}"), "ComplexerLabel(1)");
+    assert!(id.downcast::<ComplexerLabel<usize>>().is_none());
+    if let Some(label) = id.downcast::<ComplexerLabel<i128>>() {
+        assert_eq!(label.0, 1);
+    }
 }
 
 #[derive(SystemLabel)]

crates/bevy_ecs/src/schedule/label.rs

@@ -1,4 +1,4 @@
-use std::fmt::Debug;
+use std::{fmt::Debug, hash::Hash};
 
 use bevy_utils::{define_label, StableHashMap};
 use downcast_rs::{impl_downcast, Downcast};
@@ -32,14 +32,14 @@ define_label!(
 );
 
 /// Data structure used to intern a set of labels for a given type.
-pub struct Labels<L>(RwLock<StableHashMap<u64, L>>);
+pub struct Labels<K: Hash + Eq + Copy, L>(RwLock<StableHashMap<K, L>>);
 
 /// The type returned from [`Labels::get`](Labels#method.get).
 ///
 /// Will hold a lock on the string interner for type `L`, until this value gets dropped.
 pub type LabelGuard<'a, L> = parking_lot::MappedRwLockReadGuard<'a, L>;
 
-impl<L> Labels<L> {
+impl<K: Hash + Eq + Copy, L> Labels<K, L> {
     #[inline]
     pub const fn new() -> Self {
         Self(RwLock::new(StableHashMap::with_hasher(
@@ -48,7 +48,7 @@ impl<L> Labels<L> {
     }
 
     /// Interns a value, if it was not already interned in this set.
-    pub fn intern(&self, key: u64, f: impl FnOnce() -> L) {
+    pub fn intern(&self, key: K, f: impl FnOnce() -> L) {
         use parking_lot::RwLockUpgradableReadGuard as Guard;
 
         // Acquire an upgradeable read lock, since we might not have to do any writing.
@@ -70,11 +70,11 @@ impl<L> Labels<L> {
     /// optionally returning a value.
     ///
     /// Returns `None` if there is no interned label with that key.
-    pub fn scope<U>(&self, key: u64, f: impl FnOnce(&L) -> U) -> Option<U> {
+    pub fn scope<U>(&self, key: K, f: impl FnOnce(&L) -> U) -> Option<U> {
         self.0.read().get(&key).map(f)
     }
 
-    pub fn get(&self, key: u64) -> Option<LabelGuard<L>> {
+    pub fn get(&self, key: K) -> Option<LabelGuard<L>> {
         RwLockReadGuard::try_map(self.0.read(), |map| map.get(&key)).ok()
     }
 }

crates/bevy_macro_utils/src/lib.rs

@@ -345,25 +345,29 @@ fn derive_hashed_label(
         path
     };
 
+    let boxed_ty = quote! { ::std::boxed::Box::<dyn #downcast_debug_path + ::std::marker::Send + ::std::marker::Sync> };
+
     quote! {
         // NOTE: We need to box the interned objects in order to deal with generics.
         // Unfortunately, we cannot associate data with generic impls, so they all must be
         // stored together.
         // FIXME: Make this concrete for types without generics.
         static #interner_ident
-            : #interner_type_path <::std::boxed::Box::<dyn #downcast_debug_path + ::std::marker::Send + ::std::marker::Sync>>
+            : #interner_type_path <(u64, ::std::any::TypeId), #boxed_ty>
             = #interner_type_path::new();
 
         impl #impl_generics #trait_path for #ident #ty_generics #where_clause {
             #[inline]
             fn data(&self) -> u64 {
                 let hash = #compute_hash_path(self);
-                #interner_ident .intern(hash, || ::std::boxed::Box::new(::std::clone::Clone::clone(self)));
+                let key = (hash, ::std::any::TypeId::of::<Self>());
+                #interner_ident .intern(key, || ::std::boxed::Box::new(::std::clone::Clone::clone(self)));
                 hash
             }
             fn fmt(hash: u64, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
+                let key = (hash, ::std::any::TypeId::of::<Self>());
                 #interner_ident
-                    .scope(hash, |val| ::std::fmt::Debug::fmt(val, f))
+                    .scope(key, |val| ::std::fmt::Debug::fmt(val, f))
                     .ok_or(::std::fmt::Error)?
             }
         }
@@ -372,8 +376,9 @@ fn derive_hashed_label(
             type Output = #guard_type_path <'static, Self>;
             fn downcast_from(label: #id_path) -> Option<Self::Output> {
                 let hash = <#id_path as #trait_path>::data(&label);
+                let key = (hash, ::std::any::TypeId::of::<Self>());
                 <#guard_type_path<_>>::try_map(
-                    #interner_ident .get(hash)?,
+                    #interner_ident .get(key)?,
                     |val| <dyn #downcast_debug_path>::downcast_ref::<Self>(&**val),
                 ).ok()
             }