Let’s talk about building state machines in Rust using lifetimes. It’s a cool way to make your code safer and faster.

Rust’s lifetime system is pretty unique. It lets us bake state transitions right into our types. This means we can catch a lot of errors at compile time, not when our program is running.

I’ve found that using lifetimes for state machines is super helpful when I’m working on things like network protocols or game logic. It’s all about making sure we can’t accidentally do something that doesn’t make sense for the current state.

Here’s a simple example to get us started:

struct Door<'a> {
    state: &'a str,
}

impl<'a> Door<'a> {
    fn new() -> Door<'static> {
        Door { state: "closed" }
    }

    fn open(self) -> Door<'static> {
        Door { state: "open" }
    }

    fn close(self) -> Door<'static> {
        Door { state: "closed" }
    }
}

In this code, we’re using lifetimes to represent the state of a door. The open and close methods consume the door and return a new one with a different state.

But we can take this further. Let’s make it impossible to open an already open door:

struct Closed;
struct Open;

struct Door<State> {
    _state: State,
}

impl Door<Closed> {
    fn new() -> Self {
        Door { _state: Closed }
    }

    fn open(self) -> Door<Open> {
        Door { _state: Open }
    }
}

impl Door<Open> {
    fn close(self) -> Door<Closed> {
        Door { _state: Closed }
    }
}

Now, if we try to open an open door, the compiler will stop us. It’s pretty neat how we can use the type system to enforce these rules.

But what if we want to do something more complex? Let’s say we’re modeling a vending machine. It might have states like “idle”, “item selected”, “payment received”, and “dispensing”. We can use associated types to make this work:

trait State {
    type Next;
    fn next(self) -> Self::Next;
}

struct Idle;
struct ItemSelected;
struct PaymentReceived;
struct Dispensing;

impl State for Idle {
    type Next = ItemSelected;
    fn next(self) -> Self::Next {
        ItemSelected
    }
}

impl State for ItemSelected {
    type Next = PaymentReceived;
    fn next(self) -> Self::Next {
        PaymentReceived
    }
}

impl State for PaymentReceived {
    type Next = Dispensing;
    fn next(self) -> Self::Next {
        Dispensing
    }
}

impl State for Dispensing {
    type Next = Idle;
    fn next(self) -> Self::Next {
        Idle
    }
}

struct VendingMachine<S: State> {
    state: S,
}

impl<S: State> VendingMachine<S> {
    fn next(self) -> VendingMachine<S::Next> {
        VendingMachine { state: self.state.next() }
    }
}

This setup ensures that our vending machine always moves through the states in the correct order. We can’t accidentally skip a step or go backwards.

One thing I love about this approach is that it’s zero-cost. The compiler can optimize away all of this type-level machinery, leaving us with just the bare metal operations we need.

But what if we want to allow multiple possible transitions from a state? We can use enums for this:

enum DoorState {
    Open,
    Closed,
    Locked,
}

struct Door {
    state: DoorState,
}

impl Door {
    fn new() -> Self {
        Door { state: DoorState::Closed }
    }

    fn open(&mut self) {
        match self.state {
            DoorState::Closed => self.state = DoorState::Open,
            DoorState::Open => println!("Door is already open"),
            DoorState::Locked => println!("Can't open a locked door"),
        }
    }

    fn close(&mut self) {
        match self.state {
            DoorState::Open => self.state = DoorState::Closed,
            DoorState::Closed => println!("Door is already closed"),
            DoorState::Locked => println!("Door is locked"),
        }
    }

    fn lock(&mut self) {
        match self.state {
            DoorState::Closed => self.state = DoorState::Locked,
            DoorState::Open => println!("Can't lock an open door"),
            DoorState::Locked => println!("Door is already locked"),
        }
    }
}

This approach gives us more flexibility, but we lose some of the compile-time guarantees. It’s a trade-off we often have to make in real-world programming.

Now, let’s look at how we can use generics to make our state machines more flexible. Imagine we’re modeling a network connection:

trait ConnectionState {}

struct Disconnected;
struct Connected;
struct Authenticated;

impl ConnectionState for Disconnected {}
impl ConnectionState for Connected {}
impl ConnectionState for Authenticated {}

struct Connection<S: ConnectionState> {
    _state: S,
}

impl Connection<Disconnected> {
    fn new() -> Self {
        Connection { _state: Disconnected }
    }

