From 61c9a40fdedda31a686a24447688004615ac1ba7 Mon Sep 17 00:00:00 2001
From: Telzhaak <da.ahrens@outlook.de>
Date: Sat, 30 Jan 2021 21:10:14 +0100
Subject: [PATCH] [Bugfix] add_stage now checks Stage existence (#1346)

add_stage now checks stage existence
---
 crates/bevy_ecs/src/schedule/mod.rs | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/crates/bevy_ecs/src/schedule/mod.rs b/crates/bevy_ecs/src/schedule/mod.rs
index 8920fbf246..9c9ada93af 100644
--- a/crates/bevy_ecs/src/schedule/mod.rs
+++ b/crates/bevy_ecs/src/schedule/mod.rs
@@ -58,15 +58,14 @@ impl Schedule {
 
     pub fn add_stage<S: Stage>(&mut self, name: &str, stage: S) -> &mut Self {
         self.stage_order.push(name.to_string());
-        self.stages.insert(name.to_string(), Box::new(stage));
+        let prev = self.stages.insert(name.to_string(), Box::new(stage));
+        if prev.is_some() {
+            panic!("Stage already exists: {}.", name);
+        }
         self
     }
 
     pub fn add_stage_after<S: Stage>(&mut self, target: &str, name: &str, stage: S) -> &mut Self {
-        if self.stages.get(name).is_some() {
-            panic!("Stage already exists: {}.", name);
-        }
-
         let target_index = self
             .stage_order
             .iter()
@@ -75,16 +74,15 @@ impl Schedule {
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {}.", target));
 
-        self.stages.insert(name.to_string(), Box::new(stage));
         self.stage_order.insert(target_index + 1, name.to_string());
+        let prev = self.stages.insert(name.to_string(), Box::new(stage));
+        if prev.is_some() {
+            panic!("Stage already exists: {}.", name);
+        }
         self
     }
 
     pub fn add_stage_before<S: Stage>(&mut self, target: &str, name: &str, stage: S) -> &mut Self {
-        if self.stages.get(name).is_some() {
-            panic!("Stage already exists: {}.", name);
-        }
-
         let target_index = self
             .stage_order
             .iter()
@@ -93,8 +91,11 @@ impl Schedule {
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {}.", target));
 
-        self.stages.insert(name.to_string(), Box::new(stage));
         self.stage_order.insert(target_index, name.to_string());
+        let prev = self.stages.insert(name.to_string(), Box::new(stage));
+        if prev.is_some() {
+            panic!("Stage already exists: {}.", name);
+        }
         self
     }