Challenging the State Pattern with Rust Enums

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 });
}