From fdae9f7f48200e259d5a416db0dd6fbb6445113c Mon Sep 17 00:00:00 2001
From: Janne Valkealahti <janne.valkealahti@gmail.com>
Date: Sun, 10 Dec 2017 11:39:41 +0000
Subject: [PATCH] Allow to define extended state variables during persist

- Now in AbstractPersistingStateMachineInterceptor it is possible
  to use a Function to define which variables are persisted. This
  is for cases where it is not desirable to persist some variables
  for variours reasons.
- Relates to #423
---
 ...ractPersistingStateMachineInterceptor.java | 25 ++++++++++++-
 .../statemachine/support/Function.java        | 35 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 spring-statemachine-core/src/main/java/org/springframework/statemachine/support/Function.java

diff --git a/spring-statemachine-core/src/main/java/org/springframework/statemachine/persist/AbstractPersistingStateMachineInterceptor.java b/spring-statemachine-core/src/main/java/org/springframework/statemachine/persist/AbstractPersistingStateMachineInterceptor.java
index 72ecfb7b..e52a2e32 100644
--- a/spring-statemachine-core/src/main/java/org/springframework/statemachine/persist/AbstractPersistingStateMachineInterceptor.java
+++ b/spring-statemachine-core/src/main/java/org/springframework/statemachine/persist/AbstractPersistingStateMachineInterceptor.java
@@ -34,10 +34,12 @@ import org.springframework.statemachine.state.State;
 import org.springframework.statemachine.support.AbstractStateMachine;
 import org.springframework.statemachine.support.DefaultExtendedState;
 import org.springframework.statemachine.support.DefaultStateMachineContext;
+import org.springframework.statemachine.support.Function;
 import org.springframework.statemachine.support.StateMachineInterceptor;
 import org.springframework.statemachine.support.StateMachineInterceptorAdapter;
 import org.springframework.statemachine.transition.Transition;
 import org.springframework.statemachine.transition.TransitionKind;
+import org.springframework.util.Assert;
 
 /**
  * Base class for {@link StateMachineInterceptor} persisting {@link StateMachineContext}s.
@@ -54,6 +56,8 @@ import org.springframework.statemachine.transition.TransitionKind;
 public abstract class AbstractPersistingStateMachineInterceptor<S, E, T> extends StateMachineInterceptorAdapter<S, E>
 		implements StateMachinePersist<S, E, T> {
 
+	private Function<StateMachine<S, E>, Map<Object, Object>> extendedStateVariablesFunction = new AllVariablesFunction<>();
+
 	@Override
 	public void preStateChange(State<S, E> state, Message<E> message, Transition<S, E> transition, StateMachine<S, E> stateMachine) {
 		// try to persist context and in case of failure, interceptor
@@ -96,6 +100,17 @@ public abstract class AbstractPersistingStateMachineInterceptor<S, E, T> extends
 	 */
 	public abstract StateMachineContext<S, E> read(T contextObj) throws Exception;
 
+	/**
+	 * Sets the function creating extended state variables.
+	 *
+	 * @param extendedStateVariablesFunction the extended state variables function
+	 */
+	public void setExtendedStateVariablesFunction(
+			Function<StateMachine<S, E>, Map<Object, Object>> extendedStateVariablesFunction) {
+		Assert.notNull(extendedStateVariablesFunction, "'extendedStateVariablesFunction' cannot be null");
+		this.extendedStateVariablesFunction = extendedStateVariablesFunction;
+	}
+
 	/**
 	 * Builds the state machine context.
 	 *
@@ -105,7 +120,7 @@ public abstract class AbstractPersistingStateMachineInterceptor<S, E, T> extends
 	 */
 	protected StateMachineContext<S, E> buildStateMachineContext(StateMachine<S, E> stateMachine, State<S, E> state) {
 		ExtendedState extendedState = new DefaultExtendedState();
-		extendedState.getVariables().putAll(stateMachine.getExtendedState().getVariables());
+		extendedState.getVariables().putAll(extendedStateVariablesFunction.apply(stateMachine));
 
 		ArrayList<StateMachineContext<S, E>> childs = new ArrayList<StateMachineContext<S, E>>();
 		S id = null;
@@ -151,4 +166,12 @@ public abstract class AbstractPersistingStateMachineInterceptor<S, E, T> extends
 		// TODO: can this be empty as then we'd get error?
 		return ids2[ids2.length-1];
 	}
+
+	private static class AllVariablesFunction<S, E> implements Function<StateMachine<S, E>, Map<Object, Object>> {
+
+		@Override
+		public Map<Object, Object> apply(StateMachine<S, E> stateMachine) {
+			return stateMachine.getExtendedState().getVariables();
+		}
+	}
 }
diff --git a/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/Function.java b/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/Function.java
new file mode 100644
index 00000000..df5660bc
--- /dev/null
+++ b/spring-statemachine-core/src/main/java/org/springframework/statemachine/support/Function.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2017 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.statemachine.support;
+
+/**
+ * Mimics Java 8's Function interface to allow deferring a computation.
+ *
+ * @author Janne Valkealahti
+ *
+ * @param <T> the type of the input to the function
+ * @param <R> the type of the result of the function
+ */
+public interface Function<T, R> {
+
+	/**
+	 * Applies this function to the given argument.
+	 *
+	 * @param t the function argument
+	 * @return the function result
+	 */
+	R apply(T t);
+}