Optimize label comparisons and remove restrictions on types that can be labels

benches/benches/bevy_ecs/scheduling/schedule.rs

@@ -64,14 +64,17 @@ pub fn build_schedule(criterion: &mut Criterion) {
     // Use multiple different kinds of label to ensure that dynamic dispatch
     // doesn't somehow get optimized away.
     #[derive(Debug, Clone, Copy)]
-    struct NumLabel(usize);
+    struct NumLabel(u64);
     #[derive(Debug, Clone, Copy, SystemLabel)]
     struct DummyLabel;
 
     impl SystemLabel for NumLabel {
-        fn as_str(&self) -> &'static str {
-            let s = self.0.to_string();
-            Box::leak(s.into_boxed_str())
+        #[inline]
+        fn data(&self) -> u64 {
+            self.0
+        }
+        fn fmt(data: u64, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+            f.debug_tuple("NumLabel").field(&data).finish()
         }
     }
 
@@ -82,10 +85,6 @@ pub fn build_schedule(criterion: &mut Criterion) {
     // Method: generate a set of `graph_size` systems which have a One True Ordering.
     // Add system to the stage with full constraints. Hopefully this should be maximimally
     // difficult for bevy to figure out.
-    // Also, we are performing the `as_label` operation outside of the loop since that
-    // requires an allocation and a leak. This is not something that would be necessary in a
-    // real scenario, just a contrivance for the benchmark.
-    let labels: Vec<_> = (0..1000).map(|i| NumLabel(i).as_label()).collect();
 
     // Benchmark graphs of different sizes.
     for graph_size in [100, 500, 1000] {
@@ -109,12 +108,12 @@ pub fn build_schedule(criterion: &mut Criterion) {
                 // Build a fully-connected dependency graph describing the One True Ordering.
                 // Not particularly realistic but this can be refined later.
                 for i in 0..graph_size {
-                    let mut sys = empty_system.label(labels[i]).before(DummyLabel);
+                    let mut sys = empty_system.label(NumLabel(i)).before(DummyLabel);
                     for a in 0..i {
-                        sys = sys.after(labels[a]);
+                        sys = sys.after(NumLabel(a));
                     }
                     for b in i + 1..graph_size {
-                        sys = sys.before(labels[b]);
+                        sys = sys.before(NumLabel(b));
                     }
                     app.add_system(sys);
                 }

crates/bevy_app/src/app.rs

@@ -379,9 +379,9 @@ impl App {
         stage_label: impl StageLabel,
         system: impl IntoSystemDescriptor<Params>,
     ) -> &mut Self {
-        use std::any::TypeId;
+        let stage_label = stage_label.as_label();
         assert!(
-            stage_label.type_id() != TypeId::of::<StartupStage>(),
+            !stage_label.is::<StartupStage>(),
             "add systems to a startup stage using App::add_startup_system_to_stage"
         );
         self.schedule.add_system_to_stage(stage_label, system);
@@ -414,9 +414,9 @@ impl App {
         stage_label: impl StageLabel,
         system_set: SystemSet,
     ) -> &mut Self {
-        use std::any::TypeId;
+        let stage_label = stage_label.as_label();
         assert!(
-            stage_label.type_id() != TypeId::of::<StartupStage>(),
+            !stage_label.is::<StartupStage>(),
             "add system sets to a startup stage using App::add_startup_system_set_to_stage"
         );
         self.schedule
@@ -952,7 +952,7 @@ impl App {
     pub fn sub_app_mut(&mut self, label: impl AppLabel) -> &mut App {
         match self.get_sub_app_mut(label) {
             Ok(app) => app,
-            Err(label) => panic!("Sub-App with label '{:?}' does not exist", label.as_str()),
+            Err(label) => panic!("Sub-App with label '{:?}' does not exist", label.as_label()),
         }
     }
 
@@ -974,7 +974,7 @@ impl App {
     pub fn sub_app(&self, label: impl AppLabel) -> &App {
         match self.get_sub_app(label) {
             Ok(app) => app,
-            Err(label) => panic!("Sub-App with label '{:?}' does not exist", label.as_str()),
+            Err(label) => panic!("Sub-App with label '{:?}' does not exist", label.as_label()),
         }
     }
 

crates/bevy_derive/src/lib.rs

@@ -82,12 +82,20 @@ pub fn derive_enum_variant_meta(input: TokenStream) -> TokenStream {
 
 /// Generates an impl of the `AppLabel` trait.
 ///
-/// This works only for unit structs, or enums with only unit variants.
-/// You may force a struct or variant to behave as if it were fieldless with `#[app_label(ignore_fields)]`.
+/// For unit structs and enums with only unit variants, a cheap implementation can easily be created.
+///
+/// More complex types must be boxed and interned
+/// - opt in to this by annotating the entire item with `#[app_label(intern)]`.
+///
+/// Alternatively, you may force a struct or variant to behave as if
+/// it were fieldless with `#[app_label(ignore_fields)]`.
+/// This is especially useful for [`PhantomData`](core::marker::PhantomData) fields.
 #[proc_macro_derive(AppLabel, attributes(app_label))]
 pub fn derive_app_label(input: TokenStream) -> TokenStream {
     let input = syn::parse_macro_input!(input as syn::DeriveInput);
     let mut trait_path = BevyManifest::default().get_path("bevy_app");
     trait_path.segments.push(format_ident!("AppLabel").into());
-    derive_label(input, &trait_path, "app_label")
+    let mut id_path = BevyManifest::default().get_path("bevy_app");
+    id_path.segments.push(format_ident!("AppLabelId").into());
+    derive_label(input, &trait_path, &id_path, "app_label")
 }

crates/bevy_ecs/examples/derive_label.rs

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::{fmt::Debug, hash::Hash, marker::PhantomData};
 
 use bevy_ecs::prelude::*;
 
@@ -18,6 +18,57 @@ fn main() {
         GenericLabel::<f64>::One.as_label(),
         GenericLabel::<char>::One.as_label(),
     );
+
+    assert_eq!(format!("{:?}", UnitLabel.as_label()), "UnitLabel");
+    assert_eq!(format!("{:?}", WeirdLabel(1).as_label()), "WeirdLabel");
+    assert_eq!(format!("{:?}", WeirdLabel(2).as_label()), "WeirdLabel");
+    assert_eq!(
+        format!("{:?}", GenericLabel::<f64>::One.as_label()),
+        "GenericLabel::One::<f64>"
+    );
+
+    // Working with labels that need to be heap allocated.
+    let label = ComplexLabel {
+        people: vec!["John", "William", "Sharon"],
+    };
+    // Convert to to a LabelId. Its type gets erased.
+    let id = label.as_label();
+    assert_eq!(
+        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 = 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);
+    } else {
+        unreachable!();
+    }
+
+    // Different types with the same type constructor.
+    let id2 = ComplexerLabel(1_u32).as_label();
+    // The debug representations are the same...
+    assert_eq!(format!("{id:?}"), format!("{id2:?}"));
+    // ...but they do not compare equal...
+    assert_ne!(id, id2);
+    // ...nor can you downcast between monomorphizations.
+    assert!(id2.downcast::<ComplexerLabel<i128>>().is_none());
+    if let Some(label) = id2.downcast::<ComplexerLabel<u32>>() {
+        assert_eq!(label.0, 1);
+    } else {
+        unreachable!();
+    }
 }
 
 #[derive(SystemLabel)]
@@ -60,3 +111,13 @@ pub struct BadLabel2 {
     #[system_label(ignore_fields)]
     x: (),
 }*/
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, SystemLabel)]
+#[system_label(intern)]
+pub struct ComplexLabel {
+    people: Vec<&'static str>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, SystemLabel)]
+#[system_label(intern)]
+pub struct ComplexerLabel<T>(T);

crates/bevy_ecs/macros/src/lib.rs

@@ -436,8 +436,14 @@ pub fn derive_world_query(input: TokenStream) -> TokenStream {
 
 /// Generates an impl of the `SystemLabel` trait.
 ///
-/// This works only for unit structs, or enums with only unit variants.
-/// You may force a struct or variant to behave as if it were fieldless with `#[system_label(ignore_fields)]`.
+/// For unit structs and enums with only unit variants, a cheap implementation can easily be created.
+///
+/// More complex types must be boxed and interned
+/// - opt in to this by annotating the entire item with `#[system_label(intern)]`.
+///
+/// Alternatively, you may force a struct or variant to behave as if
+/// it were fieldless with `#[system_label(ignore_fields)]`.
+/// This is especially useful for [`PhantomData`](core::marker::PhantomData) fields.
 #[proc_macro_derive(SystemLabel, attributes(system_label))]
 pub fn derive_system_label(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
@@ -446,26 +452,44 @@ pub fn derive_system_label(input: TokenStream) -> TokenStream {
     trait_path
         .segments
         .push(format_ident!("SystemLabel").into());
-    derive_label(input, &trait_path, "system_label")
+    let mut id_path = bevy_ecs_path();
+    id_path.segments.push(format_ident!("schedule").into());
+    id_path.segments.push(format_ident!("SystemLabelId").into());
+    derive_label(input, &trait_path, &id_path, "system_label")
 }
 
 /// Generates an impl of the `StageLabel` trait.
 ///
-/// This works only for unit structs, or enums with only unit variants.
-/// You may force a struct or variant to behave as if it were fieldless with `#[stage_label(ignore_fields)]`.
+/// For unit structs and enums with only unit variants, a cheap implementation can easily be created.
+///
+/// More complex types must be boxed and interned
+/// - opt in to this by annotating the entire item with `#[stage_label(intern)]`.
+///
+/// Alternatively, you may force a struct or variant to behave as if
+/// it were fieldless with `#[stage_label(ignore_fields)]`.
+/// This is especially useful for [`PhantomData`](core::marker::PhantomData) fields.
 #[proc_macro_derive(StageLabel, attributes(stage_label))]
 pub fn derive_stage_label(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
     let mut trait_path = bevy_ecs_path();
     trait_path.segments.push(format_ident!("schedule").into());
     trait_path.segments.push(format_ident!("StageLabel").into());
-    derive_label(input, &trait_path, "stage_label")
+    let mut id_path = bevy_ecs_path();
+    id_path.segments.push(format_ident!("schedule").into());
+    id_path.segments.push(format_ident!("StageLabelId").into());
+    derive_label(input, &trait_path, &id_path, "stage_label")
 }
 
 /// Generates an impl of the `AmbiguitySetLabel` trait.
 ///
-/// This works only for unit structs, or enums with only unit variants.
-/// You may force a struct or variant to behave as if it were fieldless with `#[ambiguity_set_label(ignore_fields)]`.
+/// For unit structs and enums with only unit variants, a cheap implementation can easily be created.
+///
+/// More complex types must be boxed and interned
+/// - opt in to this by annotating the entire item with `#[ambiguity_set_label(intern)]`.
+///
+/// Alternatively, you may force a struct or variant to behave as if
+/// it were fieldless with `#[ambiguity_set_label(ignore_fields)]`.
+/// This is especially useful for [`PhantomData`](core::marker::PhantomData) fields.
 #[proc_macro_derive(AmbiguitySetLabel, attributes(ambiguity_set_label))]
 pub fn derive_ambiguity_set_label(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
@@ -474,13 +498,24 @@ pub fn derive_ambiguity_set_label(input: TokenStream) -> TokenStream {
     trait_path
         .segments
         .push(format_ident!("AmbiguitySetLabel").into());
-    derive_label(input, &trait_path, "ambiguity_set_label")
+    let mut id_path = bevy_ecs_path();
+    id_path.segments.push(format_ident!("schedule").into());
+    id_path
+        .segments
+        .push(format_ident!("AmbiguitySetLabelId").into());
+    derive_label(input, &trait_path, &id_path, "ambiguity_set_label")
 }
 
 /// Generates an impl of the `RunCriteriaLabel` trait.
 ///
-/// This works only for unit structs, or enums with only unit variants.
-/// You may force a struct or variant to behave as if it were fieldless with `#[run_criteria_label(ignore_fields)]`.
+/// For unit structs and enums with only unit variants, a cheap implementation can easily be created.
+///
+/// More complex types must be boxed and interned
+/// - opt in to this by annotating the entire item with `#[run_criteria_label(intern)]`.
+///
+/// Alternatively, you may force a struct or variant to behave as if
+/// it were fieldless with `#[run_criteria_label(ignore_fields)]`.
+/// This is especially useful for [`PhantomData`](core::marker::PhantomData) fields.
 #[proc_macro_derive(RunCriteriaLabel, attributes(run_criteria_label))]
 pub fn derive_run_criteria_label(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
@@ -489,7 +524,12 @@ pub fn derive_run_criteria_label(input: TokenStream) -> TokenStream {
     trait_path
         .segments
         .push(format_ident!("RunCriteriaLabel").into());
-    derive_label(input, &trait_path, "run_criteria_label")
+    let mut id_path = bevy_ecs_path();
+    id_path.segments.push(format_ident!("schedule").into());
+    id_path
+        .segments
+        .push(format_ident!("RunCriteriaLabelId").into());
+    derive_label(input, &trait_path, &id_path, "run_criteria_label")
 }
 
 pub(crate) fn bevy_ecs_path() -> syn::Path {

crates/bevy_ecs/src/schedule/label.rs

@@ -1,6 +1,7 @@
-pub use bevy_ecs_macros::{AmbiguitySetLabel, RunCriteriaLabel, StageLabel, SystemLabel};
 use bevy_utils::define_label;
 
+pub use bevy_ecs_macros::{AmbiguitySetLabel, RunCriteriaLabel, StageLabel, SystemLabel};
+
 define_label!(
     /// A strongly-typed class of labels used to identify [`Stage`](crate::schedule::Stage)s.
     StageLabel,

crates/bevy_ecs/src/schedule/mod.rs

@@ -211,9 +211,10 @@ impl Schedule {
             )
         }
 
+        let stage_label = stage_label.as_label();
         let stage = self
-            .get_stage_mut::<SystemStage>(&stage_label)
-            .unwrap_or_else(move || stage_not_found(&stage_label.as_label()));
+            .get_stage_mut::<SystemStage>(stage_label)
+            .unwrap_or_else(move || stage_not_found(&stage_label));
         stage.add_system(system);
         self
     }
@@ -281,11 +282,9 @@ impl Schedule {
         label: impl StageLabel,
         func: F,
     ) -> &mut Self {
-        let stage = self.get_stage_mut::<T>(&label).unwrap_or_else(move || {
-            panic!(
-                "stage '{:?}' does not exist or is the wrong type",
-                label.as_label()
-            )
+        let label = label.as_label();
+        let stage = self.get_stage_mut::<T>(label).unwrap_or_else(move || {
+            panic!("stage '{label:?}' does not exist or is the wrong type")
         });
         func(stage);
         self
@@ -304,9 +303,9 @@ impl Schedule {
     /// # let mut schedule = Schedule::default();
     /// # schedule.add_stage("my_stage", SystemStage::parallel());
     /// #
-    /// let stage = schedule.get_stage::<SystemStage>(&"my_stage").unwrap();
+    /// let stage = schedule.get_stage::<SystemStage>("my_stage").unwrap();
     /// ```
-    pub fn get_stage<T: Stage>(&self, label: &dyn StageLabel) -> Option<&T> {
+    pub fn get_stage<T: Stage>(&self, label: impl StageLabel) -> Option<&T> {
         self.stages
             .get(&label.as_label())
             .and_then(|stage| stage.downcast_ref::<T>())
@@ -325,9 +324,9 @@ impl Schedule {
     /// # let mut schedule = Schedule::default();
     /// # schedule.add_stage("my_stage", SystemStage::parallel());
     /// #
-    /// let stage = schedule.get_stage_mut::<SystemStage>(&"my_stage").unwrap();
+    /// let stage = schedule.get_stage_mut::<SystemStage>("my_stage").unwrap();
     /// ```
-    pub fn get_stage_mut<T: Stage>(&mut self, label: &dyn StageLabel) -> Option<&mut T> {
+    pub fn get_stage_mut<T: Stage>(&mut self, label: impl StageLabel) -> Option<&mut T> {
         self.stages
             .get_mut(&label.as_label())
             .and_then(|stage| stage.downcast_mut::<T>())

crates/bevy_ecs/src/schedule/stage.rs

@@ -1624,12 +1624,12 @@ mod tests {
                         *systems[index_a]
                             .labels()
                             .iter()
-                            .find(|a| a.type_id() == std::any::TypeId::of::<&str>())
+                            .find(|a| a.is::<&str>())
                             .unwrap(),
                         *systems[index_b]
                             .labels()
                             .iter()
-                            .find(|a| a.type_id() == std::any::TypeId::of::<&str>())
+                            .find(|a| a.is::<&str>())
                             .unwrap(),
                     )
                 })