Event-driven finite state machine DSL in Kotlin

Question

I wrote a DSL for defining a finite state machine. My robotics team wants to use a state machine to manage the state of the subsystems on our robot, like the arms, drive train, lifter, etc. Each subsystem is implemented as a FSM and subscribes to events published by buttons and sensors. I'm using an event bus as the middleman.

The design takes inspiration from Akka's FSM actor (written in Scala). JetBrain's Anko and Spek wonderfully showcase Kotlin's DSL capabilities, but outside of these libraries there are surprisingly few resources on how one would go about writing a DSL.

Thus, while this code works, it feels a clunky and I can't help but think I've missed something. I've explicitly listed some of the smells below, but I appreciate advice on all fronts.

Design goals and concerns

• Scoping. The appeal of a DSL lies in its restrictions. However, with the current structure we're leaking the build method. Although harmless, it should never be called while in the DSL. Is there any way to not expose this in the DSL or is this something we just have to accept as unavoidable?
  • On this note, do extension functions help here?
• Transitions. Transitions boil down to a function of (S, E) -> S (with side effects). In the DSL it's exposed as the whenState parameter for Transition.(E) -> Intent<S> where Intent is a sealed class with only the options for Goto<S> or Stay. The builder then maps that into a transition function. Can we change the design to expose only the goto and stay helper functions in Transition and keep private their respective Intent constructors?
• Event-driven. Not related to the DSL, but I went so far as to make StateMachine<S, E> directly implement (E) -> Unit in order to have a shared interface with my event bus; now I can just pass in a direct instance of the state machine as a listener instead of a method reference. I'm surprised I can do this. Although it's awesome, is this bad practice?

As far as state machines go, this would be my first true attempt at writing one. Any suggestions there may be on the state machine itself, including naming conventions, would also be of great help. Thanks!

Example scenario

Currently, here's how you would define the above elevator/lift using the DSL:

Lift.kt

class Lift(private val motor: Motor) : Stateful<LiftState, LiftEvent>(IDLE, {
  fun stop() = motor.stop()
  fun motorForward() = motor.setSpeed(1.0)
  fun motorBackward() = motor.setSpeed(-1.0)

  whenState(IDLE, ::stop) {
    when (it) {
      is GoDown -> goto(GOING_DOWN)
      is GoUp -> goto(GOING_UP)
      else -> stay
    }
  }

  whenState(GOING_UP, ::motorForward) {
    when (it) {
      is HitTop -> goto(AT_TOP)
      is GoDown -> goto(GOING_DOWN)
      is Halt -> goto(IDLE)
      else -> stay
    }
  }

  whenState(GOING_DOWN, ::motorBackward) {
    when (it) {
      is HitBottom -> goto(AT_BOTTOM)
      is GoUp -> goto(GOING_UP)
      is Halt -> goto(IDLE)
      else -> stay
    }
  }

  whenState(AT_TOP, ::stop) {
    when (it) {
      is GoDown -> goto(GOING_DOWN)
      else -> stay
    }
  }

  whenState(AT_BOTTOM, ::stop) {
    when (it) {
      is GoUp -> goto(GOING_UP)
      else -> stay
    }
  }
})

LiftState.kt

enum class KickerState {
  IDLE, GOING_DOWN, GOING_UP, AT_BOTTOM, AT_TOP
}

LiftEvent.kt

sealed class KickerEvent {
  object GoDown : KickerEvent()
  object GoUp : KickerEvent()
  object Halt : KickerEvent()
  object HitBottom : KickerEvent()
  object HitTop : KickerEvent()
}

Usage

// Register some triggers
downButton.onTriggered { EventBus.post(LifterEvent.GO_DOWN) }
upButton.onTriggered { EventBus.post(LifterEvent.GO_UP) }

// Register the listener
val lift = Lift(motor)
EventBus.subscribe(lift)

The code

StateMachine.kt

