Type State Pattern

Type State Pattern

Özgün Özerk
Özgün Özerk

Safe API

Imagine we are writing a library for a game.

For this library, we will expose a player object. The player object will have the following methods:

  • new() -> to create new Player object
  • buy_item()
  • sell_item()
  • list_items()
  • die()
  • get_resurrected()


pretty usual stuff so far. Where things get interesting is here: we want some methods to be callable on the dead player objects, some on the alive player objects, and some on the both.

Dead:

  • get_resurrected()


Alive:

  • die()
  • sell_item()
  • buy_item()


both:

  • list_items() -> we can list the items of the dead players, why not…

Designing the API

Say, our Player object is defined as such:

struct Player {
    username: String,
    health: u8,
    items: Vec<String>,
}

now, if we implement all the above methods for the Player struct, like this:

impl Player {
    fn die(&self) {
        todo!()
    }

    // and other methods below
}

we can call die() multiple times on the Player, which makes no sense.

I’ll use die() as an example throughout this article.

But the same principles apply to all other methods.

Going back to our API, it is definitely not safe! We should do better…

Solution candidate 1:

we can introduce an additional field to the struct alive:

struct Player {
    username: String,
    health: u8,
    items: Vec<String>,
    alive: bool // it is here in case you are blind
}


and we can do runtime checks when necessary:

impl Player {
    fn die(&self) {
        if self.alive {
            todo!()
        } else {
            panic!("🤷‍♀️ cannot die twice 🤷‍♀️")
        }
    }
}


this will work… But it’s disgusting. And surprisingly, not only because of the emoji:

  • runtime check? we can do compile-time check 🤓
  • we introduced another field to the struct and now it will occupy more memory (kind of tolerable)
  • we increased runtime complexity. I hate the fact that we will now do an extra check before each call to the methods of Player to be able to determine if we can call them in the first place… (intolerable)

Solution candidate 2:

One approach for achieving compile-time checks, might be having two different structs for alive and dead player, and have the necessary methods implemented for them respectively.

struct AlivePlayer {
    username: String,
    health: u8,
    items: Vec<String>,
}

struct DeadPlayer {
    username: String,
    health: u8,
    items: Vec<String>,
}

impl DeadPlayer {
    // no die method here
}

impl AlivePlayer {
    fn die(self) -> DeadPlayer  {
        todo!()
    }
}


better than the previous solution. But still bad

why it’s bad:

  • the API is not clean. We are storing the same fields in both Dead and Alive player, while they are both just Players.
  • the end-user has to know when to create an instance of Alive player and Dead player. It might be simple to guess in this example, but imagine much more complex/abstract types. If possible, our API should be responsible for when to use which type, not the end user.

Solution candidate 3:

This is called type state pattern, and I really like it.

We will introduce type as a generic parameter for our Player struct, and we will define methods accordingly.

struct Player<State> {
    username: String,
    health: u8,
    items: Vec<String>,
    _state: PhantomData<State>, // i'll explain this, wait
}

struct Alive;
struct Dead;

impl Player<Alive> {
    pub fn die(self) -> Player<Dead> {
        Player {
            username: self.username,
            health: 0,
            items: self.items,
            _state: PhantomData,
        }
    }

    // and other methods
}

impl Player<Dead> {
    // no `die` method here :)
}


What’s going on here:

  • by making our struct generic, we can have an alive and a dead version of it.
  • alive and dead are represented by unit structs, allowing us to use them as types.
  • die() method is only implemented for the Player<Alive>, and it is returning Player<Dead> after consuming the self. In other words, alive player is gone.
  • for the Player<Dead>, we don’t have the die() method, making it impossible to call die() on a dead player.
  • we achieve compile-time warranty of die() method being not callable on a dead player.

What the hell is state: PhantomData<State> in the Player struct?:

If I hadn’t put that, it wouldn’t compile.

Rust does not allow us to declare struct Player<State> and not use State in any of the fields.

But, we don’t want to use the State in the fields. We don’t want it to put extra memory on runtime.

