Implement StageLabel for Box<dyn StageLabel>

Author: vladbat00

crates/bevy_ecs/src/schedule/label.rs

@@ -6,6 +6,12 @@ define_label!(SystemLabel);
 define_label!(AmbiguitySetLabel);
 define_label!(RunCriteriaLabel);
 
+impl StageLabel for Box<dyn StageLabel> {
+    fn dyn_clone(&self) -> Box<dyn StageLabel> {
+        self.as_ref().dyn_clone()
+    }
+}
+
 pub(crate) type BoxedStageLabel = Box<dyn StageLabel>;
 pub(crate) type BoxedSystemLabel = Box<dyn SystemLabel>;
 pub(crate) type BoxedAmbiguitySetLabel = Box<dyn AmbiguitySetLabel>;

crates/bevy_ecs/src/schedule/mod.rs

@@ -141,7 +141,7 @@ impl Schedule {
             .stage_order
             .iter()
             .enumerate()
-            .find(|(_i, stage_label)| &***stage_label == target)
+            .find(|(_i, stage_label)| stage_label.dyn_clone() == target.dyn_clone())
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));
 
@@ -175,7 +175,7 @@ impl Schedule {
             .stage_order
             .iter()
             .enumerate()
-            .find(|(_i, stage_label)| &***stage_label == target)
+            .find(|(_i, stage_label)| stage_label.dyn_clone() == target.dyn_clone())
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));
 
@@ -369,3 +369,26 @@ impl Stage for Schedule {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_adding_after_boxed_stage() {
+        let mut schedule = Schedule::default();
+        schedule.add_stage("first", SystemStage::single_threaded());
+        let stage = schedule.iter_stages().next().unwrap().0.dyn_clone();
+        // shouldn't panic
+        schedule.add_stage_after(stage, "second", SystemStage::single_threaded());
+    }
+
+    #[test]
+    fn test_adding_before_boxed_stage() {
+        let mut schedule = Schedule::default();
+        schedule.add_stage("first", SystemStage::single_threaded());
+        let stage = schedule.iter_stages().next().unwrap().0.dyn_clone();
+        // shouldn't panic
+        schedule.add_stage_before(stage, "second", SystemStage::single_threaded());
+    }
+}