Type Safe State Machines

For chuckles, let’s make a type-safe state-machine in Swift. I would like to exploit types in order to prevent invalid states being generated. This is a prelude to implementing state-charts. My first steps are to implement state-machines, then later work on nesting them in order to construct state-charts.

I won’t explain what state-machines are; there are plenty of better sources for that! Instead I’m just going to write some code. Please note that the code examples require Swift 4.0.

Putting things in Order

As a motivating example, I am going to model the state of an online order. To begin with, here are the basic set of states.

• Pending: Customer has not placed the order yet
• Billed: Has been placed and billed
• Packed: Order is packed pending shipping
• Shipped: Sent to customer
• Cancelled: Cancelled before being shipped

The state transitions are also quite basic.

• Pending to Billed
• Billed to Packed
• Packed to Shipped
• Packed to Billed
• Billed to Cancelled
• Packed to Cancelled

An order has to be billed and packed before it’s shipped and it can only be cancelled if it is billed or packed. It cannot be cancelled while it is pending or if it has been shipped.

Adding Types

Let’s represent our order as a struct. The type-safety comes from representing it’s states as actual types. So, Order becomes generic.

struct Order<T> {

}

The state types will be phantom-types so lets create them as enums without initializers.

enum Pending {}
enum Billed {}
enum Packed {}
enum Shipped {}
enum Cancelled {}

Now we can represent an order in any of the states.

let pending = Order<Pending>()
let shipped = Order<Shipped>()

The next task is to implement the transitions from state to state. These are getters which are added to the Order using constrained-extensions.

extension Order where T == Pending {
  var place: Order<Billed> {
    return Order<Billed>()
  }
}

This getter will only be available where the state type is Pending.

let update = Order<Pending>().place // OK!
let ohno = Order<Billed>().place // Will not work

This technique is exactly the same for all the other states and transitions.

Getting Fancy

Now we can extend the set of states and transitions, while also exploiting protocols in order to save us some work. We’ll add some extra states.

• LostInTransit
• DamagedInTransit
• Refunded

With the following transitions.

• Shipped to LostInTransit
• Shipped to DamagedInTransit
• LostInTransit to Refunded
• DamagedInTransit to Refunded

Adding those extra states and transitions is easy. But we can see some redundancy in the transitions. Both LostInTransit and DamagedInTransit are able to transition to Refunded. Let’s use a protocol to add a shared transition.

protocol Refundable {}

enum LostInTransit: Refundable {}
enum DamagedInTransit: Refundable {}
enum Refunded {}

extension Order where T: Refundable {
  var refund: Order<Refunded> {
    return Order<Refunded>()
  }
}

We also have some redundancy with the transition to Cancelled, so we can do something similar there as well.

The Full Thing

Here is all the working code, which captures our states and valid transitions.

protocol Cancellable {}
protocol Billable {}
protocol Refundable {}

enum Pending: Cancellable, Billable {}
enum Billed: Cancellable {}
enum Packed: Cancellable, Billable {}
enum Shipped {}
enum LostInTransit: Refundable {}
enum DamagedInTransit: Refundable {}
enum Refunded {}
enum Cancelled {}

struct Order<T> {

}

extension Order where T: Billable {
  var bill: Order<Billed> {
    return Order<Billed>()
  }
}

extension Order where T: Cancellable {
  var cancel: Order<Cancelled> {
    return Order<Cancelled>()
  }
}

extension Order where T == Billed {
  var pack: Order<Packed> {
    return Order<Packed>()
  }
}

extension Order where T == Packed {
  var ship: Order<Shipped> {
    return Order<Shipped>()
  }
}

extension Order where T == Shipped {
  var lost: Order<LostInTransit> {
    return Order<LostInTransit>()
  }

  var damaged: Order<DamagedInTransit> {
    return Order<DamagedInTransit>()
  }
}

extension Order where T: Refundable {
  var refund: Order<Refunded> {
    return Order<Refunded>()
  }
}

With all these states and transitions, we can now have complex and safe transitions.

let order = Order<Pending>().bill.pack.ship.lost.refund

Further Enhancements

Here are some potential enhancements.

• Use the state types to store data; turn them into structs and store them on the machine
• Store the history of transitions on the machine
• Add vars which allow the machine to be queried e.g. var isRefundable: Bool
• Turn the transitions into functions which take arguments and return a result-type; in this way they can perform validations