From 60e8f0745aab824da6a77fe1cd0bd7acdfb3e56e Mon Sep 17 00:00:00 2001
From: Janne Valkealahti <janne.valkealahti@gmail.com>
Date: Fri, 10 Apr 2015 09:14:58 +0100
Subject: [PATCH] Type cast access for ext state variables.

---
 .../statemachine/ExtendedState.java                |  9 +++++++++
 .../statemachine/support/DefaultExtendedState.java | 14 ++++++++++++++
 .../src/main/java/demo/showcase/Application.java   | 12 ++++++------
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/spring-statemachine-core/src/main/java/org/springframework/statemachine/ExtendedState.java b/spring-statemachine-core/src/main/java/org/springframework/statemachine/ExtendedState.java
index f58828e1..b1bfe0d9 100644
--- a/spring-statemachine-core/src/main/java/org/springframework/statemachine/ExtendedState.java
+++ b/spring-statemachine-core/src/main/java/org/springframework/statemachine/ExtendedState.java
@@ -34,4 +34,13 @@ public interface ExtendedState {
 	 */
 	Map<Object, Object> getVariables();
 
+	/**
+	 * Gets a variable which is automatically casted into a type.
+	 *
+	 * @param key the variable key
+	 * @param type the variable type
+	 * @return the variable
+	 */
+	<T> T get(Object key, Class<T> type);
+
 }
diff --git a/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultExtendedState.java b/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultExtendedState.java
index 7fe9cdcb..a072b416 100644
--- a/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultExtendedState.java
+++ b/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultExtendedState.java
@@ -42,4 +42,18 @@ public class DefaultExtendedState implements ExtendedState {
 		return variables;
 	}
 
+	@SuppressWarnings("unchecked")
+	@Override
+	public <T> T get(Object key, Class<T> type) {
+		Object value = this.variables.get(key);
+		if (value == null) {
+			return null;
+		}
+		if (!type.isAssignableFrom(value.getClass())) {
+			throw new IllegalArgumentException("Incorrect type specified for variable '" +
+					key + "'. Expected [" + type + "] but actual type is [" + value.getClass() + "]");
+		}
+		return (T) value;
+	}
+
 }
diff --git a/spring-statemachine-samples/showcase/src/main/java/demo/showcase/Application.java b/spring-statemachine-samples/showcase/src/main/java/demo/showcase/Application.java
index 3e69b9cd..6faa0779 100644
--- a/spring-statemachine-samples/showcase/src/main/java/demo/showcase/Application.java
+++ b/spring-statemachine-samples/showcase/src/main/java/demo/showcase/Application.java
@@ -166,16 +166,16 @@ public class Application  {
 		@Override
 		public void execute(StateContext<States, Events> context) {
 			Map<Object, Object> variables = context.getExtendedState().getVariables();
-			Object foo = variables.get("foo");
-			if (foo instanceof Integer && ((Integer)foo) == 0) {
+			Integer foo = context.getExtendedState().get("foo", Integer.class);
+			if (foo == null) {
+				log.info("Init foo to 0");
+				variables.put("foo", 0);
+			} else if (foo == 0) {
 				log.info("Switch foo to 1");
 				variables.put("foo", 1);
-			} else if (foo instanceof Integer && ((Integer)foo) == 1) {
+			} else if (foo == 1) {
 				log.info("Switch foo to 0");
 				variables.put("foo", 0);
-			} else {
-				log.info("Init foo to 0");
-				variables.put("foo", 0);
 			}
 		}
 	}