    fn connect(self) -> Connection<Connected> {
        println!("Connecting...");
        Connection { _state: Connected }
    }
}

impl Connection<Connected> {
    fn authenticate(self) -> Connection<Authenticated> {
        println!("Authenticating...");
        Connection { _state: Authenticated }
    }

    fn disconnect(self) -> Connection<Disconnected> {
        println!("Disconnecting...");
        Connection { _state: Disconnected }
    }
}

impl Connection<Authenticated> {
    fn send_data(&self, data: &str) {
        println!("Sending data: {}", data);
    }

    fn disconnect(self) -> Connection<Disconnected> {
        println!("Disconnecting...");
        Connection { _state: Disconnected }
    }
}

This setup allows us to model complex state transitions while still maintaining type safety. We can’t accidentally send data on a connection that isn’t authenticated, for example.

One of the cool things about using Rust’s type system for state machines is that it serves as documentation. Anyone reading your code can immediately see what states are possible and how they transition.

But there’s a downside to this approach: it can lead to code duplication. If we have methods that should be available in multiple states, we might end up copying them. There are ways around this, like using trait objects, but they come with their own trade-offs.

Let’s look at a more advanced example. We’ll model a simple ATM:

struct Card;
struct Pin;
struct Amount;

struct IdleState;
struct CardInsertedState;
struct PinEnteredState;
struct TransactionState;

struct ATM<State> {
    _state: State,
}

impl ATM<IdleState> {
    fn new() -> Self {
        ATM { _state: IdleState }
    }

    fn insert_card(self, _: Card) -> ATM<CardInsertedState> {
        println!("Card inserted");
        ATM { _state: CardInsertedState }
    }
}

impl ATM<CardInsertedState> {
    fn enter_pin(self, _: Pin) -> ATM<PinEnteredState> {
        println!("PIN entered");
        ATM { _state: PinEnteredState }
    }

    fn eject_card(self) -> ATM<IdleState> {
        println!("Card ejected");
        ATM { _state: IdleState }
    }
}

impl ATM<PinEnteredState> {
    fn select_transaction(self) -> ATM<TransactionState> {
        println!("Transaction selected");
        ATM { _state: TransactionState }
    }

    fn cancel(self) -> ATM<IdleState> {
        println!("Transaction cancelled");
        ATM { _state: IdleState }
    }
}

impl ATM<TransactionState> {
    fn withdraw(self, _: Amount) -> ATM<IdleState> {
        println!("Cash withdrawn");
        ATM { _state: IdleState }
    }

    fn check_balance(self) -> ATM<TransactionState> {
        println!("Balance displayed");
        self
    }

    fn finish(self) -> ATM<IdleState> {
        println!("Transaction finished");
        ATM { _state: IdleState }
    }
}

This ATM model ensures that operations happen in the correct order. You can’t withdraw money without inserting a card and entering a PIN first.

One thing to keep in mind when using this technique is that it can make your types more complex. This can lead to longer compile times and more complex error messages. It’s a trade-off between safety and simplicity.

There’s also the question of how to handle errors. In a real ATM, the PIN might be incorrect, or there might not be enough money for a withdrawal. We can model this using Result types:

impl ATM<CardInsertedState> {
    fn enter_pin(self, pin: Pin) -> Result<ATM<PinEnteredState>, ATM<CardInsertedState>> {
        if pin_is_correct(&pin) {
            println!("PIN entered correctly");
            Ok(ATM { _state: PinEnteredState })
        } else {
            println!("Incorrect PIN");
            Err(self)
        }
    }
}

impl ATM<TransactionState> {
    fn withdraw(self, amount: Amount) -> Result<ATM<IdleState>, ATM<TransactionState>> {
        if sufficient_funds(&amount) {
            println!("Cash withdrawn");
            Ok(ATM { _state: IdleState })
        } else {
            println!("Insufficient funds");
            Err(self)
        }
    }
}

This approach allows us to handle errors while still maintaining our state machine structure.

In conclusion, using Rust’s lifetime system to create zero-cost state machines is a powerful technique. It allows us to catch errors at compile time, serve as self-documenting code, and potentially improve performance. However, it’s not without its challenges. It can lead to more complex types and potential code duplication. As with many things in programming, it’s about choosing the right tool for the job. For systems where correctness is crucial and performance is a concern, this approach can be a game-changer. But for simpler applications, a more traditional approach might be more appropriate. The key is to understand the trade-offs and make an informed decision based on your specific needs.