The TypeState Pattern

Preface: In this post I’ll go through TypeState, a programming design pattern that I’m learning. Using the TypeState pattern can provide added security to our code, in the form of compile-time checks ¹. I want to point to the following excellent resources which have guided my learning so far: ², and ³, and ⁴ .

Code: The complete source code for this post is available at the end of this post.

The Art of casting Spells

When learning a new concept I often try to come up with an imaginary example. I feel that it helps me remember things, and generally makes learning more fun. Today, the concept is TypeState and I’m imagining a wizard casting spells.

In the classic Role-Playing Game, Dungeons & Dragons (or D&D RPG), magic is a powerful force.
In order to create a balanced game, there has to be limits to the use of magic.
The rules of D&D states that:

It takes Magical Energy to cast spells.
A wizard has limited magical energy and needs to “recharge” once it has been spent
The way a wizard can “recharge” magical energy is with a full nights sleep

Magic in Dungeons & Dragons is sorted using a ranking system, with spells being categorized by how powerful they are. There is a greater “cost” in magical energy to cast a 5th level spell than a 1st level spell, and a wizard has to attain a certain level to be able to cast some powerful spells.

Next, let’s have a look at how we can model a wizards magic energy using types.

State Machine: Magic Energy

We can take the the constraints placed by the D&D rules on Wizards and their access to magic energy, and express it as a state machine:

Rested
“If magic energy is not exhausted” -> “Can cast spells”
“Else if magic energy is exhausted” => Tired
Tired
“If has no magic energy to cast spells with” -> “Cannot cast spells”
“If go to sleep” => Sleep
Sleep
“If hasn’t slept a full nights sleep” -> “Cannot cast spells”
“If has slept a full nights sleep” => Rested

I’m not a D&D player mayself but let’s assume our state machine is correct. So, next, let’s use Rust’s mighty type system and the TypeState pattern to express the constraints previously laid out.
Given a struct Wizard, our goal is:

Depending on the state of the Wizard, some methods should be callable.
Depending on the state of the Wizard, some methods should not be callable.
We want to declare how a wizard can move from one state to another.

Okay, neat - but how is the TypeState pattern actually implemented?
Before we look at the final code, let’s look at a problem I faced while coding this.

Peano Arithmetic to the rescue

While working on this post I got stuck on a problem: how can the magic energy level of a wizard be persisted when transitioning between two typestates, Wizard<Rested> and Wizard<Tired>?

If the wizard has magic energy left then it’s able to continue casting spells and vice versa - no magic energy, no spell casting, bro. I explored having an enum being returned when wizard.cast_spell() is cast. Unfortunately, the Rust compiler wasn’t having it.

// We want to check if the wizard has enough magic energy to cast spells.
enum EnergyCheck {
    CanCastAgain(Wizard<Rested>),
    NeedsRest(Wizard<Tired>),
}

// Every time the wizard casts a spell we reduce its magic_energy.
// If the energy is reduced to zero, the wizard becomes tired and
// cannot cast any spells until rested.
impl Wizard<Rested> {
    fn cast_spell(mut self) -> EnergyCheck {
        match self.energy > 0 {
            true => {
                self.energy = self.energy - 1;
                println!(
                    "The rested Wizard cast a spell and now
                    has {} magic energy points left.",
                    self.energy
                );
                return EnergyCheck::CanCastAgain(self);
            }
            false => {
                println!(
                    "The Wizard is exhausted and 
                    has to sleep to restore its magic energy."
                );

                EnergyCheck::NeedsRest(Wizard::<Tired> {
                    energy: self.energy,
                    rested_energy: self.rested_energy,
                    _state: marker::PhantomData,
                })
            }
        }
    }
}

After tinkering with this for a while I decided to ask on the Rust Discord server. Thanks to Rust Discord members; Morgane for a solution, and Kyuuhachi an improvement using Peano Arithmetic.

This Rust code defines a type-level natural number system using traits and structs, commonly used in functional programming and type-level programming to represent and manipulate natural numbers at the type level. Here’s a breakdown of the code:

// We define the trait Nat to represent Natural numbers at the type level.
// VAL defines the numeric type of the number system.
trait Nat { const VAL: u8; }

// We represent the value zero.
struct Z;

// S representes a _successor type_, the next natural number after N.
// If N == 1, then S<N> == 2.
struct S<N: Nat>(PhantomData<N>);

// We implement the trait of being a natural number for Z with the value 0.
// This makes Z the base case of this type-level number system.
impl Nat for Z { const VAL: u8 = 0; }

// We implement the trait of being the next natural number for S<T>.
// S<N> gets its value from T::VAL + 1, where T::VAL is the value 
// of the previous neighboring number.
// Z::VAL = 0
// S<Z> => 0 + 1 = 1
// S<S<Z>> => (0 + 1) + 1 = 2
impl<N: Natu> Nat for S<N> { const VAL: u8 = N::VAL +1; }

// Usage:
// type Zero = Z; represents 0
// type One = S<Zero>; represents 1
// println!("Usage {}", Zero::VAL);
// println!("Usage {}", One::VAL);

