Support downcasting interned labels

Author: joseph-gio

crates/bevy_ecs/examples/derive_label.rs

@@ -41,17 +41,38 @@ fn main() {
         format!("{id:?}"),
         r#"ComplexLabel { people: ["John", "William", "Sharon"] }"#
     );
+    // Try to downcast it back to its concrete type.
+    if let Some(complex_label) = id.downcast::<ComplexLabel>() {
+        assert_eq!(complex_label.people, vec!["John", "William", "Sharon"]);
+    } else {
+        // The downcast will never fail in this example, since the label is always
+        // created from a value of type `ComplexLabel`.
+        unreachable!();
+    }
 
     // Generic heap-allocated labels.
     let id = WrapLabel(1_i128).as_label();
     assert_eq!(format!("{id:?}"), "WrapLabel(1)");
+    assert!(id.downcast::<WrapLabel<usize>>().is_none());
+    if let Some(label) = id.downcast::<WrapLabel<i128>>() {
+        assert_eq!(label.0, 1);
+    } else {
+        unreachable!();
+    }
 
     // Different types with the same type constructor.
     let id2 = WrapLabel(1_u32).as_label();
     // The debug representations are the same...
     assert_eq!(format!("{id:?}"), format!("{id2:?}"));
-    // ...but they do not compare equal.
+    // ...but they do not compare equal...
     assert_ne!(id, id2);
+    // ...nor can you downcast between monomorphizations.
+    assert!(id2.downcast::<WrapLabel<i128>>().is_none());
+    if let Some(label) = id2.downcast::<WrapLabel<u32>>() {
+        assert_eq!(label.0, 1);
+    } else {
+        unreachable!();
+    }
 }
 
 #[derive(SystemLabel)]

crates/bevy_macro_utils/src/lib.rs

@@ -370,6 +370,12 @@ fn derive_interned_label(
         path
     };
     let interner_ident = format_ident!("{}_INTERN", ident.to_string().to_uppercase());
+    let downcast_trait_path = {
+        let mut path = manifest.get_path("bevy_utils");
+        path.segments.push(format_ident!("label").into());
+        path.segments.push(format_ident!("LabelDowncast").into());
+        path
+    };
 
     Ok(quote! {
         static #interner_ident : #interner_type_expr = #interner_type_path::new();
@@ -384,5 +390,12 @@ fn derive_interned_label(
                 ::std::fmt::Debug::fmt(&*val, f)
             }
         }
+
+        impl #impl_generics #downcast_trait_path for #ident #ty_generics #where_clause {
+            type Output = #guard_type_path <'static, Self>;
+            fn downcast_from(idx: u64) -> Option<Self::Output> {
+                #interner_ident .get(idx as usize)
+            }
+        }
     })
 }

crates/bevy_utils/src/label.rs

@@ -3,6 +3,7 @@
 use std::{
     any::Any,
     hash::{Hash, Hasher},
+    ops::Deref,
 };
 
 use crate::Interner;
@@ -49,6 +50,15 @@ where
     }
 }
 
+/// Trait for implementors of `*Label` types that support downcasting.
+pub trait LabelDowncast {
+    // FIXME: use "return position impl Trait in traits" when that stabilizes.
+    /// The type returned from [`downcast_from`](#method.downcast_from).
+    type Output: Deref<Target = Self>;
+    /// Attempts to convert data from a label to type `Self`. Returns a reference-like type.
+    fn downcast_from(data: u64) -> Option<Self::Output>;
+}
+
 #[doc(hidden)]
 pub struct VTable {
     // FIXME: When const TypeId stabilizes, inline the type instead of using a fn pointer for indirection.
@@ -163,6 +173,33 @@ macro_rules! define_label {
             pub fn is<L: $label_name>(self) -> bool {
                 self.type_id() == ::std::any::TypeId::of::<L>()
             }
+            /// Attempts to downcast this label to type `L`.
+            ///
+            /// As an anti-footgun measure, the returned reference-like type is `!Send + !Sync`
+            /// -- often it is a mutex guard type, so it should be contained to one thread,
+            /// and should not be held onto for very long.
+            ///
+            /// This method is not available for all types of labels.
+            pub fn downcast<L>(self) -> Option<impl ::std::ops::Deref<Target = L>>
+            where
+                L: $label_name + $crate::label::LabelDowncast
+            {
+                // Wraps a deref type and forces it to be !Send + !Sync
+                struct NonSendSyncDeref<L>(L, ::std::marker::PhantomData<*mut u8>);
+                impl<L: ::std::ops::Deref> ::std::ops::Deref for NonSendSyncDeref<L> {
+                    type Target = <L as ::std::ops::Deref>::Target;
+                    fn deref(&self) -> &Self::Target {
+                        &*self.0
+                    }
+                }
+
+                if self.is::<L>() {
+                    let val = L::downcast_from(self.data())?;
+                    Some(NonSendSyncDeref(val, ::std::marker::PhantomData))
+                } else {
+                    None
+                }
+            }
         }
 
         impl $label_name for &'static str {