class StateMachine<S, in E> private constructor(
    initialState: S,
    private val states: Map<S, State>,
    private val transitions: Map<S, (E) -> S>
) : EventListener<E> {
  private var currentState = initialState

  override fun invoke(event: E) {
    // Look up the transition function
    val transition = transitions[currentState] ?: error("No transition function matching ($currentState, $event)")

    // Resolve the transition function in order to find our destination state
    val destinationKey = transition(event)
    val destinationState = states[destinationKey] ?: error("Cannot find target state with key $destinationKey")
    destinationState.onEntry()

    // Set the current state to the target state
    currentState = destinationKey
  }

  private data class State(val onEntry: () -> Any)

  // The DSL
  class Builder<S, E>(private val initialState: S) {
    private val states: MutableMap<S, State> = mutableMapOf()
    private val transitions: MutableMap<S, (E) -> S> = mutableMapOf()

    inner class Transition {
      fun goto(targetState: S): Intent<S> = Intent.GoTo(targetState)
      val stay: Intent<S> = Intent.Stay()
    }

    fun whenState(key: S, onEntry: () -> Any = {}, intentFunc: Transition.(E) -> Intent<S>) {
      // Add the state to the map, and throw an exception if there happened to be a previous state there
      states.put(key, State(onEntry))?.let { error("Attempting to add already-defined state $key.") }

      // Interprets the intent function into a pure transition function understood by the state machine
      val transitionTable: (E) -> S = { event: E ->
        val intent = Transition().intentFunc(event)
        val nextState = when (intent) {
          is Intent.Stay -> key
          is Intent.GoTo<S> -> intent.targetState
        }
        nextState
      }

      // Add a transition table for the specified state, and throw an exception if one was already registered
      transitions.put(key, transitionTable)?.let { error("State $key already has a transition table.") }
    }

    fun build(): StateMachine<S, E> {
      return StateMachine(initialState, states, transitions)
    }
  }
}

EventListener.kt

// typealias for (E) -> Unit
interface EventListener<in E> : (E) -> Unit

Stateful.kt

// uses Spek's initializer hack to expose the DSL builder
abstract class Stateful<out S, in E>(
    initialState: S,
    init: StateMachine.Builder<S, E>.() -> Unit)
: EventListener<E> by StateMachine.Builder<S, E>(initialState).apply(init).build()

EventBus.kt

object EventBus {
  private val subject = PublishRelay.create<Any>().toSerialized()
  val stream: Observable<Any> = subject // downcast when exposing

  // inlined because we want to take advantage of reified types
  inline fun <reified T : Any> subscribe(noinline listener: (T) -> Unit) {
    stream.ofType(T::class.java).subscribe(listener) // inline function only lets us use public APIs
  }

  fun publish(event: Any) = subject.call(event)
}

Answer 1

I'd design the StateMachine differently.

The StateMachine does not need to know about all the possible Transitions between the States. All the StateMachine should care about is the current State it is in.

The State itself should know about all possible Transitions (I call them Edges). So if an Event occurs, the StateMachine asks the current State, whether there is an appropriate Edge to go down. This Edge then tells the StateMachine what the current State of the StateMachine is after the Transition.

Just like the Strategy-Pattern: Wikipedia: Strategy Pattern

I've created a little Proof-of-Concept:

There is the class State:

class State(val name: String) {    
    private val edgeList = mutableListOf<Edge>()

    fun edge(name: String, targetState: String, init: Edge.() -> Unit) {
        val edge = Edge(name, targetState)
        edge.init()

        edgeList.add(edge)    
    }

    private val stateEnterAction = mutableListOf<(State) -> Unit>()

    //Add an action which will be called when the state is entered
    fun action(action: (State) -> Unit) {
        stateEnterAction.add(action)
    }

    fun enterState() {
        stateEnterAction.forEach { it(this) }
    }

    //Get the appropriate Edge for the Event
    fun getEdgeForEvent(event: Event): Edge {
        return edgeList.first { it.canHandleEvent(event) }
    }
}

It contains all of the possible Transitions(Edges). Furthermore, every State stores some actions, which will be called upon entering the State. (e.g. set motor-speed, stop motor)

Here is the Edge Class:

class Edge(val name: String, val targetState: String) {   
    lateinit var eventHandler: (Event) -> Boolean

    private val actionList = mutableListOf<(Edge) -> Unit>()

    fun action(action: (Edge) -> Unit) {
        actionList.add(action)
    }

    //Invoke when you go down the edge to another state
    fun enterEdge(retrieveState: (String) -> State): State {
        actionList.forEach { it(this) }

        return retrieveState(targetState)
    }