Did this solve everything?
Well, almost. The current solution has no forks in the state-path. When I have more time I’d like to rewrite the code and use an Enum to model a fork of the state path depending on how much magic energy the wizard has after casting a spell. If the wizard has more then 0 energy, return Wizard<Rested>, if the wizard has 0 energy left, return Wizard<Tired>.
But for now, this will have to do.

🧙‍♂️ Creating a Wizard

Style: I’m personally not a fan of examples where the code is taken out of context and interspaced with text - I prefer to have one single blob of source code with explanatory comments sprinkled throughout. Please note that these comments are a bit chatty, but only because this is a blog post and not source from an actual project.

// File: main.rs
use std::marker::PhantomData;

////////////////////////////////////////////////////////////////////////////////
// TypeStates
#[derive(Debug)]
struct Rested<Energy> {
    _energy: PhantomData<Energy>,
}

#[derive(Debug)]
struct Tired<Energy> {
    _energy: PhantomData<Energy>,
}

#[derive(Debug)]
struct Sleeping;

////////////////////////////////////////////////////////////////////////////////
// Peano Arithmetic
trait Nat {
    const VAL: u8;
}
struct Z;
impl Nat for Z {
    const VAL: u8 = 0;
}

struct S<N: Nat>(PhantomData<N>);
impl<N: Nat> Nat for S<N> {
    const VAL: u8 = N::VAL + 1;
}

// Quick macro to give Peano numbers less of a nested look.
// Turns this:
// let gandalf = new_wizard::<S<S<S<Z>>>>();
//
// ...into this:
// let gandalf = new_wizard::<Three>();
macro_rules! define_peano {
    ($($name:ident = $value:ty),*) => {
        $(
            type $name = $value;
        )*
    };
}

define_peano! {
    Zero = Z,
    One = S<Z>,
    Two = S<One>,
    Three = S<Two>
}

////////////////////////////////////////////////////////////////////////////////
// Transition between states based on method calls and energy level.
#[derive(Debug)]
struct Wizard<State, Energy> {
    _state: PhantomData<State>,
    _energy: PhantomData<Energy>,
}

////////////////////////////////////////////////////////////////////////////////
// Failed to get this impl design to work. Not sure how to solve this.
/*
impl<T> Wizard<Rested<T>, T> {
    fn new() -> Self {
        Wizard {
            _state: PhantomData,
            _energy: PhantomData,
        }
    }
}
*/

////////////////////////////////////////////////////////////////////////////////
// Constuctor
fn new_wizard<T>() -> Wizard<Rested<T>, T> {
    Wizard {
        _state: PhantomData,
        _energy: PhantomData,
    }
}

// Return the magic energy level of the wizard.
impl<N: Nat, Energy> Wizard<Rested<N>, Energy> {
    const fn energy(&self) -> u8 {
        N::VAL
    }
}

impl<N: Nat, Energy> Wizard<Rested<S<N>>, Energy> {
    fn cast_spell(self) -> Wizard<Rested<N>, Energy> {
        println!("Spell was cast!");
        Wizard {
            _state: PhantomData,
            _energy: PhantomData,
        }
    }
}

impl<Energy> Wizard<Rested<Z>, Energy> {
    fn go_to_sleep(self) -> Wizard<Sleeping, Energy> {
        Wizard {
            _state: PhantomData,
            _energy: PhantomData,
        }
    }
}

impl<Energy> Wizard<Sleeping, Energy> {
    fn wakeup(self) -> Wizard<Rested<Energy>, Energy> {
        Wizard {
            _state: PhantomData,
            _energy: PhantomData,
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
fn main() {
    let gandalf = new_wizard::<Three>();
    println!("Wizard created with energy: {}", gandalf.energy(),);

    let gandalf = gandalf.cast_spell();
    println!("Gandal's energy is: {}", gandalf.energy());

    let gandalf = gandalf.cast_spell();
    println!("Wizard's energy is: {}", gandalf.energy());

    println!("Wizard trying to cast spell with zero energy.");

    let gandalf = gandalf.cast_spell();

    if gandalf.energy() == 0 {
        println!("Wizard has 0 energy and is tired. Needs sleep.");
    }

    let gandalf = gandalf.go_to_sleep();
    println!("Wizard is sleeping.");

    let gandalf = gandalf.wakeup();
    println!("Wizard woke up from sleeping.");
    println!("Wizard slept and now has energy: {}", gandalf.energy());

    let gandalf = gandalf.cast_spell();
    println!("Wizard cast a first spell.");
    println!("Wizard now has energy: {}", gandalf.energy());
}

In computer science, compile-time checking for undefined behavior refers to the detection of potential issues and violations of the language rules during the compilation process. The compiler analyzes the code and performs static checks to identify situations that may lead to undefined behavior at runtime. ↩︎
https://cliffle.com/blog/rust-typestate ↩︎
https://yoric.github.io/post/rust-typestate/ ↩︎
https://doc.rust-lang.org/std/marker/struct.PhantomData.html ↩︎

The TypeState Pattern

December 10, 2024 by CryptoPatrick

The Art of casting Spells

State Machine: Magic Energy