State Machine

state_machine.rb

require 'json'

# these were actually in a namespace, but whatever
StateMachineError = Class.new StandardError

class InvalidStateMachine < StateMachineError
  def self.check(seen, known)
    return if known.include? seen
    raise new(seen, known)
  end
  def initialize(seen, known)
    super "state machine references #{seen.inspect}, which is not in: #{known.inspect}"
  end
end

class UnknownEvent < StateMachineError
  attr_reader :event
  def initialize(event)
    @event = event
    super "Unknown event: #{event.inspect}"
  end
end

class InvalidTransition < StandardError
  attr_reader :state, :event, :transitions
  def initialize(state, event, transitions)
    @state = state
    @event = event
    @transitions = transitions
    super "State #{ state.to_json } has no event #{ event.to_json }, it only has #{ transitions.keys.to_json }"
  end
end


class StateMachine
  def initialize(states, events, transitions)
    @states = states.map(&:intern)
    @events = events.map(&:intern)
    transitions.each do |from_state, state_transitions|
      InvalidStateMachine.check from_state, @states
      state_transitions.each do |event, to_state|
        InvalidStateMachine.check to_state, @states
        InvalidStateMachine.check event, @events
      end
    end
    defaults = @events.map { |event| [ event, nil ] }.to_h
    @machine = @states.each_with_object({}) do |state, machine|
      state_transitions = transitions[state] || {}
      machine[state] = { **defaults, **state_transitions }
    end
  end

  def transition(state, event)
    state = state&.intern
    event = event&.intern
    transitions = @machine.fetch state
    raise UnknownEvent, event unless transitions.key? event
    new_state = transitions.fetch event
    raise InvalidTransition.new(state, event, transitions) unless new_state
    new_state
  end

  def to_graphviz
    gv = "digraph G {\n"
    gv << "  node [ shape = rect ]\n"
    @machine.each do |old_state, transitions|
      transitions.each do |event, new_state|
        next unless new_state
        gv << "  #{old_state.to_s.dump} -> #{new_state.to_s.dump}"
        gv << " [ label = #{event.to_s.dump} ]\n"
      end
    end
    gv << "}\n"
  end
end

state_machine_spec.rb

require "state_machine"

RSpec.describe StateMachine do
  describe "raises InvalidStateMachine on instantiation when" do
    example "transition states aren't in the list of known states" do
      expect { StateMachine.new [:a], [], b: {} }
        .to raise_error(InvalidStateMachine)
      expect { StateMachine.new [:a], [:b], a: { b: :c } }
        .to raise_error(InvalidStateMachine)
    end
    example "transition events aren't in the list of known events" do
      expect { StateMachine.new %I[a], [:b], a: { c: :a } }
        .to raise_error(InvalidStateMachine)
    end
  end

  describe "#transition" do
    it "raises UnknownEvent if it was not instantiated with that event" do
      machine = StateMachine.new %I[a], [:b], a: { b: :a }
      expect { machine.transition :a, :c }.to raise_error(UnknownEvent)
    end
    it "raises InvalidTransition if that state cannot be transitioned by that event" do
      machine = StateMachine.new %I[a], [:b], a: {}
      expect { machine.transition :a, :b }.to raise_error(InvalidTransition)
    end
    it "returns the output state, based on the input state and the event" do
      machine = StateMachine.new(
        %I[a b c d],
        %I[x y],
        a: { x: :b, y: :c },
        d: { x: :a, y: :b }
      )
      expect(machine.transition :a, :x).to eq :b
      expect(machine.transition :a, :y).to eq :c
      expect(machine.transition :d, :x).to eq :a
      expect(machine.transition :d, :y).to eq :b
    end
  end

  describe "#to_graphviz" do
    it "renders the state machine as a graphviz directed graph" do
      machine = StateMachine.new(
        %I[a b c d],
        %I[x y],
        a: { x: :b, y: :c },
        d: { x: :a, y: :b }
      )
      expect(machine.to_graphviz).to eq <<~GRAPHVIZ
      digraph G {
        node [ shape = rect ]
        "a" -> "b" [ label = "x" ]
        "a" -> "c" [ label = "y" ]
        "d" -> "a" [ label = "x" ]
        "d" -> "b" [ label = "y" ]
      }
      GRAPHVIZ
    end
  end
end