Swift State Machines, Part 2

Feb 1, 2015

This is a two-parter. Part 1 covered some quick background about why our objects are broken, why that’s not our fault, and how we can fix them using state machines.

Alright. Now let’s build us one!

The Swift State Machine

The biggest impediment to using existing state machine libraries is getting them configured. Before we can get off the ground in some of these packages, we have to define not only all our states, but also all of their transitions. Often this requires a lot of ugly factories, redundant boilerplate, or even XML.

In our Swift state machine, we’re going to get rid of this cruft the Cocoa way: delegation! And as we’ll see, Swift has some nice features that help us keep our delegates slim as well.

But let’s start with the machine itself. All we really need is a property to hold the current state and an initializer:

class StateMachine{
  //We'll define StateType in a moment...
  var state:StateType   
  init(initialState:StateType){
    state = initialState
  }
}

Easy, right? Now let’s add our delegate. We’ll start with its protocol. We want it to do three things:

Tell us what the valid states are,
Decide whether a given transition should be allowed to happen, and
Optionally do something after the transition.

We’ll do all this with an associated type, a method that returns a Bool, and a Void method that gives the delegate a chance to do some processing.

protocol StateMachineDelegateProtocol:class{
  typealias StateType
  func shouldTransitionFrom(from:StateType, to:StateType)->Bool
  func didTransitionFrom(from:StateType, to:StateType)
}

[UPDATE: See Part 3 for some new ideas on whether didTransitionFrom(to:) should live in the delegate or not.]

Note that we’d probably like didTransitionFrom(to:) to be optional. That’s not possible in pure Swift protocols yet, but it’s planned for a future release.

Also check out that we’ve made this a class-only protocol. This is de rigueur for delegate protocols. After all, on a philosophical level, a delegate with value semantics doesn’t really make any sense. More pragmatically, delegate properties should almost always be weak or unowned. And values, of course, can’t be either.

So let’s add a delegate to our machine. This gets a little hairy, genericly-speaking, because we have to use the delegate protocol as a generic parameter (see the previous re: Generic Delegate Protocols for details):

class StateMachine<P:StateMachineDelegateProtocol>{
  private unowned let delegate:P
  var state:P.StateType

  init(initialState:P.StateType, delegate:P){
    state = initialState
    self.delegate = delegate
  }
}

Now we have to wire up calls to our delegate to both confirm we want to switch state, and to let it do something after we have successfully switched state.

Swift doesn’t provide us a way to conditionally set a stored property in its setter¹. So we’re going to make a public computed property that sits in front of a private stored property and conditionally sets it depending on the return from our delegate:

private var _state:P.StateType

var state:P.StateType{
  get{ return _state }
  set{
    if delegate.shouldTransitionFrom(_state, to:newValue){
      _state = newValue
    }
  }
}

Then we’ll add an observer to our stored property that calls our delegate whenever it’s been set:

private var _state:P.StateType{
  didSet{ delegate.didTransitionFrom(oldValue, to:_state) }
}

Finally, becasue setting the initial state of our machine isn’t tecnically a “transition”, we don’t want to have to call shouldTransitionFrom(to) when we do it. In our initializer, then, we’ll assign initialState directly to the stored property instead of our computed one. And we’re all set!

class StateMachine<P:StateMachineDelegateProtocol>{
  private unowned let delegate:P
  private var _state:P.StateType{
    didSet{ delegate.didTransitionFrom(oldValue, to:_state) }
  }
  var state:P.StateType{
    get{ return _state }
    set{
      if delegate.shouldTransitionFrom(_state, to:newValue){
        _state = newValue
      }
    }
  }

  init(initialState:P.StateType, delegate:P){
    _state = initialState
    self.delegate = delegate
  }
}

Transitional Awareness

That’s a nice looking machine… so how do we use it?

Because our delegate methods are so general, there’s no one right answer. But as I teased earlier, Swift has some nice tools for implementing methods like this. Namely enums, tuples, and the switch statement.

The first thing we’ll need to do is define our states, and that means assigning a concrete enum to our protocol’s associated type:

class MyClass:StateMachineDelegateProtocol{
  enum AsyncNetworkState{
    case Ready, Fetching, Saving
    case Success(NSDictionary)
    case Failure(NSError)
  }
  typealias StateType = AsyncNetworkState
}

Next, we’ll let our machine know when a given transition is allowed by implementing shouldTransitionFrom(to:). We could use a bunch of if statements or a big nested switch to iterate over all possible pairings of our states². But instead, we’re going to wrap all our params up in a tuple, switch on that, and use default and some of Swift’s cool pattern matching to make short work of it:

func shouldTransitionFrom(from:StateType, to:StateType)->Bool{
  switch (from, to){
  case (.Ready, .Fetching), (.Ready, .Saving):
    return true
  case (.Fetching, .Success), (.Fetching, .Failure):
    return true
  case (.Saving, .Success), (.Saving, .Failure):
    return true
  case (_, .Ready):
    return true
  default:
    return false
  }
}

Let’s start at the bottom. By default, any transition between states we don’t expressly define simply won’t happen. That should already make us feel warmer and fuzzier compared to the status quo.

Moving back to the top, we specifically approve transitions from Ready to any of our async states, and then from any of our async states to either Success or Failure. There’s no real reason to break these up on multiple lines — I often don’t. They’re separated here just for readability.

Finally we use Swift’s wildcard pattern to approve any transition to our Ready state — all in a single line.

And that’s it! Even this liberally formatted example squeaks by at 14 lines of very readable case statements.

Doing Stuff With State

We’ll follow the same general pattern to actually get stuff done when the state changes:

func didTransitionFrom(from:StateType, to:StateType){
  switch (from, to){
  case (.Ready, .Fetching):
    task = myAPI.fetchRequestWithCompletion{ ... }
  case (.Ready, .Saving):
    task = myAPI.saveRequestWithCompletion{ ... }
  }
}

Here, we’re sending off the appropriate request whenever the Fetching or Saving is transitioned to from a Ready state. Note that, in contrast to our example from Part 1, we’re now relying on the simple state of our system rather than the complex value of our task property to manage our requests. If we imagine a new implementation of our fetch action that looks something like:

func fetchThing(params:NSDictionary){
	machine.state = .Fetching
}

We can see that, because Fetching → Fetching is an invalid transition not allowed by our delegate, it doesn’t matter how many times we call fetchThing. We’ll still only get a single request sent³.

The last trick up our state machine’s sleeve is that we can use associated values in our enums to actually pass data around as we switch state. A fuller implementation of the above might look something like:

func didTransitionFrom(from:StateType, to:StateType){
  switch (from, to){
  case (.Ready, .Fetching):
    myAPI.fetchRequestWithCompletion{json, error in
	  if let someError = error{
	    machine.state = .Failure(someError)
	  } else{
	    machine.state .Success(json)
	  }
    }
  case (_, .Failure(let error)):
	displayGeneralError(error)
	machine.state = .Ready
  case (.Fetching, .Success(let json)):
    parseFetchSpecificJSON(json)
	machine.state = .Ready
  case (_, .Ready):
    updateInterface()
}

Our request’s completion handler, like most actions, can be reduced to nothing more than state transitions. And once we reduce it so, reasoning about our app’s interactions⁴ becomes much simpler. We no longer have to ask ourselves “Which actions cause the interface to get updated?” It’s a simple matter to see what state updates the UI (the Ready state in the example above) and what states are allowed to transition to it (Success and Failure).

Wrap Up

Once again, state machines don’t make this stuff easy. We’ve still got to reason through the use cases and write the code. But what state machines will do is keep our code abstract by decoupling concepts that never should have been coupled in the first place, opening the door to whole new worlds of refactoring.

Giving ourselves a dedicated structure to store state is just the first step.

The inline code above is a little disjointed in the name of constructing a narrative, so I’ve made a sample gist available of all the stuff we’ve talked about. Pull requests welcome!

[UPDATE: There are some great comments on this gist! Be sure to check them out and also Swift State Machines, Part 3, which responds to some of them.]

1: If our property has a setter, then it is, by definition, a computed property. And computed properties have no storage.↩︎

2: Five states mapped over both the from and to params in our delegate would be 25 (5²) distinct transition types. Ick!↩︎

3: That is, until our state returns to Ready, as that is the only state allowed to transition to Fetching.↩︎

4: Extra especially its asynchronous interactions.↩︎