From 64dca14b3c98528ca524c66670ef13a2c1047cae Mon Sep 17 00:00:00 2001
From: sapir <yasapir@gmail.com>
Date: Sun, 13 Dec 2020 19:26:12 +0200
Subject: [PATCH] Fix hang on missing state update handler

---
 crates/bevy_ecs/src/schedule/state.rs | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/schedule/state.rs b/crates/bevy_ecs/src/schedule/state.rs
index 629f926336b28..d6c6963032aa0 100644
--- a/crates/bevy_ecs/src/schedule/state.rs
+++ b/crates/bevy_ecs/src/schedule/state.rs
@@ -98,7 +98,7 @@ impl<T: Resource + Clone> Stage for StateStage<T> {
     }
 
     fn run(&mut self, world: &mut World, resources: &mut Resources) {
-        loop {
+        let current_stage = loop {
             let (next_stage, current_stage) = {
                 let mut state = resources
                     .get_mut::<State<T>>()
@@ -132,14 +132,17 @@ impl<T: Resource + Clone> Stage for StateStage<T> {
                 {
                     enter_next.run(world, resources);
                 }
-            } else if let Some(update_current) = self
-                .stages
-                .get_mut(&current_stage)
-                .and_then(|stage| stage.update.as_mut())
-            {
-                update_current.run(world, resources);
-                break;
+            } else {
+                break current_stage;
             }
+        };
+
+        if let Some(update_current) = self
+            .stages
+            .get_mut(&current_stage)
+            .and_then(|stage| stage.update.as_mut())
+        {
+            update_current.run(world, resources);
         }
     }
 }