Luckily, PhantomData is is a zero-sized-type, meaning that, it won’t be in the compiled code. No runtime effects will be present, totally zero cost! It just acts like it owns/uses State, so that compiler can shut up about:

In the end, we did introduce a new field to the Player struct, but only in compile-time. It does not exist in runtime.

Very cool, but not cool enough.

The small problem:

What happens if the user wanted to mess around and provided their own type to the Player struct:

struct Player<State> {
    username: String,
    health: u8,
    items: Vec<String>,
    _state: PhantomData<State>,
}

fn main() {
    let player = Player::<u8> {
        username: "potato".to_string(),
        health: 100,
        items: vec!["frostmourne".to_string()],
        _state: PhantomData,
    };
}


The Final Solution

By applying a trait bound to the generic, we can also solve the above problem:

struct Alive;
struct Dead;

trait PlayerState {}

impl PlayerState for Alive {}
impl PlayerState for Dead {}

struct Player<State>
where
    State: PlayerState,
{
    username: String,
    health: u8,
    items: Vec<String>,
    _state: PhantomData<State>,
}


Now, the user won’t be able to do this:

let player = Player::<u8> {
        username: "potato".to_string(),
        health: 100,
        items: vec!["frostmourne".to_string()],
        _state: PhantomData,
    };


The Final Final Solution

Actually, the user still can do this:

let player = Player::<u8> {
    username: "potato".to_string(),
    health: 100,
    items: vec!["frostmourne".to_string()],
    _state: PhantomData,
}

by implementing PlayerState for u8:

impl PlayerState for u8 {}

To prevent this, we can use a neat trick called sealed traits. If this is the first time you hear about sealed traits, here are 2 great articles about it:

  • https://predr.ag/blog/definitive-guide-to-sealed-traits-in-rust/
  • https://rust-lang.github.io/api-guidelines/future-proofing.html

Basically, we “seal” the trait, so that no one can implement it.

We will rewrite our PlayerState trait as such:

mod private {
    pub trait Sealed {}
}

trait PlayerState: private::Sealed {}

struct Alive;
struct Dead;

impl private::Sealed for Alive {}
impl private::Sealed for Dead {}

impl PlayerState for Alive {}
impl PlayerState for Dead {}

Now, the user can’t implement PlayerState for any type, because the PlayerState depends on Sealed trait, which is private and can only be accessed internally from our library.

Full Code

Here, have a big chunk of code:

#![allow(unused)]

use std::marker::PhantomData;

mod private {
    pub trait Sealed {}
}

trait PlayerState: private::Sealed {}

struct Alive;
struct Dead;

impl private::Sealed for Alive {}
impl private::Sealed for Dead {}

impl PlayerState for Alive {}
impl PlayerState for Dead {}

struct Player<State = Alive>
where
    State: PlayerState,
    // or `Player<State: PlayerState = Alive>`
{
    username: String,
    health: u8,
    items: Vec<String>,
    _state: PhantomData<State>,
}

impl Player<Alive> {
    pub fn die(self) -> Player<Dead> {
        Player {
            username: self.username,
            health: 0,
            items: self.items,
            _state: PhantomData,
        }
    }

    pub fn sell_item(&mut self, item_idx: usize) {
        let _ = &self.items.remove(item_idx);
    }

    pub fn buy_item(&mut self, item: String) {
        self.items.push(item);
    }
}

impl Player<Dead> {
    pub fn get_resurrected(self) -> Player<Alive> {
        Player {
            username: self.username,
            health: 100,
            items: self.items,
            _state: PhantomData,
        }
    }
}

impl<State: PlayerState> Player<State> {
    pub fn list_items(&self) -> &[String] {
        self.items.as_slice()
    }
}

impl Player {
    pub fn new(username: String) -> Self {
        Player {
            username,
            health: 100,
            items: Default::default(),
            _state: PhantomData,
        }
    }
}

