Easy State Management: Unleashing the Power of the State Pattern in Rust

Introduction

The state pattern is a behavourial state pattern, which allows an object to change its behaviour when its internal state changes. Since this sounds quite cryptic, let’s have a look at the diagram:

A short breakdown:

1. Context has two component in this simplified version: a field called state, which is an interface type, and an operation() method.
2. When the operation() method is called, the Context object calls the operation() on its state field. Since the state field can hold any implementation of the State interface, the behaviour of Context can change.

As you can see this is not the most complex of Design patterns.

Implementation in Rust

In an empty directory open your commandline or terminal and type:

cargo new state_pattern
cd state_pattern

Implementation with smart pointers

There are several ways to implement this pattern in Rust. We will start by implementing this using smart pointers. First we define the State trait:

trait State {
    fn handle(&self);
}

This quite straightforward, just one method which does some handling.

Now implement StateA:

struct StateA;

impl State for StateA {
    fn handle(&self) {
        println!("Handling state A");
    }
}

Some explanation:

1. StateA is an empty struct. This is just for simplicity’s sake.
2. We explicitly define the handle() method for StateA. This does nothing more than print out a message.

The definition of StateB is similar:

struct StateB;

impl State for StateB {
    fn handle(&self) {
        println!("Handling state B");
    }
}

Now we define the Context:

struct Context {
    state:Box<dyn State>
}

impl Context {
    fn new(state: Box<dyn State>)->Self {
        Self {
            state: state
        }
    }

    fn set_state(&mut self,state:Box<dyn State>) {
        self.state=state;
    }
    
}

impl State for Context {
    fn handle(&self) {
        self.state.handle();
    }
}

A line by line explanation:

1. The Context is a struct which holds one field in this example, state, which is a Box smartpointer to a dyn State. This is needed because Rust needs to know the size of an object before it can be placed on the stack. We need a dyn here, because State is a trait object.
2. We define a constructor, new, which gets a Box-ed dyn State object and assigns it to the state.
3. We also define a set_state() method, so we can change the state.
4. We also implement the State trait for Context, this allows us to pass down the call to handle() to the state.

Now, time for a test run:

fn main() {
    let state=StateA;
    let mut context=Context::new(Box::new(state));

    context.handle();

    let second_state=StateB;
    context.set_state(Box::new(second_state));
    context.handle();
}

We will go through this:

1. We construct an object of type StateA. Remember StateA implements the State trait.
2. Then we construct the context, and pass it the Box-ed instance of the state struct.
3. We call handle().
4. Next we construct a new state called second_state of type StateB
5. We set the state of the context by passing a Box-ed instance of the second_state to the set_state method.
6. Again we call handle()

I found that the Box-ed parameters clutter up the code, so we will do something about it:

Use explicit implementations

We can use impl trait types to get rid of the Box-ed parameters, like this. First we rewrite the Context:

struct Context<'a> {
    state:&'a dyn State
}

impl<'a> Context<'a> {
    fn new(state: &'a dyn State)->Self {
        Self {
            state: state
        }
    }

    fn set_state(&mut self,state:&'a impl State) {
        self.state=state;
    }
    
}

impl<'a> State for Context<'a> {
    fn handle(&self) {
        self.state.handle();
    }
}

1. Because we are using unboxed types which can go out of scope before we are done using it, we need a lifetime parameter which I simple named ‘a
2. In the Context implementation the constructor is mostly the same, but with the addition of the lifetime parameter.
3. The set_state() method gets an &’a impl State parameter instead of a Box-ed pointer. This means this method takes any struct which implement the State trait.
4. The State implementation for the Context just gets the additional lifetime parameter.

The main method also has some changes:

fn main() {
    let state=&StateA;
    let mut context=Context::new(state);

    context.handle();

    let second_state=&StateB;
    context.set_state(second_state);
    context.handle();
}

Line by line:

1. We create a reference to a struct of type StateA
2. We create a Context object and pass it the state variable. Notice that context has the mut keyword, since the set_state() method can change it.
3. We call handle() on context
4. Next we create second_state of type StateB
5. We set the state of context using set_state().
6. We call handle() again on the context

As you can see this method has less clutter.

Conclusion

This pattern was quite straightforward to implement. Rust is quite flexible in this way, as you can tell from the two different implementations. Also the Rust compiler is very helpful, and it came up with many a useful suggestion while writing this code.