diff --git a/src/fsm_generator.h b/src/fsm_generator.h
index ef79c2cd..a0d585e2 100644
--- a/src/fsm_generator.h
+++ b/src/fsm_generator.h
@@ -127,7 +127,7 @@ namespace etl
     etl::fsm_state_id_t get_state_id() const
     {
       return state_id;
-    }    
+    }
 
   protected:
 
@@ -143,7 +143,7 @@ namespace etl
 
     virtual fsm_state_id_t process_event(etl::imessage_router& source, const etl::imessage& message) = 0;
 
-    virtual void on_enter_state() {}; // By default, do nothing.
+    virtual fsm_state_id_t on_enter_state() { return state_id; }; // By default, do nothing.
     virtual void on_exit_state() {};  // By default, do nothing.
 
     // The state id.
@@ -171,9 +171,9 @@ namespace etl
     }
 
     //*******************************************
-    inline void on_enter_state(etl::ifsm_state &state)
+    inline fsm_state_id_t on_enter_state(etl::ifsm_state &state)
     {
-      state.on_enter_state();
+      return state.on_enter_state();
     }
 
     //*******************************************
@@ -258,21 +258,26 @@ namespace etl
 
     //*******************************************
     void receive(etl::imessage_router& source, const etl::imessage& message)
-    {       
+    {
       etl::fsm_state_id_t next_state_id = fsm_helper::process_event(*p_state, source, message);
-
       ETL_ASSERT(next_state_id < number_of_states, ETL_ERROR(etl::fsm_state_id_exception));
 
+      etl::ifsm_state* p_next_state = state_list[next_state_id];
+
       // Have we changed state?
-      if (next_state_id != p_state->get_state_id())
+      if (p_next_state != p_state)
       {
-        //p_state->on_exit_state();
-        fsm_helper::on_exit_state(*p_state);
+        do
+        {
+          p_state = p_next_state;
+          fsm_helper::on_exit_state(*p_state);
 
-        p_state = state_list[next_state_id];
-        ETL_ASSERT(p_state != nullptr, ETL_ERROR(etl::fsm_null_state_exception));
+          next_state_id = fsm_helper::on_enter_state(*p_state);
+          ETL_ASSERT(next_state_id < number_of_states, ETL_ERROR(etl::fsm_state_id_exception));
 
-        fsm_helper::on_enter_state(*p_state);
+          p_next_state = state_list[next_state_id];
+
+        } while (p_next_state != p_state); // Have we changed state again?
       }
     }
 
@@ -375,12 +380,12 @@ namespace etl
   // To generate to header file, run this at the command line.
   // Note: You will need Python and COG installed.
   //
-  // python -m cogapp -d -e -ofsm.h -DHandlers=<n> fsm_generator.h 
+  // python -m cogapp -d -e -ofsm.h -DHandlers=<n> fsm_generator.h
   // Where <n> is the number of messages to support.
   //
   // e.g.
   // To generate handlers for up to 16 events...
-  // python -m cogapp -d -e -ofsm.h -DHandlers=16 fsm_generator.h 
+  // python -m cogapp -d -e -ofsm.h -DHandlers=16 fsm_generator.h
   //
   // See CreateFSM.bat
   //***************************************************************************