How to build type-safe State Machines using type-state

This will be a short one, but I hope it'll make you want to go refactor everything.

There are a lot of ways of writing state machines in typed languages, all varying in the degree of type-safety. Sure you can put everything in a record full of this option or have one variant per state and a bunch of repeated things...but today I wanna show you a way of doing state machines with something called type-state.

What this buys you is that you can offer a nice uniform API for your state machines while keeping the internal states cleanly separated.

Let's start from the beginning.

What is type-state?

Type state is kind of a weird term. You don't really have state in your types, right? Like you don't have a counter that keeps incrementing.

type counter = int
let inc : counter -> counter = fun x -> x + 1

let one: counter = inc 0
let two: counter+1?  = inc one

You can bend backwards to do this with some type schemes (which is a fancy word for "how to do something" or "pattern" that you may read around). Like a common way of doing this is to use type-params to keep adding types to something. For example:

type zero = Zero
type 'a inc = Inc of 'a

(* our `inc` function wraps whatever in an `inc` type *)
let inc : 'a -> 'a inc = fun x -> Inc x

let one: zero inc = inc Zero
let two: zero inc inc = inc one
let three: zero inc inc inc = inc two

And this sort of thing can be suuuper useful if you want to statically validate something. For example, let's say you have a fuel type and you can only accelerate your car if you have enough fuel. And enough fuel is defined as "one unit of fuel" or in types as "it's wrapped by one fuel type".

So we can start with a car without fuel, fuel it up, and then consume the fuel. And in the next example, you can see how the type annotation matches the amount of fuel a car has.

type no_fuel = No_fuel
type 'a fuel = Fuel of 'a

let fuel_up x = (Fuel (Fuel x))

let run (Fuel x) = x

let empty_car = No_fuel
let full_car: no_fuel fuel fuel = empty_car |> fuel_up
let full_car: no_fuel fuel = run full_car
let full_car: no_fuel = run full_car

(* this last one will be a type error! *)
let full_car: no_fuel = run full_car

This fails because the last full_car that worked becomes a car with no fuel, and our run function requires a car with fuel.

But alas, this can get super complex quickly because:

1. We're not used to this kind of programming in the types (unlike in programming languages like Idris or TypeScript)
2. We don't have good debugging tools for this, and type errors can become strange quickly, especially if we use polymorphic variants, objects, or GADTs.

Okay, there's more to say about type-state, but this should be enough for now: type-state lets you rule out certain behaviors by putting in your types, information about the values.

Moving on...

Type-state state machines

Alright, we're ready to go. The pattern is pretty straightforward:

1. We need a type for our thing, which will be a state machine
2. We need a type parameter for the state the machine is in
3. We want to save that state as a discrete value in our state machine type

type 'state fsm = { state: 'state }

That's it. That's the pattern. If it seems small it's because it is, but there's so much power to this pattern! Let's explore with an example.

Requesting Permissions as a Type-state machine

We'll build a tiny Permissions module that will help us check if a User has the right permissions to access a resource. There will be 3 states: Requested, Granted, and Denied. We'll always start on Requested, and we can move to Granted or Denied. If we are on Granted we should have access to the resource itself, if we are on Denied we should have access to some reason for why we were denied access.

So we start with our basics, the pattern:

type 'state t = { state: 'state };

Next, we're going to add the 3 states as distinct types:

type id
type scope

type 'resource granted = { resource : 'resource }
type denied = { reason : string }
type requested = { scopes : scope list }

Let's put this together with our t type. We will need to introduce a new type parameter for the 'resource, so we know what 'resource to use when we create our 'resource granted state. We'll also add some more metadata that is specific to a permission request but is shared across all states.

type ('state, 'resource) t = {
  (* the state of the permission request *)
  state : 'state;
  (* other shared metadata *)
  resource_id : id;
  user_id : id;
}

And now we can build our API on top of this, using our 'state type to guide the user to the functions they can use:

1. We want to create a new Permissions Request that will be requested t
2. We want to call a function that returns either a granted t with our resource, or a denied t with a reason
3. If we get a denied t we want to be able to see the reasons
4. If we get a granted t we want to be able to use the resource.

Let's get to work!

We'll start with a constructor function, and a function to transition to our final states:

let make ~resource_id ~user_id ~scopes =
  { state = { scopes }; resource_id; user_id }

let request_access t =
  match run_request t with
  | Ok resource -> Ok { t with state = { resource } }
  | Error reason -> Error { t with state = { reason } }

The basic machinery is done. We create, and the type system will infer correctly that the state should be requested, because a requested record includes scopes.

Then, our request_access function will return either (granted, 'resource) t or a (denied, 'resource) t wrapped in a result. We can use any variant for this, even a Future, but a result is already present so we'll go with that here.

Great, the next step is implementing our operations:

(* extract the reason from a `denied` state *)
let reason { state = { reason }; _ } = reason

(* do something with the resource in a `granted` state *)
let with_resource { state = {resource}; _ } fn = fn resource

(* get the resource out of a granted permission *)
let get { state = {resource}; _ } fn = resource

And done. The types get inferred nicely here as well because the records are all disjoint. Now we can use our little state machine:

(* this would be our resource *)
module Album =  struct
  type t = string
  let print = print_string
end

let _ =
  let req: (_, Album.t) Permission_request.t =
    Permission_request.make
      ~resource_id:"spotify:album:5SYItU4P7NIwiI6Swug4GE"
      ~user_id:"user:2pVEM9qgPvPeMslgnGDDOr"
      ~scopes:[ "listen"; "star"; "playlist/add"; "share" ]
  in

  match Permission_request.request_access req with
  | Ok res ->
      Permission_request.with_resource res Album.print;
      let album = Permission_request.get res in
      Album.print album;
  | Error res ->
      print_string (Permission_request.reason res)

Pretty neat, right?

For completeness sake, here's our Permission_request module:

module Permission_request = struct
  type id = string
  type scope = string
  type 'resource granted = { resource : 'resource }
  type denied = { reason : string }
  type requested = { scopes : scope list }

  type ('state, 'resource) t = {
    (* the state of the permission request *)
    state : 'state;
    (* other shared metadata *)
    resource_id : id;
    user_id : id;
  }

  let make ~resource_id ~user_id ~scopes =
    { state = { scopes }; resource_id; user_id }

  (* TODO(@you): you can implement here the logic for checking
   * if you actually have permissions for this request :)
   *)
  let run_request : (requested, 'resource') t -> ('resource, string) result =
   fun _t -> Error "unimplemented!"

  let request_access t =
    match run_request t with
    | Ok resource -> Ok { t with state = { resource } }
    | Error reason -> Error { t with state = { reason } }

  let reason { state = { reason }; _ } = reason
  let with_resource { state = { resource }; _ } fn = fn resource
  let get { state = {resource}; _ } = resource
end

Conclusions: Type-state is Great

Type-state is just another tool in your bat-belt to build great APIs with type-safety, and good ergonomics.

It has a cost, like everything else, but it enables you to grow your states and maintain them separately by using submodules, in a very natural way. After all they are different types!

I picked up this pattern from Rust libraries that model the state of sockets, database connections, and many other things, by doing exactly the same thing. Of course, having traits in Rust means we can extend the methods for a specific type-state, which means narrowing down even further the information you have to process when you go type req. and get the autocompletion list.

But at least in OCaml we can implement the pattern safely and maybe with time our autocompletion will get smarter and do the same filtering for us!

That's it. That's type-states for OCaml.

Have you implemented typed state machines in some other ways? Have anything to add or challenge? I'd love to hear it! Join the x.com thread.

Happy Cameling! 🐫