Safe operations with typestate in Rust

• 2021-12-17
• 5 minutes
• Rust

Table of contents

1. First approach with runtime checks
2. Static analysis with typestate
3. Drawbacks

In my last article, I explored how you can use enum using Rust to create unambiguous JSON representations. In this one, I’ll look at another interesting feature that Rust’s types and traits system gives us.

For example, let’s say you are building a system that represents an order flow. It’s an example I love to use because we can easily image what are the most important steps in the process. Once an ordered is placed, someone needs to package the item in one or multiple boxes, then ships it to the end customer. We basically have four main states:

• Pending: the order has been created, but not yet submitted for packaging.
• Packaging: the order is being packaged in a warehouse.
• InDelivery: the order is being delivered to the customer.
• Completed: the order has been processed.

There can be a lot of other steps in between in a real life scenario, such as payment verification, or a packaged order waiting to be picked up, or cancelled because the items are not in stock. We’ll just look at those four states for the sake of simplicity here.

First approach with runtime checks

We could create an enum with a variant for each state, then have that as a field in our order struct. We can also use the fact that Rust supports data in variants to add information, such as the tracking ID when an order is being delivered.

/// Enum representing all the possible states for an order
enum OrderState {
    Pending,
    Packaging,
    InDelivery(tracking_id: String),
    Completed,
}

struct Order {
    // Order data
    pub id: String,
    pub user: User,
    pub products: Vec<Product>,
    // Current state of the order
    pub state: OrderState,
}

To transition an order from one state into another, we could implement functions that will take care of changing the state. For example, a ship function would change it into the InDelivery state.

However, there is a catch: we cannot transition from any state to any state. There is a logical order that we must follow. We can’t ship an order that completed, or submit a shipped order for packaging. For this, we can add runtime checks:

impl Order {
    fn ship(&mut self, tracking_id: String) -> Result<(), Error> {
        // Only ship if the previous state is Packaging
        if order.state != OrderState::Packaging {
            return Err::OrderStateError;
        }

        // Change the order state
        self.state = OrderState::InDelivery(tracking_id);

        Ok(())
    }
}

While this will catch issues and correctly prevent us from performing invalid transitions, it does it at runtime. That means that the code is already in production or in a testing environment.

Static analysis with typestate

This is where typestate analysis comes in. The idea of typestate is to embed state information into the type itself, and restrict which operations we can perform on a given type-state combination. Instead of adding runtime checks, we can now perform those when we are compiling our application.

Instead of using an enum, we can use a trait and structs implementing it.

/// Trait representing an order state
trait OrderState {}

// Struct that will implement that trait
#[derive(PartialEq, Clone, Debug)]
pub struct Pending;
#[derive(PartialEq, Clone, Debug)]
pub struct Packaging;
#[derive(PartialEq, Clone, Debug)]
pub struct InDelivery {
    pub tracking_id: String,
}
#[derive(PartialEq, Clone, Debug)]
pub struct Completed;

// Trait implementations
impl OrderState for Pending {}
impl OrderState for Packaging {}
impl OrderState for InDelivery {}
impl OrderState for Completed {}

/// New order struct that uses the trait instead

pub struct Order<S: OrderState> {
    // Order data
    pub id: String,
    pub user: User,
    pub products: Vec<Products>,

    // Current state of the order
    //
    // Because `InDelivery` has concrete data, we have to store what implements
    // the trait here. Otherwise, we could use `PhantomData`, which is zero-sized.
    pub state: S,
}

We can now restrict implementations to only Orders that use a specific implementation of our OrderState trait. For example, we can implement the ship function only if the Order is in the Packaging state. The benefit of that approach is that the code won’t compile if we try to ship a completed or newly created order. We have removed the ability to create potential bugs that wouldn’t match our order flow, without adding runtime checks.

impl Order<Packaging> {
    /// Ship the order
    pub fn ship(self, tracking_id: String) -> Order<InDelivery> {
        Order {
            id: self.id,
            user: self.user,
            products: self.products,
            state: InDelivery {
                tracking_id,
            },
        }
    }
}

If you’re interested to see a working implementation, you can find this repository on GitHub that models the order flow discussed in this article. You can also look at the static guarantee section of the Rust embedded book, which explains how typestate is used to prevent sending data to an input pin, or reading from an output pin on hardware devices.

Drawbacks

As a closing note, I’d like to point out that this approach is not without drawbacks.

The first one is that it adds a lot of complexity and repetition depending on how complex your state machine is. If you can accept 3 states out of 7, it’s easy to do with a match pattern in Rust. With typestate, you might need to duplicate implementation for each concrete implementation.

Furthermore, depending on how you implement this with the rest of your application, you might end up either with multiple copies of function (one for each trait variant) due to static dispatch, or with recreating the same runtime logic with dynamic dispatch. See this section of the Rust book to understand more about how Rust dispatches traits.