State Machines in Ruby

Update: Added exceptions for non-existent transitions, and cleaned up a bit.

I’m not a fan of the currently popular Ruby state machine gems (state_machine, AASM). State machines are simple, and these are complicated gems with large APIs.

I had a tinker with implementing the simplest state machine I could in Ruby, with just the features I needed. What follows is what fell out of that tinkering. It’s all in roughly hand-tested form, so it comes with no warranty of any kind. I might expand this into a gem at some point.

Transition Tables

First up, it seems like a transition table would be the easiest way to specify a state machine in Ruby. We don’t need a DSL for this:

transition_table = TransitionTable.new(
  # State         Input     Next state      Output
  [:awaiting_foo, :foo] => [:awaiting_bar,  :do_stuff],
  [:awaiting_foo, :bar] => [:awaiting_foo,  :do_nothing],
  [:awaiting_bar, :bar] => [:all_done,      :do_other_stuff]
)

You’ll see why I use this hash structure below. (It might be nicer to turn this into simple nested arrays, and have the TransitionTable convert to a hash.)

What we want for actually running the state machine is a transition function: a function that takes the current state and some input, and returns the next state and some output.

We can make the transition table behave like that pretty easily:

class TransitionTable
  class TransitionError < RuntimeError
    def initialize(state, input)
      super "No transition from state #{state.inspect} for input #{input.inspect}"
    end
  end

  def initialize(transitions)
    @transitions = transitions
  end

  def call(state, input)
    @transitions.fetch([state, input])
  rescue KeyError
    raise TransitionError.new(state, input)
  end
end

The State Machine

Now we’ve got a way of specifying states and transitions, we need something to run the state machine.

class StateMachine
  def initialize(transition_function, initial_state)
    @transition_function = transition_function
    @state = initial_state
  end

  attr_reader :state

  def send_input(input)
    @state, output = @transition_function.call(@state, input)
    output
  end
end

We can fire up a state machine like this:

state_machine = StateMachine.new(transition_table, :awaiting_foo)

state_machine.state             # => :awaiting_foo
state_machine.send_input(:foo)  # => :do_stuff
state_machine.state             # => :awaiting_bar

Calling Methods on Transitions

I want to include a state machine in my model, and have methods on the model called when transitions happen. I can do something like this:

class MyModel
  STATE_TRANSITIONS = TransitionTable.new(
    # State         Input     Next state      Output
    [:awaiting_foo, :foo] => [:awaiting_bar,  :do_stuff],
    [:awaiting_foo, :bar] => [:awaiting_foo,  nil],
    [:awaiting_bar, :bar] => [:all_done,      :do_other_stuff]
  )

  def initialize
    @state_machine = StateMachine.new(STATE_TRANSITIONS, :awaiting_foo)
  end

  def handle_event(event)
    action = @state_machine.send_input(event)
    send(action) unless action.nil?
  end

  def do_stuff
    # ...
  end

  def do_other_stuff
    # ...
  end
end

Drawing a Picture

It’d be great if we could get a state diagram out of this, so:

class GraphvizFormatter
  HEADER = "digraph G {"
  FOOTER = "}"

  def format(transition_table)
    edges = transition_table.map do |current_state, input, next_state, output|
      %{  "#{escape(current_state)}" -> "#{escape(next_state)}" [ label=" #{escape(input)}/#{escape(output)}" ]}
    end
    [HEADER, *edges, FOOTER].join("\n")
  end

private

  def escape(string)
    string.to_s.gsub('"', '\\"')
  end
end

That requires being able to iterate over the TransitionTable, so we’ll need:

class TransitionTable
  include Enumerable

  def each
    @transitions.each_pair {|key, value| yield *key, *value }
  end
end

Throwing this and Graphviz at the above transition table gives: State diagram

Further Tinkering

There are a bunch more things I’d like to tinker with: