DEV Community: Joshua

Enhancing Rust Enums in the State Pattern

Joshua — Sun, 16 Jun 2024 16:39:46 +0000

Recap

In my previous article, I discussed how Rust enums should be strongly considered when the solution benefits from a state machine. The strongest argument for this is the fact that the Rust compiler will inform you when a state variant isn't covered in a match expression.

This was stemmed from a state pattern example provided by the official Rust book. The sample scenario they used was an article that was required to go through the status of Draft, PendingReview, and Approved.

Using a similar strategy as I explained in my original article, I came up with the following code:

post.rs

use crate::state::{ArticleState, ArticleTransition};

#[derive(Default)]
pub struct Post {
  state: ArticleState,
  content: String,
}

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

  pub fn content(&self) -> &str {
    self.state.content(&self.content)
  }

  pub fn request_review(&mut self) {
    self.state.update_state(ArticleTransition::RequestReview);
  }

  pub fn approve(&mut self) {
    self.state.update_state(ArticleTransition::Approve);
  }
}

state.rs

enum State {
  Draft,
  PendingReview,
  Published,
}

pub enum ArticleTransition {
  RequestReview,
  Approve,
}

pub struct ArticleState {
  state: State,
}

impl Default for ArticleState {
  fn default() -> Self {
    Self {
       state: State::Draft,
    }
  }
}

use ArticleTransition as T;
use State as S;
impl ArticleState {
  pub fn update_state(&mut self, transition: ArticleTransition) {
     match (&self.state, transition) {
       // Handle RequestReview
       (S::Draft, T::RequestReview) => self.state = S::PendingReview,
       (_, T::RequestReview) => (),

       // Handle Approve
       (S::PendingReview, T::Approve) => self.state = S::Published,
       (_, T::Approve) => (),
     }
  }

  pub fn content<'a>(&self, article_post: &'a str) -> &'a str {
    match self.state {
       S::Published => article_post,
       _ => "",
    }
  }
}

New Requirements

While this code satisfies the test cases described by the Rust book, the following section shares some plausible modifiers to the requirements of our code. The new requirements are as follows:

Add a reject method that changes the post’s state from PendingReview back to Draft.
Require two calls to approve before the state can be changed to Published.
Allow users to add text content only when a post is in the Draft state. Hint: have the state object responsible for what might change about the content but not responsible for modifying the Post.

The first and third requirements are fairly straight forward in execution with our enums and match pattern. However the second case exposes yet another feature we can utilize in Rust.

Because an article in the PendingReview state require two approvals before moving to Published, we will need a reference the current count of approvals.

Enums Can Contain Fields

Some initial considerations may include a mutable field in our state struct. But this would be disconnected from our enum definition. We only really need an approval count when the article is in the state of PendingReview.

Another feature Rust enums provide is the ability to store fields internally inside a single enum variant. Let's try this with our PendingReview state.

enum State {
  Draft,
  PendingReview { approvals: u8 },
  Published,
}

Now we have a stateful field for our PendingReview state that holds the current count of approvals. This will require that any use of PendingReview must also acknowledge the approvals field.

We can now update our match arm for PendingReview to include the reference to approvals and run a simple condition of whether we should set the state to Published.

(S::PendingReview { approvals }, T::Approve) => {
  let current_approvals = approvals + 1;
  if current_approvals >= 2 {
    self.state = S::Published
  } else {
    self.state = S::PendingReview {
      approvals: current_approvals,
    }
  }
}

And that's it. By adding a field to a single enum variant, we now have additional context for this single expression without the need to store any reference in the root of our State struct.

This also makes our code easier to understand as an unfamiliar contributor can easily deduce that the approvals field only matters when our state is in PendingReview.

Full Code Example

post.rs

use crate::state::{ArticleState, ArticleTransition};

#[derive(Default)]
pub struct Post {
  state: ArticleState,
  content: String,
}

