DEV Community

Cover image for Enhancing Rust Enums in the State Pattern
digclo
digclo

Posted on

Enhancing Rust Enums in the State Pattern

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);
  }
}
Enter fullscreen mode Exit fullscreen mode

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,
       _ => "",
    }
  }
}
Enter fullscreen mode Exit fullscreen mode

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,
}
Enter fullscreen mode Exit fullscreen mode

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,
    }
  }
}
Enter fullscreen mode Exit fullscreen mode

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)
  }
}
Enter fullscreen mode Exit fullscreen mode

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)
  }
}
Enter fullscreen mode Exit fullscreen mode

Top comments (0)