    fun canHandleEvent(event: Event): Boolean {
        return eventHandler(event)
    }
}

Every Edge has actions too, which will be called when you transition to another State. Also the Edge returns the new State after the transition is done.

And then there is the StateMachine with its Building-Function:

class StateMachine(private val initialStateName: String) {
    private lateinit var currentState: State

    private val stateList = mutableListOf<State>()

    fun state(name: String, init: State.() -> Unit) {
        val state = State(name)
        state.init()

        stateList.add(state)    
    }

    fun getStateByName(name: String): State {
        val result = stateList.firstOrNull { it.name == name}
            ?: throw NoSuchElementException(name)

        return result
    }

    fun initialize() {
        currentState = getStateByName(initialStateName)

        currentState.enterState()
    }

    fun eventOccured(event: Event) {
        val edge = currentState.getEdgeForEvent(event)

        if(edge is Edge) {
            val newState = edge.enterEdge { getStateByName(it) }

            newState.enterState()

            currentState = newState
        }
    }
}

fun buildStateMachine(initialStateName: String, init: StateMachine.() -> Unit): StateMachine {
    val stateMachine = StateMachine(initialStateName)

    stateMachine.init()

    return stateMachine
}

The different Events that occur:

sealed class Event {
    class GoUpEvent: Event() { }
    class GoDownEvent: Event() { }
    class HaltEvent: Event() { }
    class TopLimitHitEvent: Event() { }
    class BottomLimitHitEvent: Event() { }
}

Now you have your own small DSL to Build a complete StateMachine:

fun main(args: Array<String>) {

    val stateMachine = buildStateMachine("Idle") {
        state(name="Idle") {
            action {
                println("Entered state ${it.name}")
            }

            action {
                println("Set motor-speed to 0")
            }

            edge(name = "Go Up", targetState = "Going Up") {
                eventHandler = { it is Event.GoUpEvent }
            }

            edge(name = "Go Down", targetState = "Going Down") {
                eventHandler = { it is Event.GoDownEvent }
            }
        }

        state(name="Going Up") {
            action {
                println("Entered state ${it.name}")
                println("Set motor-speed to 1.0")
            }


            edge(name="Top Limit Hit", targetState = "At Top") {
                eventHandler = { it is Event.TopLimitHitEvent }    
            }

            edge(name="Halt", targetState = "Idle") {
                eventHandler = { it is Event.HaltEvent }
            }
        }

        state(name="Going Down") {
            action {
                println("Entered state ${it.name}")
            }

            action {
                println("Set motor-speed to -1.0")
            }

            edge(name="Bottom Limit Hit", targetState = "At Bottom") {
                eventHandler = { it is Event.BottomLimitHitEvent }
            }

            edge(name="Halt", targetState = "Idle") {
                eventHandler = { it is Event.HaltEvent }
            }
        }

        state(name="At Top") {
            action {
                println("Entered state ${it.name}")
                println("Set motor-speed to 0")
           }

            edge(name = "Go Down", targetState = "Going Down") {
                eventHandler = { it is Event.GoDownEvent }
            }
        }

        state(name="At Bottom") {
            action {
                println("Entered state ${it.name}")
                println("Set motor-speed to 0")
            }

            edge(name = "Go Up", targetState = "Going Up") {
                eventHandler = { it is Event.GoUpEvent }
            }
        }
    }

    stateMachine.initialize()


    stateMachine.eventOccured(Event.GoUpEvent())
    stateMachine.eventOccured(Event.TopLimitHitEvent())
    stateMachine.eventOccured(Event.GoDownEvent())
    stateMachine.eventOccured(Event.BottomLimitHitEvent())
    stateMachine.eventOccured(Event.GoUpEvent())
    stateMachine.eventOccured(Event.HaltEvent())
}

And if you run this Program, the output is:

Entered state Idle
Set motor-speed to 0
Entered state Going Up
Set motor-speed to 1.0
Entered state At Top
Set motor-speed to 0
Entered state Going Down
Set motor-speed to -1.0
Entered state At Bottom
Set motor-speed to 0
Entered state Going Up
Set motor-speed to 1.0
Entered state Idle
Set motor-speed to 0