impl Post {
  pub fn add_text(&mut self, text: &str) {
    let is_mutable = self.state.is_text_mutable();
    match is_mutable {
      true => {
        self.content.push_str(text);
      }
      false => (),
    }
  }

  pub fn content(&self) -> &str {
    self.state.content(&self.content)
  }

  pub fn request_review(&mut self) {
    self.state.update_state(ArticleTransition::RequestReview);
  }

  pub fn approve(&mut self) {
    self.state.update_state(ArticleTransition::Approve);
  }

  pub fn reject(&mut self) {
    self.state.update_state(ArticleTransition::Reject)
  }
}

state.rs

enum State {
  Draft,
  PendingReview { approvals: u8 },
  Published,
}

pub enum ArticleTransition {
  RequestReview,
  Approve,
  Reject,
}

pub struct ArticleState {
  state: State,
}

impl Default for ArticleState {
  fn default() -> Self {
    Self {
       state: State::Draft,
    }
  }
}

use ArticleTransition as T;
use State as S;

impl ArticleState {
  pub fn update_state(&mut self, transition: ArticleTransition) {
    match (&self.state, transition) {
      // Handle RequestReview
      (S::Draft, T::RequestReview) => self.state = S::PendingReview { approvals: 0 },
      (_, T::RequestReview) => (),

      // Handle Approve
      (S::PendingReview { approvals }, T::Approve) => {
        let current_approvals = approvals + 1;
        if current_approvals >= 2 {
          self.state = S::Published
        } else {
          self.state = S::PendingReview {
            approvals: current_approvals,
          }
        }
      }
      (_, T::Approve) => (),
      (S::PendingReview { .. }, T::Reject) => self.state = S::Draft,
      (_, T::Reject) => (),
    }
  }

  pub fn content<'a>(&self, article_post: &'a str) -> &'a str {
    match self.state {
      S::Published => article_post,
      _ => "",
    }
  }

  pub fn is_text_mutable(&self) -> bool {
    matches!(self.state, S::Draft)
  }
}

Challenging the State Pattern with Rust Enums

Joshua — Thu, 13 Jun 2024 13:59:31 +0000

In the official Rust book, there's a section that attempts to provide an example of the State design pattern in order to showcase some of Rust's OOP muscles.

If you are not familiar with the state pattern, I suggest reading up on it before continuing.

As I read this example, I found its design to be odd. I started wondering why the example wasn't taking advantage of enums. Then like magic the book included this figure text:

You may have been wondering why we didn’t use an enum with the different possible post states as variants. That’s certainly a possible solution, try it and compare the end results to see which you prefer! One disadvantage of using an enum is every place that checks the value of the enum will need a match expression or similar to handle every possible variant. This could get more repetitive than this trait object solution.

Reference: The Rust Book

After reading this I still disagreed that structs were the better choice. So I decided to take the challenge to build my own state machine using Rust's enums.

Before we get started, let's first look at an example using structs.

A State Machine with Structs

The scenario being covered in the Rust book involved the different states of an article post. However, I came across a more interesting example from Refactoring Guru which involves a music player.

The music player will have four buttons:

Play
Stop
Prev Track
Next Track

Some of these buttons (Play and Stop) should behave differently depending on the current state of the music player. While others (Prev/Next Track) should behave the same regardless of state.

The State Struct

Refactoring Guru's example follows a similar strategy as the Rust Book, by creating a collection of State structs, each one handling their own behavior for each button of the music player. The state struct also has a mutable reference to the music player so it can apply the necessary side effect for each action.

*A lot of these code snippets will be edited for the sake of brevity, but I will include a link to the full code snippet provided by Refactoring Guru.

Ref: https://refactoring.guru/design-patterns/state/rust/example#example-0--state-rs

pub trait State {
  fn play(self: Box<Self>, player: &mut Player) -> Box<dyn State>;
  fn stop(self: Box<Self>, player: &mut Player) -> Box<dyn State>;
}

impl State for StoppedState {
  fn play(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Play/Pause" button for the "Stopped" state.
  }

  fn stop(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Stop" button for the "Stopped" state.
  }
}

