Moving the link between the states from the constructor of the state machine to the states classes.

1 files changed, 12 insertions, 4 deletions

diff --git a/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractYesNoState.java b/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractYesNoState.java
index 67701c8db..67baf6bb6 100644
--- a/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractYesNoState.java
+++ b/EssentialsUpdate/src/com/earth2me/essentials/update/states/AbstractYesNoState.java
@@ -4,11 +4,17 @@ package com.earth2me.essentials.update.states;
 public abstract class AbstractYesNoState extends AbstractState
 {
 	private boolean answer = false;
-	private final transient AbstractState yesState;
-	private final transient AbstractState noState;
+	private final transient Class<? extends AbstractState> yesState;
+	private final transient Class<? extends AbstractState> noState;
 
-	public AbstractYesNoState(final AbstractState yesState, final AbstractState noState)
+	public AbstractYesNoState(final StateMap states, final Class<? extends AbstractState> nextState)
 	{
+		this(states, nextState, nextState);
+	}
+
+	public AbstractYesNoState(final StateMap states, final Class<? extends AbstractState> yesState, final Class<? extends AbstractState> noState)
+	{
+		super(states);
 		this.yesState = yesState;
 		this.noState = noState;
 	}
@@ -16,7 +22,9 @@ public abstract class AbstractYesNoState extends AbstractState
 	@Override
 	public AbstractState getNextState()
 	{
-		return answer ? yesState : noState;
+		return answer
+			   ? (yesState == null ? null : getState(yesState))
+			   : (noState == null ? null : getState(noState));
 	}
 
 	@Override