Skip to content

Instantly share code, notes, and snippets.

@sebtoun
Created July 20, 2018 15:51
Show Gist options
  • Save sebtoun/03fadfb0aa3b8e33d8ac6243ad3fbbf8 to your computer and use it in GitHub Desktop.
Save sebtoun/03fadfb0aa3b8e33d8ac6243ad3fbbf8 to your computer and use it in GitHub Desktop.
Simple State machine for unity
using System;
using System.Collections;
using System.Collections.Generic;
#if USE_GAMELOGIC_EXTENSIONS
using Gamelogic.Extensions;
#endif
using UnityEngine;
/// <summary>
/// A lightweight state machine that support either update or coroutine pattern for states update.
/// State is represented with an enum.
/// </summary>
/// <typeparam name="TState">The enum type that represents a state.</typeparam>
[ Serializable ]
public class SimpleStateMachine<TState> where TState : struct
{
/// <summary>
/// State transition is used when transitionning to the next state of a state machine with a parameter
/// that will be send to the next state's OnEnter handler.
/// </summary>
public struct StateTransition
{
/// <summary>
/// The state to transition to
/// </summary>
public TState State;
/// <summary>
/// The parameter that will be sent to the OnEnter handler.
/// </summary>
public readonly object Parameter;
public StateTransition( TState state, object parameter = null )
{
State = state;
Parameter = parameter;
}
public static implicit operator StateTransition( TState state )
{
return new StateTransition( state );
}
}
private struct StateHandlers
{
public Func<StateTransition> OnUpdate;
public Func<IEnumerable> UpdateCoroutine;
public Action<TState, object> OnEnter;
public Action<TState> OnExit;
}
private readonly Dictionary<TState, StateHandlers> _states;
#if USE_GAMELOGIC_EXTENSIONS
[ SerializeField, ReadOnly ]
#endif
private TState _currentState;
private IEnumerator _currentStateCoroutine;
/// <summary>
/// Creates a new empty state machine.
/// </summary>
public SimpleStateMachine()
{
_states = new Dictionary<TState, StateHandlers>();
}
/// <summary>
/// Registers a new state within this state machine.
/// </summary>
/// <remarks>A state should be registered exactly once.</remarks>
/// <param name="state">The state that is to be registered.</param>
/// <param name="onUpdate">Optionnal OnUpdate method that will be called each frame this state is active. The returned value is the next state to transition to.</param>
/// <param name="onEnter">Optionnal OnEnter method that will be called one time when this state is entered and before any Update.</param>
/// <param name="onExit">Optionnal OnExit method that will be called one time when this state is leaved and after any Update.</param>
/// <param name="coroutine">Optionnal Coroutine on which MoveNext will be called each frame this state is active. The coroutine should yield the next state to transition.</param>
/// <exception cref="ArgumentException"></exception>
public void RegisterState( TState state,
Func<StateTransition> onUpdate = null,
Action<TState, object> onEnter = null,
Action<TState> onExit = null,
Func<IEnumerable> coroutine = null )
{
if ( IsDefaultState( state ) )
{
throw new ArgumentException( "Cannot register default state" );
}
if ( onUpdate != null && coroutine != null )
{
Debug.LogWarning(
"Setting both onUpdate and coroutine is not supported. But we'll try to: update will be called before coroutine movenext" );
}
_states.Add( state, new StateHandlers()
{
OnUpdate = onUpdate,
OnEnter = onEnter,
OnExit = onExit,
UpdateCoroutine = coroutine
} );
}
/// <summary>
/// Registers a new state within this state machine.
/// </summary>
/// <remarks>A state should be registered exactly once.</remarks>
/// <param name="state">The state that is to be registered.</param>
/// <param name="onUpdate">Optionnal OnUpdate method that will be called each frame this state is active. The returned value is the next state to transition to.</param>
/// <param name="onEnter">Optionnal OnEnter method that will be called one time when this state is entered and before any Update.</param>
/// <param name="onExit">Optionnal OnExit method that will be called one time when this state is leaved and after any Update.</param>
/// <param name="coroutine">Optionnal Coroutine on which MoveNext will be called each frame this state is active. The coroutine should yield the next state to transition.</param>
/// <exception cref="ArgumentException"></exception>
public void RegisterState( TState state,
Func<TState> onUpdate = null,
Action<TState> onEnter = null,
Action<TState> onExit = null,
Func<IEnumerable> coroutine = null )
{
RegisterState( state,
onUpdate: onUpdate == null ? (Func<StateTransition>) null : ( () => (StateTransition) onUpdate() ),
onEnter: onEnter == null ? (Action<TState, object>) null : ( prevState, unused ) => onEnter( prevState ),
onExit: onExit,
coroutine: coroutine
);
}
/// <summary>
/// Should be called each frame to allow current state update and transitions handling.
/// </summary>
public void Update()
{
if ( IsDefaultState( _currentState ) )
{
Debug.LogWarning( "Trying to update a state machine that probably has not been started" );
return;
}
var onUpdate = _states[ _currentState ].OnUpdate;
if ( onUpdate != null )
{
var transition = onUpdate();
if ( !IsDefaultState( transition.State ) )
{
TransitionTo( transition.State, transition.Parameter );
}
}
if ( _currentStateCoroutine != null )
{
var hasNext = _currentStateCoroutine.MoveNext();
var current = _currentStateCoroutine.Current;
if ( current != null && !IsDefaultState( (TState) current ) )
{
TransitionTo( (TState) current );
}
else if ( !hasNext )
{
Debug.LogWarning( "State coroutine is over and has not transitionned" );
}
}
}
private bool IsDefaultState( TState state )
{
return EqualityComparer<TState>.Default.Equals( state, default( TState ) );
}
/// <summary>
/// Force transition to a given state. You can use this method to force a transtion from outise any state update.
/// Transitions are generally handled in state update using handlers return value to make sure states are exited
/// in a determined way.
/// </summary>
/// <param name="nextState">The state to transition to.</param>
/// <param name="userData">Optionnal OnEnter handler parameter.</param>
/// <exception cref="ArgumentException"></exception>
public void TransitionTo( TState nextState, object userData = null )
{
StateHandlers handlers;
if ( IsDefaultState( nextState ) || !_states.TryGetValue( nextState, out handlers ) )
{
throw new ArgumentException( "Trying to transition to default or unregistered state" );
}
var onExit = _states[ _currentState ].OnExit;
if ( onExit != null ) onExit( nextState );
var prevState = _currentState;
_currentState = nextState;
var onEnter = handlers.OnEnter;
if ( onEnter != null ) onEnter( prevState, userData );
var coroutine = handlers.UpdateCoroutine;
_currentStateCoroutine = coroutine != null ? coroutine().GetEnumerator() : null;
}
/// <summary>
/// Starts this state machine in the givent initial state.
/// </summary>
/// <param name="initialState">The state the machine will be upon start.</param>
/// <param name="userData">Optionnal OnEnter handler parameter.</param>
/// <exception cref="ArgumentException"></exception>
public void Start( TState initialState, object userData = null )
{
StateHandlers handlers;
if ( IsDefaultState( initialState ) || !_states.TryGetValue( initialState, out handlers ) )
{
throw new ArgumentException( "Trying to initialize state machine with default or unregistered state" );
}
_currentState = initialState;
var onEnter = handlers.OnEnter;
if ( onEnter != null ) onEnter( default( TState ), userData );
var coroutine = handlers.UpdateCoroutine;
_currentStateCoroutine = coroutine != null ? coroutine().GetEnumerator() : null;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment