1 files changed, 12 insertions, 8 deletions

diff --git a/src/FSM.hpp b/src/FSM.hpp
index 9570e26..346d0f0 100644
--- a/src/FSM.hpp
+++ b/src/FSM.hpp
@@ -6,27 +6,31 @@
 
 template<class T>
 class FSM {
-	T previousState;
-	T &state;
-	std::map<Transaction, std::function<void()>> handlers;
 public:
 	using Transaction = std::pair<T, T>;
+	using Handler = std::function<void(T &)>;
 
-	FSM(T &value) : state(value), previousState(value) {}
+	FSM(T initialState) : state(initialState), previousState(initialState) {}
 
 	~FSM() = default;
 
 	void Update() {
-		auto handler = handlers[Transaction{previousState, state}];
+		auto &handler = handlers[Transaction{previousState, state}];
 		if (handler)
-			handler();
+			handler(state);
+		previousState = state;
 	}
 
-	void RegisterHandler(T state, std::function<void()> handler) {
+	void RegisterHandler(T state, Handler handler) {
 		handlers[Transaction{state, state}] = handler;
 	}
 
-	void RegisterTransactionHandler(Transaction transaction, std::function<void()> handler) {
+	void RegisterTransactionHandler(Transaction transaction, Handler handler) {
 		handlers[transaction] = handler;
 	}
+
+private:
+	T previousState;
+	T state;
+	std::map<Transaction, Handler> handlers;
 };
\ No newline at end of file