Challenging the State Pattern with Rust Enums
Tags: rust, designpatterns, learning
Published: 2024-06-13 13:59:31

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](https://refactoring.guru/design-patterns/state) 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](https://doc.rust-lang.org/stable/book/ch17-03-oo-design-patterns.html#why-not-an-enum)_

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](https://refactoring.guru/design-patterns/state/rust/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
```rust
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

```rust
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

```rust
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`.

```rust
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`.

```rust
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.

```rust
    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.

```rust
  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.

```rust
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.

```rust
#[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.

```rust
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.

```rust
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)
```rust
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_
```rust
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_
```rust
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 });
}
```