Added a reference to the StateMachineController on the State which removes the need to get the controller via GetComponent. The Controller can then be cast to a derived implementation to retrieve references from it directly rather than having to add and retrieve references from SharedData which comes with an overhead

Author: PaulNonatomic

Runtime/StateGraph/State.cs

@@ -18,6 +18,7 @@ public abstract class State : ScriptableObject
 		public StateMachine StateMachine { get; set; }
 		public ISharedData SharedData { get; set; }
 		public TransitionEventData TransitionData { get; set; }
+		public StateMachineController Controller { get; set; }
 
 		/// <summary>
 		/// The OnStart method is mapped the state machines Awake method

Runtime/StateGraph/StateMachine.cs

@@ -28,16 +28,16 @@ public class StateMachine
 		private Dictionary<JumpId, StateNodeModel> _jumpNodeLookup = new();
 		private CancellationTokenSource _cancellationTokenSource = new();
 
-		public StateMachine(StateMachineModel model, GameObject gameObject)
+		public StateMachine(StateMachineModel model, StateMachineController controller)
 		{
 			SharedData = new SharedData();
-			Initialize(model, gameObject);
+			Initialize(model, controller);
 		}
 
-		public StateMachine(StateMachineModel model, GameObject gameObject, ISharedData sharedData = null)
+		public StateMachine(StateMachineModel model, StateMachineController controller, ISharedData sharedData = null)
 		{
 			SharedData = sharedData ?? new SharedData();
-			Initialize(model, gameObject);
+			Initialize(model, controller);
 		}
 
 		public void Update()
@@ -128,10 +128,10 @@ public void OnDestroy()
 			Model = null;
 		}
 
-		private void Initialize(StateMachineModel model, GameObject gameObject)
+		private void Initialize(StateMachineModel model, StateMachineController controller)
 		{
 			Model = StateMachineModel.CreateInstance(model);
-			Model.Initialize(gameObject, this, SharedData);
+			Model.Initialize(controller, this, SharedData);
 
 			CreateNodeLookupTable();
 			CreateTransitionLookupTable();

Runtime/StateGraph/StateMachineController.cs

@@ -146,7 +146,7 @@ private void CreateUniqueId()
 		private void CreateStateMachine()
 		{
 			if(!_model) return;
-			_stateMachine = new StateMachine(_model, gameObject);
+			_stateMachine = new StateMachine(_model, this);
 		}
 	}
 }

Runtime/StateGraph/StateMachineModel.cs

@@ -135,9 +135,9 @@ protected override void ValidateSubAssets()
 			StateMachineValidator.Validate(this);
 		}
 
-		public void Initialize(GameObject gameObject, StateMachine stateMachine, ISharedData sharedData)
+		public void Initialize(StateMachineController controller, StateMachine stateMachine, ISharedData sharedData)
 		{
-			if (!gameObject || stateMachine == null) return;
+			if (!controller || stateMachine == null) return;
 
 			foreach (var stateNode in Nodes)
 			{
@@ -146,7 +146,8 @@ public void Initialize(GameObject gameObject, StateMachine stateMachine, IShared
 				var instantiatedState = Instantiate(stateNode.State);
 				if (!instantiatedState) continue;
 
-				instantiatedState.GameObject = gameObject;
+				instantiatedState.GameObject = controller.gameObject;
+				instantiatedState.Controller = controller;
 				instantiatedState.StateMachine = stateMachine;
 				instantiatedState.SharedData = sharedData;
 				stateNode.State = instantiatedState;

Runtime/StateGraph/States/BaseParallelSubStateMachineState.cs

@@ -86,7 +86,7 @@ protected virtual void CreateStateMachines()
 			{
 				if (model == null) continue;
 
-				var subSubStateMachine = new StateMachine(model, GameObject, SharedData);
+				var subSubStateMachine = new StateMachine(model, Controller, SharedData);
 				subSubStateMachine.SetParent(StateMachine);
 				SubStateMachines.Add(subSubStateMachine);
 			}

Runtime/StateGraph/States/BaseSubStateMachineState.cs

@@ -97,7 +97,7 @@ protected virtual void CreateStateMachine()
 		{
 			if(!_model) return;
 
-			SubStateMachine = new StateMachine(_model, GameObject, SharedData);
+			SubStateMachine = new StateMachine(_model, Controller, SharedData);
 			SubStateMachine.SetParent(StateMachine);
 		}