Enhancing Rust Enums in the State Pattern
Tags: rust, designpatterns, learning
Published: 2024-06-16 16:39:46

# Recap

In my previous [article](https://dev.to/digclo/state-pattern-with-rust-enums-61g), 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](https://doc.rust-lang.org/stable/book/ch17-03-oo-design-patterns.html#implementing-an-object-oriented-design-pattern). 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_

```rust
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_
```rust
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](https://doc.rust-lang.org/stable/book/ch17-03-oo-design-patterns.html#trade-offs-of-the-state-pattern) 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.

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

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