Simple State machine for unity

SimpleStateMachine.cs

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;
    }
}