State Design Pattern in Rust

The State design pattern is a behavioral pattern that allows an object to alter its behavior when its internal state changes. The object appears to change its class. This pattern is useful for managing objects that have a complex lifecycle with distinct stages, where the behavior in each stage is significantly different.

A classic example is a blog post, which might transition through states like Draft, PendingReview, and Published. The actions you can perform on the post (like editing content or approving it) depend entirely on its current state.

In Rust, there are two common ways to implement the State pattern:

  1. A traditional object-oriented approach using trait objects.
  2. A more idiomatic Rust approach using enums and match expressions.

This article will explore both implementations.

1. The OOP-style State Pattern with Trait Objects

This approach closely follows the classic Gang of Four pattern. We define a State trait that all concrete state types will implement. The main object, our Post, will hold a Box<dyn State> and delegate behavior to it.

The core idea is that the state objects themselves control the state transitions.

// The context object that holds the current state.
pub struct Post {
    // We use Option to temporarily take ownership of the state for transitions.
    state: Option<Box<dyn State>>,
    content: String,
}

impl Post {
    pub fn new() -> Post {
        Post {
            // The post starts in the Draft state.
            state: Some(Box::new(Draft {})),
            content: String::new(),
        }
    }

    pub fn add_text(&mut self, text: &str) {
        self.content.push_str(text);
    }

    pub fn content(&self) -> &str {
        // The state determines how to get the content.
        // We use as_ref to get a reference to the Box.
        self.state.as_ref().unwrap().content(self)
    }

    pub fn request_review(&mut self) {
        // take() moves the state out of the Option, leaving None.
        if let Some(s) = self.state.take() {
            // The old state determines the new state.
            self.state = Some(s.request_review())
        }
    }

    pub fn approve(&mut self) {
        if let Some(s) = self.state.take() {
            self.state = Some(s.approve())
        }
    }
}

// The State trait defines the behavior for all states.
trait State {
    // These methods take ownership of the old state and return the new state.
    fn request_review(self: Box<Self>) -> Box<dyn State>;
    fn approve(self: Box<Self>) -> Box<dyn State>;
    // Default implementation for content returns an empty string.
    fn content<'a>(&self, post: &'a Post) -> &'a str {
        ""
    }
}

// Concrete state: Draft
struct Draft {}
impl State for Draft {
    fn request_review(self: Box<Self>) -> Box<dyn State> {
        Box::new(PendingReview {}) // Transition to PendingReview
    }
    fn approve(self: Box<Self>) -> Box<dyn State> {
        self // No change
    }
}

// Concrete state: PendingReview
struct PendingReview {}
impl State for PendingReview {
    fn request_review(self: Box<Self>) -> Box<dyn State> {
        self // No change
    }
    fn approve(self: Box<Self>) -> Box<dyn State> {
        Box::new(Published {}) // Transition to Published
    }
}

// Concrete state: Published
struct Published {}
impl State for Published {
    fn request_review(self: Box<Self>) -> Box<dyn State> {
        self // No change
    }
    fn approve(self: Box<Self>) -> Box<Self> {
        self // No change
    }
    fn content<'a>(&self, post: &'a Post) -> &'a str {
        &post.content
    }
}

fn main() {
    let mut post = Post::new();

    post.add_text("I ate a salad for lunch today");
    assert_eq!("", post.content());

    post.request_review();
    assert_eq!("", post.content());

    post.approve();
    assert_eq!("I ate a salad for lunch today", post.content());
}

Pros:

  • Encapsulates all behavior related to a particular state within its own struct.
  • Follows the Open/Closed Principle: you can add new states without changing the Post struct or other state structs.

Cons:

  • More verbose and complex.
  • Incurs a small runtime cost due to dynamic dispatch and heap allocation (Box).
  • State transitions can be tricky to implement (e.g., the take() pattern).

2. The Idiomatic Rust State Pattern with Enums

A more common and often simpler approach in Rust is to use an enum to represent the possible states. The logic for each state is handled within a match expression inside the methods of the main Post struct.

pub struct Post {
    state: State,
    content: String,
}

// The state is now a simple enum.
enum State {
    Draft,
    PendingReview,
    Published,
}

impl Post {
    pub fn new() -> Post {
        Post {
            state: State::Draft,
            content: String::new(),
        }
    }

    pub fn add_text(&mut self, text: &str) {
        // We can only add text in the Draft state.
        if let State::Draft = self.state {
            self.content.push_str(text);
        }
    }

    pub fn content(&self) -> &str {
        match self.state {
            // Only Published posts have content.
            State::Published => &self.content,
            _ => "",
        }
    }

    pub fn request_review(&mut self) {
        match self.state {
            State::Draft => self.state = State::PendingReview,
            _ => {} // No change in other states
        }
    }

    pub fn approve(&mut self) {
        match self.state {
            State::PendingReview => self.state = State::Published,
            _ => {} // No change in other states
        }
    }
}

fn main() {
    let mut post = Post::new();

    post.add_text("I ate a salad for lunch today");
    assert_eq!("", post.content());

    post.request_review();
    assert_eq!("", post.content());

    post.approve();
    assert_eq!("I ate a salad for lunch today", post.content());
}

Pros:

  • Simpler, more concise, and easier to understand.
  • No heap allocation or dynamic dispatch, making it more performant.
  • The compiler can check for exhaustive match arms, ensuring all states are handled.

Cons:

  • Does not follow the Open/Closed Principle as strictly. Adding a new state requires modifying the match expressions in all methods on Post.
  • Can lead to large match statements if the logic is complex.

Conclusion

Both patterns are valid ways to implement the State pattern in Rust. The OOP approach with trait objects is flexible and extensible, making it a good choice for complex systems where new states might be added frequently. However, for most use cases, the idiomatic approach using enums is simpler, safer, and more performant, making it the preferred choice for many Rust developers.