How to model a state machine using generics? [dynamic dispatch?]

[jabriel (Gabe Joseph)]

I'm new to Swift, and as a learning exercise, I'm trying to make a framework for a simple state machine. But I haven't been able to find the right combination of protocols/inheritance/generics/opaque types to achieve what I want.

The overall goal is to have:

• various State structs, which hold relevant data for that state.
• various Event structs.
• each State has zero or more handle(event:) functions for specific event types, which return a new State object (representing the state after processing that event).

I'm hoping that using this might look like:

struct Start: Event {}
struct Stop: Event {}
struct ReceivedConnection: Event {
    let connection: Connection
}

struct Off: State {
    func handle(event: Start) -> Listening {
        return Listening()
    }
}

struct Listening: State {
    func handle(event: ReceivedConnection) -> Connected {
        event.connection.start()
        return Connected(connection: event.connection)
    }
    func handle(event: Stop) -> Off {
        return Off()
    }
}

struct Connected: State {
    let connection: Connection
    func handle(event: Stop) -> Off {
        connection.stop()
        return Off()
    }
}

As you can see, I'm basically hoping to use the type system to model:

• What state transitions are valid
• When it's valid to receive a kind of event
• What auxiliary data exists for each state

But I can't figure out how to implement a State that actually makes this work. I've tried inheritance, protocols using generics, and protocols using some/any. (Note that I used structs above for the example, but I don't care whether it uses value or reference types.)

I want to write something like (this does not work, for multiple reasons):

protocol State {
    func handle(event: some Event) -> any State
    // ^ Intended meaning: given a _concrete_ event, return something that's a `State`.
    // This does *not* mean that, to be a `State`, you must be able to handle *any* type of `Event`.
    // (Whether the return type is opaque, or boxed, or something else, isn't important to me.)
}

struct StateMachine {
    var state: any State

    mutating func handle(event: some Event) throws {
        // TODO: what happens when `event` is invalid for `state`?
        // How do we recognize this and throw an error?
        state = state.handle(event: event)
    }
}

The concrete States don't conform to this protocol, because they don't take an opaque some Event, but one concrete Event type.

I think what I want is basically dynamic dispatch: pick the matching method based on the concrete types at runtime. Additionally, I want to allow and handle the possibility that no methods match at runtime: this was an invalid Event, so we should throw an error.

My understanding of How to do dynamic dispatch? - #6 by anthonylatsis is that this isn't really possible in Swift. But maybe I'm missing something?

And if this isn't possible, any suggestions for a more idiomatic way to represent this?

[ibex10]

What you are trying to do is very admirable, but I am not sure if it can be done neatly in current Swift.

[aggie33]

Hmm... Yeah this would need multiple new features to be possible. It seems like each State type has its own unique list of valid Events that it can receive, and for each event it can return a different State type. First, I would question whether it even makes sense to model state as a protocol? I mean, will you ever use this in a generic context?

func takesState(_ state: some State) -> ??? {
  // what do you even do with `state`?
}

Anyhow...
To actually make this protocol, we would need some form of variadic associated types.

// just bike shedding syntax here
public protocol State {
  each associatedtype (Event, Response) where repeat each Response: State

  each func recieve(_ event: Event) -> Response
}

This might be possible at some point but definitely is not right now.

[stuchlej (Mikoláš Stuchlík)]

The issue is: you want to leverage type checking to statically verify the sanity of a state machine, correct?

TL;DR: In currently released Swift 5.8, no.

---

I don't think this is possible. In order to leverage the type checking, you would need to work with concrete types at all times (in the checked scope). This rules out using some/any, because they hide the type information at "code writing time."

Anyway. You have a States and Events and there is a subset of State ✕ Event ✕ State where 1st state is the current state, event is the event and 2nd state is the resulting state. In this case, the 2nd state is the function of 1st state and event.
The really naive solution would be to use plain old function overloading. Consider having

struct StateA {
    func handle(_ event: Event1) -> StateB { StateB() }
    func handle(_ event: Event2) -> StateC { StateC() }
}

struct StateB {
}

struct StateC {
}

struct Event1 {}
struct Event2 {}

Now, where to go from here? We would like to abstract upon this code to create a protocol, but as suggested by:

@aggie33

// just bike shedding syntax here
public protocol State {
  each associatedtype (Event, Response) where repeat each Response: State

  each func recieve(_ event: Event) -> Response
}

currently in such protocol, we can declare only "1:1" type associations, where in your case, we would need to be able to do set operation with types which is not possible.

---

And now for something completely different

You might want to look at upcoming Swift Macros and parameter packs. You might define something like:

@jabriel

protocol State {
    associatedtype PossibleOutcomes

    func handle(event: some Event) -> any State
}

where PossibleOutcomes would be "tuple of tuples" like ((Event1, StateB), (Event2, StateC)) and the handle function would be forbidden to call directly.

In this case I could imagine a macro, that would look at your PossibleOutcomes tuple of tuples and generate those overloaded functions for you.

Also, Macros can emit their own diagnostic messages to the user (since they are essentially compiler plugins) and you could somehow use them to implement the type checking you're looking for. But I think this would be a considerable effort to do correctly.

However, even if this would work, Macros are still work-in-progress feature and I would not advice sinking too much time into it before it is released.

[ibex10]

Why not just do it the old-fashion way, using classes.

@main
enum StateMachine {
    static func main () {
        var state : State = Off ()
        state = state.handle (event: Start ())
        state = state.handle (event: ReceivedConnection (connection: .init()))
        state = state.handle (event: Stop ())
        state = state.handle (event: Stop ())
    }
}

// -----------------------------------------------------
//
protocol Event: CustomStringConvertible {
}

extension Event {
    var description: String {
        "\(type (of:self))"
    }
}

struct Start: Event {}
struct Stop : Event {}
struct ReceivedConnection: Event {
    struct Connection {}
    let connection: Connection
}

// -----------------------------------------------------
//
class State {
    func handle (event: Start) -> State {
        invalid (event: event)
    }
    
    func handle (event: Stop) -> State {
        invalid (event: event)
    }
    
    func handle (event: ReceivedConnection) -> State {
        invalid (event: event)
    }
}

extension State {
    func log (event: Event) {
        print ("`\(event)` event in `\(type (of:self))`")
    }
    
    func invalid (event: Event) -> State {
        print ("`\(event)` event invalid for state `\(type (of:self))`")
        return self
    }
}

// -----------------------------------------------------
//
class Off: State {
    override func handle (event: Start) -> Listening {
        log (event:event)
        return Listening ()
    }
}

class Listening: State {
    override func handle (event: ReceivedConnection) -> Connected {
        log (event:event)
        return Connected ()
    }
    
    override func handle (event: Stop) -> Off {
        log (event:event)
        return Off ()
    }
}

class Connected: State {
    override func handle (event: Stop) -> Off {
        log (event:event)
        return Off ()
    }
}

[jabriel (Gabe Joseph)]

Wow, thanks for all the quick and helpful replies!

@ibex10

I am not sure if it can be done neatly in current Swift

@aggie33

this would need multiple new features to be possible

@stuchlej

TL;DR: In currently released Swift 5.8, no.

This is actually exactly what I was hoping to hear, so thanks for confirming that the wall I feel like I've been beating my head against is actually there.

@aggie33

First, I would question whether it even makes sense to model state as a protocol? I mean, will you ever use this in a generic context?

@aggie33, I think this is getting at the root of the problem. What I'm describing here really isn't a protocol, in the Swift sense. Given an arbitrary Event and State, the whole point is that you can't guarantee that handle is defined for that Event, which is really the opposite of a protocol.

@stuchlej

You might want to look at upcoming Swift Macros

Yeah, I think macros might be the best (only) solution here. Since what I want is basically dynamic dispatch—i.e. picking the right method at runtime based on the types of the objects—it's quite easy to implement that yourself as a big switch statement. It's just a lot of boilerplate typing.

This implementation actually works great, and is pretty straightforward. I just was hoping to not have to resort to macros for this, since as you said, they don't seem too mature yet.

// would be generated by macro
struct StateMachine {
    var state: State

    // basically
    // for event in allEvents {
    //     add handle function for this event
    //
    //     for state in statesThatRecieve[event] {
    //         add case for this state
    //     }
    //     add default case throwing StateMachineError
    // }

    mutating func handle(event: Start) throws {
        switch state {
            case let s as Off:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }

    mutating func handle(event: Stop) throws {
        switch state {
            case let s as Listening:
                state = s.handle(event: event)
            case let s as Connected:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }

    mutating func handle(event: ReceivedConnection) throws {
        switch state {
            case let s as Listening:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }
}

var sm = StateMachine(state: Off())
try? sm.handle(event: Start())
try? sm.handle(event: Start())
try? sm.handle(event: ReceivedConnection(connection: Connection()))
try? sm.handle(event: Start())
print("state machine state: \(sm.state)")

Output:

off->listening
Start() invalid for Listening()
listening->connected
Start() invalid for Connected(connection: macro.Connection())
state machine state: Connected(connection: macro.Connection())

Full code

protocol Event {}
protocol State {} // empty protocol, basically just a label

struct Start: Event {}
struct Stop: Event {}
struct ReceivedConnection: Event {
    let connection: Connection
}

struct Connection {}

struct Off: State {
    func handle(event: Start) -> Listening {
        print("off->listening")
        return Listening()
    }
}

struct Listening: State {
    func handle(event: ReceivedConnection) -> Connected {
        // event.connection.start()
        print("listening->connected")
        return Connected(connection: event.connection)
    }
    func handle(event: Stop) -> Off {
        print("\(type(of: self)) \(#function)")
        print("listening->off")
        return Off()
    }
}

struct Connected: State {
    let connection: Connection
    func handle(event: Stop) -> Off {
        // connection.stop()
        print("connected->off")
        return Off()
    }
}


struct StateMachineError: Error {
}

// would be generated by macro
struct StateMachine {
    var state: State

    // basically
    // for event in allEvents {
    //     add handle function for this event
    //
    //     for state in statesThatRecieve[event] {
    //         add case for this state
    //     }
    //     add default case throwing StateMachineError
    // }

    mutating func handle(event: Start) throws {
        switch state {
            case let s as Off:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }

    mutating func handle(event: Stop) throws {
        switch state {
            case let s as Listening:
                state = s.handle(event: event)
            case let s as Connected:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }

    mutating func handle(event: ReceivedConnection) throws {
        switch state {
            case let s as Listening:
                state = s.handle(event: event)
            default:
                print("\(event) invalid for \(state)")
                throw StateMachineError()
        }
    }
}

var sm = StateMachine(state: Off())
try? sm.handle(event: Start())
try? sm.handle(event: Start())
try? sm.handle(event: ReceivedConnection(connection: Connection()))
try? sm.handle(event: Start())
print("state machine state: \(sm.state)")

[jabriel (Gabe Joseph)]

@ibex10

Why not just do it the old-fashion way, using classes.

@ibex10 I had tried this at one point, but I did it slightly differently from you. I made the mistake of defining a function on the base class that could handle any Event:

class State {
    // intended as fallback for invalid events
    func handle(event: Event) -> State? {
        print("cannot handle \(event) from \(self)")
        return nil
    }
}
class Off: State {
    // states provide more specific implementations
    func handle(event: Start) -> Listening { ... }
}
...
class StateMachine {
    var state: State
    func handle(event: Event) {
        // this will fail every time! the compiler statically dispatches
        // to our fallback `handle` on the `State` base class, since
        // it's the only function it knows is valid for any State and Event
        // at _compile_ time.
        if let s = state.handle(event: event) {
            state = s
        }
    }
}

Your approach, with defining handlers on the base class for every possible concrete Event, works much better!

Of course, it's still a bit of boilerplate to write, especially if you had hundreds of possible events in a real system. Once again, it feels like macros might be necessary for a proper solution.

[Karl]

@jabriel

As you can see, I'm basically hoping to use the type system to model:

• What state transitions are valid
• When it's valid to receive a kind of event
• What auxiliary data exists for each state

IMO, the fundamental challenge you are facing is that the things you want are not generic. The set of valid state transitions depends on the current state, so if you abstract away the current state, clearly there will be no way for the type system to know which transitions are available from that unknown state or to prevent you from attempting invalid transitions.

You can use the type system to model these things, just not generics. I don't even think it's a question of missing language features.

In other words, I would align most closely with Mikoláš:

@stuchlej

I don't think this is possible. In order to leverage the type checking, you would need to work with concrete types at all times (in the checked scope). This rules out using some/any, because they hide the type information at "code writing time."

The alternative - to use dynamic dispatch/overloads - is a value-oriented approach, rather than a type-oriented approach.

[tkrajacic (Thomas Krajacic)]

As an alternative you might have a look at modeling your state machine as struct with an enclosed enum of your state. This has been applied in SwiftNIO, and many other packages successfully.

The great Cory Benfield gave a talk about this a few years back:

Cory Benfield - Building State Machines in Swift

[Nickolas_Pohilets (Nickolas Pokhylets)]

@jabriel

struct StateMachine {
    var state: any State

    mutating func handle(event: some Event) throws {
        // TODO: what happens when `event` is invalid for `state`?
        // How do we recognize this and throw an error?
        state = state.handle(event: event)
    }
}

You cannot handle only valid events, because of StateMachine. To do that you would need type of the event parameter depend on the value of state property. It is not possible in today’s Swift, and even if it would be - I don’t see how it would be usable. It’s just too hard for the client code to proof that event type is correct.

So effectively you need to handle all MxN combinations of (State, Event). Throwing an error counts as handling.

With this mindset, it is pretty straightforward to design a system with runtime checking for validity of events.

You can make states return optional type, with nil indicating invalid transition:

enum Event {
    case start
    case stop
    case receiveConnection(Connection)
}

protocol State {
    func handle(event: Event) -> Optional<any State>
}

struct StateMachine {
    var state: any State

    mutating func handle(event: some Event) throws {
        if let newState = state.handle(event: event) {
            state = newState
        } else {
            print("\(event) invalid for \(state)")
            throw StateMachineError()
    }
}

[scanon (Steve Canon)]

For very small state machines like this, my default approach would be to simply make everything concrete as others said--not even abstracting over events:

struct Off {
  func start() -> Listening { Listening() }
}

struct Listening {
  func stop() -> Off { Off() }
  func receive(connection id: Int) -> Receive {
    Receive(id: id)
  }
}

struct Receive {
  let id: Int
  func stop() -> Off { Off() }
}

This has the virtue that illegal transitions simply cannot be represented in the program. Sure, you can imagine making this generic, but what does that get you? Runtime errors instead of not being able to write those bugs at all.

[Joe_Groff]

In addition to what Steve said, you can also use noncopyable types to enforce that state transitions happen linearly, by making the state transitions consuming on the old state:

struct Off: ~Copyable {
  consuming func start() -> Listening { Listening() }
}

struct Listening: ~Copyable {
  consuming func stop() -> Off { Off() }
  consuming func receive(connection id: Int) -> Receive {
    Receive(id: id)
  }
}

struct Receive: ~Copyable {
  let id: Int
  consuming func stop() -> Off { Off() }
}

And then you can use enums to combine cases that share an event:

enum Stoppable: ~Copyable {
  case listening(Listening)
  case receive(Receive)

  consuming func stop() -> Off {
    switch consume self {
    case .listening(let listening):
      return listening.stop()
    case .receive(let receive):
      return receive.stop()
    }
  }
}

[jabriel (Gabe Joseph)]

Thanks for all the ideas here!

@tkrajacic

The great Cory Benfield gave a talk about this a few years back

This talk was helpful and enlightening, because it showed that the SwiftNIO approach was actually to do something more like what I was hoping to avoid: switch statements in event handlers explicitly checking what the current state is, with different code paths depending on it. (Basically the approach suggested in How to model a state machine using generics? [dynamic dispatch?] - #10 by Nickolas_Pohilets too.)

With complex real-world logic, that can turn into a very long function, with a lot of unrelated code side-by-side. Of course, you can separate the cases into their own functions, but then the switch becomes boilerplate. I was essentially hoping for a way to avoid that boilerplate by using the type system and dynamic dispatch.

In light of all this, I ended up doing away with the events-as-types idea entirely, and did something closer to what Steve suggested.

// statemachine.swift
enum StateMachineError: Error {
    case NoStateMachine(State.Type)
}

class State {
    weak var machine: StateMachine?
    
   func transition<S: State>(to state: S) throws {
        guard let machine = self.machine else {
            throw StateMachineError.NoStateMachine(type(of: self))
        }
        // TODO: how to do this atomically?
        print("transition \(self) -> \(state)")
        self.machine = nil
        machine.state = state
        state.machine = machine
    }
    
   func tryTransition<S: State>(to state: S) {
        do {
            try transition(to: state)
        } catch {
            print(error)
        }
    }
}

class StateMachine: ObservableObject {
    @Published fileprivate(set) var state: State
    
    /// Create the state machine in an initial state.
    init(state: State) {
        self.state = state
        state.machine = self
    }
}

// example-usage.swift
class Off: State {
    func startListening() {
        tryTransition(to: Listening())
    }
}

class Listening: State {
    let listener: Listener
    
    override init() {
        listener = Listener(...)
        listener.handleIncomingConnection = {[weak self] connection in
            guard let self = self else {
                return
            }
            self.tryTransition(to: Connected(connection: connection))
        }
    }
    func stop() {
        listener.cancel()
        tryTransition(to: Off())
    }
}

class Connected: State {
    let connection: Connection

    init(connection: Connection) {
        self.connection = connection
    }
    
    func stop() {
        connection.cancel()
        tryTransition(to: Off())
    }
}

class ExampleStateMachine: StateMachine {
    init() {
        super.init(state: Off())
    }
}

protocol Connection {
    func cancel() -> Void
}

protocol Listener {
    var handleIncomingConnection: (Connection) -> Void { get set }
    func cancel() -> Void
}

A big difference is that actions can only be triggered when you have a concrete State type. You can't blindly submit a StartListeningEvent; you have to first get a reference to an Off state object, so you can call startListening on it directly.
This ended up not being a big deal in my particular case, but I could see it being more annoying if one state set up a callback handler or delegate to some external API, then later transitioned to a different state, but the same delegate needed to remain attached and propagate events back into the state machine.

I might end up revisiting this as macros become more mature, since I'm still interested in that approach, but for now, what I have is working well enough.