impl State for PausedState {
  fn play(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Play/Pause" button for the "Paused" state.
  }

  fn stop(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Stop" button for the "Paused" state.
  }
}

impl State for PlayingState {
  fn play(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Play/Pause" button for the "Playing" state.
  }

  fn stop(self: Box<Self>, player: &mut Player) -> Box<dyn State> {
    // Apply logic for the "Stop" button for the "Playing" state.
  }
}

This code is pretty straightforward in what its trying to achieve. However it does feel a bit more cluttered given the duplicate function declarations, along with the added boilerplate for dynamic dispatch.

If you're not familiar with Rust's dynamic dispatch feature, (the bits about Box<dyn T>) it's not required for reading this article. But it is an important concept I would recommend to those wanting to learn Rust.

Lastly, this strategy puts more reliance on the developers' cognitive ability to remember all possible states and buttons. This may seem like a silly excuse regarding a music player, but it can be particularly challenging if you require a more complex state machine.

Actions as Strings

One thing I like about Refactoring Guru's example is their incorporation of an interactive UI which you can test your state machine on. However, when handling UI events, the example relies on static strings to define the trigger action.

Ref: https://doc.rust-lang.org/stable/book/ch17-02-trait-objects.html#trait-objects-perform-dynamic-dispatch

let mut app = cursive::default();

// ...

app.add_layer(
  Dialog::around(TextView::new("Press Play").with_name("Player Status"))
    .title("Music Player")
    .button("Play", |s| execute(s, "Play"))
    .button("Stop", |s| execute(s, "Stop"))
    .button("Prev", |s| execute(s, "Prev"))
    .button("Next", |s| execute(s, "Next")),
);

// ...

fn execute(s: &mut Cursive, button: &'static str) {
  let PlayerApplication {
    mut player,
    mut state,
  } = s.take_user_data().unwrap();

  let mut view = s.find_name::<TextView>("Player Status").unwrap();

  state = match button {
    "Play" => state.play(&mut player),
    "Stop" => state.stop(&mut player),
    "Prev" => state.prev(&mut player),
    "Next" => state.next(&mut player),
     _ => unreachable!(),
  };
}

By using static strings, we're again trusting the developers to remember all possible values when checking which button was pressed.

Bring Out the Enums Already

Now comes my attempt to refactor this using Rust's enums. The first thing I wanted to assure for this exercise is that no changes are made to the player module. Given the reason why the state pattern exists is to abstract away any stateful logic from the context object.

My first change to the state module was the definition of the State enum

enum State {
  Stopped,
  Playing,
  Paused,
}

I create a public PlayerState struct that contains one field for the current state. I also define the Default for this struct to where the state begins at Stopped.

pub struct PlayerState {
    state: State,
}

impl Default for PlayerState {
    fn default() -> Self {
        Self {
            state: State::Stopped,
        }
    }
}

The next area I believe that can benefit from enums is the execute() method inside main. Instead of passing in static strings to the execute() method, we can actually declare an enum of potential actions triggered by the UI. This enum will also be defined inside state.rs.

pub enum PlayerAction {
    Play,
    Stop,
    Prev,
    Next,
}

State Machine Logic

We have our state enum, we have our action enum, now we get to the fun part of implementing our state machine. Instead of declaring structs for each state, I simply declare a function for our single PlayerState struct. The arguments will be a mutable reference to itself, a mutable reference to the music player, and the transition event coming from the main module.

