Allow states to be automatically added to the state map.

They have to have a Constructor that accept the StateMap as argument, otherwise a RuntimeException is thrown.

1 files changed, 17 insertions, 0 deletions

diff --git a/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractState.java b/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractState.java
index 22fadc0de..5f2ff42ea 100644
--- a/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractState.java
+++ b/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractState.java
@@ -16,6 +16,23 @@ public abstract class AbstractState
 
 	public AbstractState getState(final Class<? extends AbstractState> stateClass)
 	{
+		if (!stateMap.containsKey(stateClass))
+		{
+			try
+			{
+				final AbstractState state = stateClass.getConstructor(StateMap.class).newInstance(stateMap);
+				stateMap.put(stateClass, state);
+			}
+			catch (Exception ex)
+			{
+				/*
+				 * This should never happen.
+				 * All states that are added to the map automatically,
+				 * have to have a Constructor that accepts the StateMap.
+				 */
+				throw new RuntimeException(ex);
+			}
+		}
 		return stateMap.get(stateClass);
 	}