fn main() {
    let player = Player::new("HiMom".to_owned()); // create an `Alive` player

    /* the below commented out code creates compile time error */
    /*****************/
    // player.get_resurrected(); // no method named `get_resurrected()` found for struct `Player` in the current scope

    let player = player.die(); // we can kill a an `Alive` player

    /*  CANNOT kill what's already dead */
    /*****************/
    // player.kill(); // no method named `kill` found for struct `Player<Dead>` in the current scope

    player.list_items(); // can list items for `Alive` or `Dead` player
    let player = player.get_resurrected(); // can resurrect a `Dead` player
    player.list_items(); // can list items for `Alive` or `Dead` player;

    /* we can ensure the TraitBound also in compile-time, the below code creates the error: */
    /*  the trait bound `u8: PlayerState` is not satisfied */
    /*****************/
    // let invalid_player = Player::<u8> {
    //     username: "potato".to_string(),
    //     health: 100,
    //     items: vec!["frostmourne".to_string()],
    //     state: PhantomData,
    // };
}


Some notes on the code

First note:

if this is new syntax to you: struct Player<State = Alive>, it means, if the generic type is not mentioned, assume the default value would be Alive.

So, this code:

impl Player {
    pub fn new(username: String) -> Self {
        Player {
            username,
            health: 100,
            items: Default::default(),
            _state: PhantomData,
        }
    }
}


does not mention the generic, and thus, it is actually equivalent to:

impl Player<Alive> { // notice the `Alive` here
    pub fn new(username: String) -> Self {
        Player {
            username,
            health: 100,
            items: Default::default(),
            _state: PhantomData,
        }
    }
}


Then, why did I not put new method in impl Player<Alive> block next to other alive methods? Because, I like constructors in a separate place, completely subjective reason. Do as you please.

Second note:

impl<State: PlayerState> Player<State> {
    pub fn list_items(&self) -> &[String] {
        self.items.as_slice()
    }
}


this means, as long as State implements PlayerState trait, list_items should be available. In other words, list_items is only available for Alive and Dead players.

Is it worth it?

Let’s talk about the elephant in the room: is it really worth writing all this code? I honestly think no.

The good looks ugly, it’s much less maintainable, and overall it provides a worse developer experience.

But the promises are also hard to turn your back to:

  • compile-time type safety
  • state based methods
  • better auto-completion (much better DevX in terms of this)

That’s why I wrote this crate: state-shift

Now it is worth it with State-Shift

The above code can be written in a much simpler way with state-shift.

Revised code:

use std::marker::PhantomData;

use state_shift::{require, states, switch_to, type_state};

#[type_state(state_slot = 1, default_state = Alive)]
struct Player {
    username: String,
    health: u8,
    items: Vec<String>,
}

#[states(Alive, Dead)]
impl Player {
    #[require(Alive)] // requiring the default state for the constructor
    pub fn new(username: String) -> Self {
        Player {
            username,
            health: 100,
            items: Default::default(),
        }
    }

    #[require(Alive)]
    #[switch_to(Dead)]
    pub fn die(self) -> Player {
        Player {
            username: self.username,
            health: 0,
            items: self.items,
        }
    }

    #[require(Alive)]
    pub fn sell_item(&mut self, item_idx: usize) {
        let _ = &self.items.remove(item_idx);
    }

    #[require(Alive)]
    pub fn buy_item(&mut self, item: String) {
        self.items.push(item);
    }

    #[require(Dead)]
    #[switch_to(Alive)]
    pub fn get_resurrected(self) -> Player {
        Player {
            username: self.username,
            health: 100,
            items: self.items,
        }
    }

    #[require(A)]
    pub fn list_items(&self) -> &[String] {
        self.items.as_slice()
    }
}

fn main() {
    let player = Player::new("HiMom".to_owned()); // create an `Alive` player

    let player = player.die(); // we can kill a an `Alive` player

    player.list_items(); // can list items for `Alive` or `Dead` player
    let player = player.get_resurrected(); // can resurrect a `Dead` player
    player.list_items(); // can list items for `Alive` or `Dead` player;
}

So, we get all the benefits of the type-state pattern, but without the boilerplate.

If you like this article and the state-shift crate, I’d greatly appreciate your stars on GitHub!

Type State Pattern
Prev post

Unsafe Rust
Next post

Mutability

Type State Pattern

You may also like

See all generic