Since this function will make many references to our enums, I'll add a couple of lines to shorten their names to make our match expression block a little easier to read.

    pub fn update_state(&mut self, player: &mut Player, action: PlayerAction) {
        use PlayerAction as T;
        use State as S;

We now come to our match expression. What we want to do is match the current state of the music player and the transition that will occur due to the button press. To make the comparisons easier, let's put both the state and transition value in a tuple.

  match (&self.state, action) {
    (S::Playing, T::Play) => {
      player.pause();
      self.state = S::Paused;
    }
    (_, T::Play) => {
      player.play();
      self.state = S::Playing;
    }
    (S::Stopped, T::Stop) => (),
    (_, T::Stop) => {
      player.pause();
      player.rewind();
      self.state = S::Stopped;
    }
    (_, T::Next) => player.next_track(),
    (_, T::Prev) => player.prev_track(),
  }

I tried to order this list by the transition T, covering T::Play first, then T::Stop with T::Next and T::Prev to end it.

You may have noticed the absence of any UI side effects. This is because I like to keep a separation of concerns when it comes to side effects. So, I declared a separate method that handles any UI side effects based on the current state of the module. We also update our arguments as we don't need any mutable references except for the TextView struct of our UI.

pub fn update_view(&self, player: &Player, view: &mut TextView) {
  match self.state {
    State::Stopped => view.set_content("[Stopped] Press 'Play'"),
    State::Playing => view.set_content(format!(
      "[Playing] {} - {} sec",
      player.track().title,
      player.track().duration
    )),
    State::Paused => view.set_content(format!(
      "[Paused] {} - {} sec",
      player.track().title,
      player.track().duration
    )),
  }
}

What I Like About This Strategy

By using enums in our match expression we gain a significant advantage compared to the struct example, in that the Rust compiler now takes responsibility for ensuring every condition is met for all unique combinations of states and transitions.

Additionally, we get a small performance improvement since we no longer rely on dynamic dispatch to pass arguments that implement a State trait.

Revisiting the caution about enums from the Rust Book:

One disadvantage of using an enum is every place that checks the value of the enum will need a match expression or similar to handle every possible variant. This could get more repetitive than this trait object solution.

I would argue this code is far less repetitive than defining the same struct methods for every single state.

Plus, there will be many cases where a transition only requires a special behavior in one or two states and does nothing for all remaining states. When using enums, the underscore declaration becomes useful to group together any remaining states that should share the same behavior.

Lastly, if we were working on a more complex state machine, we could extract each transition arm into its own function (or module if needed) to handle the behavior of each state for that specific transition.

The main function

Now that we've finished our state.rs module, we can apply the necessary changes to our main function. First we update our state field in our App struct.

#[derive(Default)]
struct App {
    player: Player,
    state: PlayerState,
}

Next, we update the main() function so that our button presses pass the PlayerAction enum instead of a static string.

fn main() {
  let mut app = cursive::default();

  app.set_user_data(App::default());
  app.add_layer(
    Dialog::around(TextView::new("Press Play").with_name("Player Status"))
      .title("Music Player")
      .button("Play", |s| execute(s, PlayerAction::Play))
      .button("Stop", |s| execute(s, PlayerAction::Stop))
      .button("Prev", |s| execute(s, PlayerAction::Prev))
      .button("Next", |s| execute(s, PlayerAction::Next)),
    );

  app.add_global_callback(Key::Esc, |s| s.quit());

  app.run();
}

And finally, we clean up our execute function as we now simply pass the received enum directly into our new state method to apply the necessary state machine and UI changes.

fn execute(s: &mut Cursive, action: PlayerAction) {
  let App {
    mut player,
    mut state,
  } = s.take_user_data().unwrap();

  let mut view = s.find_name::<TextView>("Player Status").unwrap();

  state.update_state(&mut player, action);
  state.update_view(&player, &mut view);

  s.set_user_data(App { player, state });
}

Conclusion

After completing this exercise, I feel confident that the usage of Rust's enums really shouldn't be overlooked. While the state pattern can be used to showcase dynamic dispatch, I believe it may mislead readers to believe it is the only way to implement the state pattern.

I hope any readers of this article will also take considerations as to how they can use Rust's full toolset to build sensible solutions.

Full Code Example

player.rs (Unchanged but including for transparency)

pub struct Track {
  pub title: String,
  pub duration: u32,
  cursor: u8,
}

impl Track {
  pub fn new(title: &str, duration: u32) -> Self {
    Self {
      title: title.into(),
      duration,
      cursor: 0,
    }
  }
}

pub struct Player {
  playlist: Vec<Track>,
  current_track: usize,
  _volume: u8,
}

impl Default for Player {
  fn default() -> Self {
    Self {
      playlist: vec![
        Track::new("Track 1", 180),
        Track::new("Track 2", 250),
        Track::new("Track 3", 130),
        Track::new("Track 4", 220),
        Track::new("Track 5", 300),
      ],
      current_track: 0,
      _volume: 25,
    }
  }
}

impl Player {
  pub fn next_track(&mut self) {
    self.current_track = (self.current_track + 1) % self.playlist.len();
  }
  pub fn prev_track(&mut self) {
    self.current_track = (self.playlist.len() + self.current_track - 1) % self.playlist.len();
  }

  pub fn play(&mut self) {
    self.track_mut().cursor = 10; // Playback imitation.
  }

  pub fn pause(&mut self) {
    self.track_mut().cursor = 43; // Paused at some moment.
  }

  pub fn rewind(&mut self) {
    self.track_mut().cursor = 0;
  }

  pub fn track(&self) -> &Track {
    &self.playlist[self.current_track]
  }

  fn track_mut(&mut self) -> &mut Track {
    &mut self.playlist[self.current_track]
  }
}

state.rs

use cursive::views::TextView;

use crate::player::Player;

enum State {
  Stopped,
  Playing,
  Paused,
}

pub enum PlayerAction {
  Play,
  Stop,
  Prev,
  Next,
}

pub struct PlayerState {
  state: State,
}

impl Default for PlayerState {
  fn default() -> Self {
    Self {
      state: State::Stopped,
    }
  }
}

impl PlayerState {
  pub fn update_state(&mut self, player: &mut Player, action: PlayerAction) {
    use PlayerAction as T;
    use State as S;
    match (&self.state, action) {
      (S::Playing, T::Play) => {
         player.pause();
         self.state = S::Paused;
      }
      (_, T::Play) => {
         player.play();
         self.state = S::Playing;
      }
      (S::Stopped, T::Stop) => (),
      (_, T::Stop) => {
         player.pause();
         player.rewind();
         self.state = S::Stopped;
      }
      (_, T::Next) => player.next_track(),
      (_, T::Prev) => player.prev_track(),
    }
  }

  pub fn update_view(&self, player: &Player, view: &mut TextView) {
    match self.state {
      State::Stopped => view.set_content("[Stopped] Press 'Play'"),
      State::Playing => view.set_content(format!(
        "[Playing] {} - {} sec",
        player.track().title,
        player.track().duration
      )),
      State::Paused => view.set_content(format!(
        "[Paused] {} - {} sec",
        player.track().title,
        player.track().duration
      )),
    }
  }
}

main.rs

mod player;
mod state;

use crate::player::Player;
use cursive::{
  event::Key,
  view::Nameable,
  views::{Dialog, TextView},
  Cursive,
};
use state::{PlayerAction, PlayerState};

#[derive(Default)]
struct App {
  player: Player,
  state: PlayerState,
}

fn main() {
  let mut app = cursive::default();

  app.set_user_data(App::default());
  app.add_layer(
    Dialog::around(TextView::new("Press Play").with_name("Player Status"))
      .title("Music Player")
      .button("Play", |s| execute(s, PlayerAction::Play))
      .button("Stop", |s| execute(s, PlayerAction::Stop))
      .button("Prev", |s| execute(s, PlayerAction::Prev))
      .button("Next", |s| execute(s, PlayerAction::Next)),
    );

  app.add_global_callback(Key::Esc, |s| s.quit());

  app.run();
}

fn execute(s: &mut Cursive, action: PlayerAction) {
  let App {
    mut player,
    mut state,
  } = s.take_user_data().unwrap();

  let mut view = s.find_name::<TextView>("Player Status").unwrap();

  state.update_state(&mut player, action);
  state.update_view(&player, &mut view);

  s.set_user_data(App { player, state });
}