Learn Rust Series (#29) - Typestate Programming: State Machines in the Type System

scipio(71)

Published in

StemSocial

Words

2987

Reading

14 min

Listen

Play

Learn Rust Series (#29) - Typestate Programming: State Machines in the Type System

What will I learn

You will learn what typestate programming is: encoding an object's current state into its type;
how to make invalid operations, like using a resource in the wrong state, a compile error rather than a runtime one;
how tiny state-marker types plus PhantomData carry state at zero runtime cost;
how a transition consumes self and hands back a value in the new state, so the old state can never be touched again;
how the same pattern powers foolproof builders and honest protocol state machines.

Requirements

A working modern computer running macOS, Windows or Ubuntu;
An installed Rust toolchain (via rustup, from rustup.rs);
The previous twenty-eight episodes, especially ownership (episode 3), generics (episode 8) and marker types (episode 25);
The ambition to learn systems programming from the ground up.

Difficulty

Beginner

Curriculum (of the `Learn Rust Series`):

Learn Rust Series (#29) - Typestate Programming: State Machines in the Type System

We close the long trait-and-type-system phase with its most striking application, and, I would argue, its most satisfying one: making illegal states unrepresentable. At the end of last episode I promised we would take the same instinct that let us seal a trait -- writing our intent down where the compiler can hold the line -- and point it at a completely different problem: not "who may implement this", but "when is this operation even allowed". That is what typestate programming does. It encodes an object's current state into its type, so the compiler only ever offers you the operations that are legal in that state. Call the wrong method for the state you are in, and the code does not compile. Not "panics at runtime", not "returns an Err you might forget to check" -- it does not compile at all.

Here is the whole idea in one sentence, and then we spend the rest of the episode earning it. A value moving through a lifecycle -- a connection that is disconnected then connected, a request that has no url and then has one, a file opened for reading versus writing -- is really a little state machine, and Rust lets you draw that state machine directly in the type system so the invalid transitions become code you literally cannot write. It is the purest form of "parse, don't validate": once you hold a value of the right type, the guarantee is already true, and no defensive check is needed because the wrong state never got that far ;-)

Let me clear last episode's homework first, as always, and then we build our first door.

Solutions to Episode 28 Exercises

Episode 28 was sealed traits and stable APIs -- using a private supertrait to control exactly who may implement your public trait, plus #[non_exhaustive] for enums and structs. Three exercises, each with full runnable code below.

Exercise 1 asked you to seal a Format trait with a pub trait Sealed {} in a private module, implement it for two of your own unit structs (Json and Xml), and then confirm that an outside type cannot join the club:

mod sealed { pub trait Sealed {} }

pub trait Format: sealed::Sealed {
    fn render(&self) -> String;
}

struct Json;
struct Xml;

impl sealed::Sealed for Json {}
impl sealed::Sealed for Xml {}

impl Format for Json { fn render(&self) -> String { "{}".into() } }
impl Format for Xml { fn render(&self) -> String { "".into() } }

fn main() {
    println!("{}", Json.render()); // {}
    println!("{}", Xml.render());  // 
    // impl Format for String {}   // would NOT compile: String: sealed::Sealed is not satisfied
}

The moment you try impl Format for String in a downstream crate, the compiler asks for String: sealed::Sealed, and there is no way to provide it, because sealed is a private module nobody outside can reach. That is the whole seal, in one bound.

Exercise 2 wanted you to add a second method to a sealed trait, updating your own impls in the same edit, and to say in a sentence why that is not a breaking change:

mod sealed { pub trait Sealed {} }

pub trait Format: sealed::Sealed {
    fn render(&self) -> String;
    fn mime(&self) -> &str; // added later, in the same release
}

struct Json;
impl sealed::Sealed for Json {}
impl Format for Json {
    fn render(&self) -> String { "{}".into() }
    fn mime(&self) -> &str { "application/json" }
}

fn main() {
    let j = Json;
    println!("{} {}", j.render(), j.mime()); // {} application/json
}

The one-sentence reason: because the trait is sealed, you own every implementation that exists, so you patch them all in the same edit and no downstream impl is ever left missing the new method -- whereas on an open trait some stranger's impl would suddenly fail to compile, which is exactly what makes it a breaking change.

Exercise 3 asked for a #[non_exhaustive] struct with a builder-style constructor (a new plus a with_* method), then to add a brand-new field and confirm old callers keep compiling:

#[non_exhaustive]
pub struct Config {
    pub level: u8,
    pub verbose: bool, // imagine this was added in a later release
}

impl Config {
    pub fn new() -> Config {
        Config { level: 1, verbose: false }
    }
    pub fn with_level(mut self, n: u8) -> Config {
        self.level = n;
        self
    }
}

fn main() {
    let c = Config::new().with_level(3); // downstream must go through this, not a literal
    println!("{} {}", c.level, c.verbose); // 3 false
}

Because #[non_exhaustive] forbids downstream struct-literal construction, callers reach Config only through new, which fills in every field. Adding verbose therefore breaks nobody. The line that would have broken, had you allowed it, is a downstream Config { level: 3 } literal -- it names every field, so a new field makes it incomplete.

Right, homework cleared. Now for the main event ;-)

The idea: state as a type parameter

The trick has three moving parts, and once you see them together the rest is variation. First, you add a generic parameter to your struct that stands for its state. Second, you fill that parameter with tiny marker types, one per state -- empty structs like struct Open; that exist only to be names. Third, you hold the marker with PhantomData, a zero-sized type from the standard library that lets a struct "carry" a type parameter it does not otherwise store any data for (we will study PhantomData and zero-sized types much more closely a bit further down the road; for now, read it as "remember this type, cost nothing at runtime"). Then you write one impl block per state, and each block exposes only the operations that are legal in that state.

A door is the canonical example, so let us build one:

use std::marker::PhantomData;

struct Open;
struct Closed;

struct Door<State> {
    _state: PhantomData<State>,
}

impl Door<Closed> {
    fn new() -> Door<Closed> {
        Door { _state: PhantomData }
    }
    fn open(self) -> Door<Open> {
        println!("opening the door");
        Door { _state: PhantomData }
    }
}

impl Door<Open> {
    fn walk_through(&self) {
        println!("walking through");
    }
    fn close(self) -> Door<Closed> {
        println!("closing the door");
        Door { _state: PhantomData }
    }
}

fn main() {
    let door = Door::<Closed>::new();
    let door = door.open();    // Closed -> Open
    door.walk_through();       // this method exists ONLY on Door
    let _door = door.close();  // Open -> Closed
    // door.walk_through();     // would NOT compile: no such method on Door
}

Look carefully at what the compiler now enforces, because it is more than it first appears. walk_through is defined only inside impl Door<Open>, so you cannot call it on a closed door -- the method genuinely does not exist for that type. open lives only on Door<Closed>, so you cannot open an already-open door. Door<Open> and Door<Closed> are, as far as the type checker is concerned, two different types that happen to share a name, and each one publishes a different menu of methods. Invalid operations are not runtime errors you must remember to guard against; they are compile errors you cannot even express. That is a categorical improvement over a runtime if self.is_open { ... } check, because a check you can forget to write is a bug waiting to happen, and a method that does not exist is not.

Transitions consume self

Notice a detail that is easy to skim past: the transition methods take self by value, not &self, and return a brand-new Door. That is not an accident, it is the load-bearing wall of the whole pattern. By consuming the old value, the transition guarantees the door cannot still be used in its previous state after it has moved on, because ownership of the old state is gone -- this is the ownership model from episode 3 doing double duty as state-machine safety:

use std::marker::PhantomData;
struct On; struct Off;
struct Switch<S> { _s: PhantomData<S> }
impl Switch<Off> {
    fn new() -> Switch<Off> { Switch { _s: PhantomData } }
    fn turn_on(self) -> Switch<On> { Switch { _s: PhantomData } } // consumes the Off switch
}
impl Switch<On> {
    fn turn_off(self) -> Switch<Off> { Switch { _s: PhantomData } }
}
fn main() {
    let s = Switch::<Off>::new();
    let s = s.turn_on();   // the Off value is moved away, gone for good
    let _s = s.turn_off();
}

Try, in your head, to reach for the old Off switch after turn_on. You cannot: the value was moved into turn_on, the borrow checker will not let you touch it, and the only thing you hold now is the fresh Switch<On>. You physically cannot keep both the old and the new state around at once. A transition is a move from one typed state into another, and that move is what makes the guarantee airtight rather than merely polite.

Carrying real data across states

The door and the switch carry no real data, which makes them clean teaching toys but slightly unrealistic. Real resources hold something -- a file handle, a socket, a title -- and that payload has to survive the transitions while the state marker changes around it. That is easy: keep your real fields as normal fields, and move them forward in each transition:

use std::marker::PhantomData;

struct Draft;
struct Published;

struct Post<State> {
    title: String,           // real data, carried through every state
    _state: PhantomData<State>,
}

impl Post<Draft> {
    fn new(title: &str) -> Post<Draft> {
        Post { title: title.to_string(), _state: PhantomData }
    }
    fn publish(self) -> Post<Published> {
        Post { title: self.title, _state: PhantomData } // move the title forward
    }
}

impl Post<Published> {
    fn view(&self) -> String {
        format!("[live] {}", self.title)
    }
}

fn main() {
    let post = Post::new("Typestate").publish();
    println!("{}", post.view()); // [live] Typestate
}

The title rides along untouched from Draft into Published, because publish moves self.title into the new value. Only the phantom state changed. And view is available only after publishing, so you can never render a post that is still a draft. This is the shape you will actually use in production: the marker types are pure compile-time bookkeeping, and your genuine data lives in ordinary fields alongside the PhantomData.

A builder that cannot be misused

The single most common real-world use of typestate is builders that must not be finished until the required pieces are in place. We saw the classic builder problem in passing when we discussed API design -- a builder that lets you call .build() before you have set a mandatory field, and then panics or returns an error at runtime. Typestate deletes that failure mode. Encode "the url has been provided" as a state, make send available only in that state, and forgetting the url stops being possible:

use std::marker::PhantomData;

struct NoUrl;
struct HasUrl;

struct Request<State> {
    url: Option<String>,
    _state: PhantomData<State>,
}

impl Request<NoUrl> {
    fn new() -> Request<NoUrl> {
        Request { url: None, _state: PhantomData }
    }
    fn url(self, u: &str) -> Request<HasUrl> {
        Request { url: Some(u.to_string()), _state: PhantomData }
    }
}

impl Request<HasUrl> {
    fn send(&self) -> String {
        format!("GET {}", self.url.as_ref().unwrap())
    }
}

fn main() {
    let response = Request::new().url("https://example.com").send();
    println!("{response}"); // GET https://example.com
    // Request::new().send(); // would NOT compile: no `send` on Request
}

Two things are worth dwelling on. First, send lives only in impl Request<HasUrl>, so Request::new().send() with no url call in between simply does not compile -- the method is not on that type. Second, and this is the quietly beautiful part, the unwrap inside send can never panic. The only route to the HasUrl state is through url, and url always sets the field, so by the time you can call send the Option is guaranteed Some. The type system has made the invariant true by construction, which means the unwrap is not a gamble, it is a statement of a fact the compiler already proved. That is what "make illegal states unrepresentable" buys you: not just blocked mistakes, but the disappearance of whole classes of defensive code.

A fixed-cycle state machine

State machines with a fixed cycle fall out of this pattern almost for free. Give each state exactly one legal transition, and the type of each state advertises only that one next:

use std::marker::PhantomData;

struct Red; struct Green; struct Yellow;
struct Light<S> { _s: PhantomData<S> }

impl Light<Red> {
    fn start() -> Light<Red> { Light { _s: PhantomData } }
    fn next(self) -> Light<Green> { println!("go"); Light { _s: PhantomData } }
}
impl Light<Green> {
    fn next(self) -> Light<Yellow> { println!("slow"); Light { _s: PhantomData } }
}
impl Light<Yellow> {
    fn next(self) -> Light<Red> { println!("stop"); Light { _s: PhantomData } }
}

fn main() {
    let l = Light::<Red>::start();
    let l = l.next(); // go
    let l = l.next(); // slow
    let _l = l.next(); // stop
}

There is no way to jump from Red straight to Yellow, because the next on Light<Red> returns a Light<Green> and nothing else. The legal path through the machine is baked into the return types, one step at a time. Notice too that each state defines its own next with a different return type -- these are not overrides of a shared method, they are three separate inherent methods that merely share a name, one per state type. The compiler picks the right one based on which state you are holding.

Modelling a real protocol

Where this genuinely pays for itself is protocols, and the network connection is the textbook case. You must not send on a connection that is not yet connected, and with typestate that rule is not a runtime assertion, it is the shape of the API:

use std::marker::PhantomData;

struct Disconnected;
struct Connected;
struct Conn<S> { _s: PhantomData<S> }

impl Conn<Disconnected> {
    fn new() -> Conn<Disconnected> { Conn { _s: PhantomData } }
    fn connect(self) -> Conn<Connected> { Conn { _s: PhantomData } }
}
impl Conn<Connected> {
    fn send(&self, msg: &str) { println!("sending {msg}"); }
    fn disconnect(self) -> Conn<Disconnected> { Conn { _s: PhantomData } }
}

fn main() {
    let c = Conn::<Disconnected>::new().connect();
    c.send("hello"); // `send` exists only on Conn
    let _c = c.disconnect();
}

Having said that, notice how disconnect consumes the connected value and returns a Conn<Disconnected> -- so once you have disconnected, send vanishes again, and you would have to connect afresh to get it back. The lifecycle is a loop drawn in types, and every point on the loop offers exactly the operations that make sense there. This is the same skeleton you would use for a file opened read-only versus read-write, a transaction that is pending versus committed, or a hardware pin configured as input versus output in embedded code. Different domains, one shape.

Where typestate earns its keep

Now the honest part, because typestate is not free and it is not for every struct. Each state is another marker type and another impl block, and that is real ceremony you are adding to your code. You want to spend it only where misuse is both likely and costly. It shines for genuine protocols and lifecycles -- the connection that is open or closed, the file opened for reading versus writing, the transaction pending or committed, the embedded pin as input or output. In every one of those, calling a method in the wrong state is a real bug with real consequences, and moving that check from runtime to compile time is worth the extra types.

For a plain value type with no meaningful lifecycle, though, typestate is over-engineering. If the "states" are really just a boolean flag nobody will misread, a normal field and an ordinary method are lighter and clearer, and a reader will thank you for not making them chase four marker types to understand a switch. The judgement, as it so often is in Rust, comes down to a single question: is the guarantee worth the machinery? For stateful resources where a wrong call is expensive, the answer is very often yes. For a humble data holder, it is usually no. Knowing which is which is the actual skill here, more than the mechanics of PhantomData.

How other languages handle this

Since quite some of you came to Rust through the Learn Python Series (as did I), a look sideways sharpens the picture -- and here the contrast is stark, because most languages simply cannot express this at all.

Python has no way to make a method appear or disappear based on state. The nearest you can do is check a flag at runtime and raise, which is exactly the failure mode typestate exists to delete:

class Connection:
    def __init__(self):
        self.connected = False

    def connect(self):
        self.connected = True

    def send(self, msg):
        if not self.connected:            # a runtime guard you must remember to write
            raise RuntimeError("not connected")
        print(f"sending {msg}")

c = Connection()
c.send("hello")   # boom at runtime: RuntimeError, only discovered when this line runs

The bug here is found when the line executes, if you are lucky enough to hit it in testing. In the Rust version, the equivalent mistake is caught before the program ever runs, by the compiler, on every build. That is the entire difference between "we test for this" and "this cannot happen".

Go is in the same boat -- an interface method set is fixed, so you cannot swap the available methods as an object changes state; you reach for a runtime if and an error return, same as Python. The languages that can do something like typestate tend to be the ones with rich type systems: Haskell with phantom types and indexed monads, TypeScript to a degree with discriminated unions narrowing which fields are present. But Rust's version is unusually clean, because ownership does half the work -- consuming self on a transition is what stops you clinging to the old state, and no garbage-collected language gives you that guarantee for free. It is the same theme we keep circling back to: Rust hands you a way to say precisely what you mean, and then the compiler makes you keep your word ;-)

What did we actually learn?

Typestate programming encodes a value's current state into its type, so each state exposes only its legal operations and calling the wrong method for the state is a compile error, not a runtime one.
The recipe is a generic state parameter filled with tiny marker types (struct Open;), held with PhantomData so the state costs nothing at runtime, plus one impl block per state.
Transitions consume self and return a value in the new state, so ownership guarantees the old state can never be used again -- episode 3's ownership model doing state-machine safety.
Real data lives in ordinary fields carried forward through each transition; the markers are pure compile-time bookkeeping.
The pattern makes builders foolproof (a required field encoded as a state means send/build cannot be reached without it, and the internal unwrap can never panic) and models honest protocols (a connection, a file mode, a transaction).
It is worth the ceremony for stateful resources where a wrong call is costly, and over-engineering for a plain value type -- judge the guarantee against the machinery.

That closes the trait-and-type-system phase we have been building for many episodes now. We have spent a long stretch teaching the compiler to hold invariants for us -- coherence, sealed traits, const generics, GATs, and now whole state machines drawn in types. Next time we change gears and put all of it to work on something bigger and more hands-on, assembling these pieces into a real, reusable library rather than isolated snippets. Bring what you have learned; you are going to need most of it. One idea at a time ;-)

Exercises

Three exercises as always, from gentle to chewier. Full solutions open the next episode, so genuinely have a go first -- typing it yourself is where the pattern actually sticks, much more than reading mine.

Add a Locked state to the door so that open is callable only from an Unlocked closed door. Give it lock and unlock transitions, and confirm the compiler rejects opening a locked door.
Extend the Request builder with a required body, encoded as its own state, so that send becomes available only once both a url and a body have been provided (hint: you will need a state that records both facts, or two chained states).
Add a caution: bool field to the traffic light, carry it through all three states like the title in our Post example, and read it inside each next so the light can behave differently when caution mode is on.