State Machine Transitions with Ownership Transfer

Hello Rustaceans!

I'm working with state machines in Rust where states are represented as enum variants and the transitions are written as match statements, and I often encounter situations where I need to transition between these states based on complex conditions. The challenge arises with non-cloneable types in these states, necessitating ownership transfer during state transitions. This rules out the use of mutable references, leading to some verbose patterns, especially in UI code with closures (e.g., using egui).

A simple approach could look something like this

struct Uncloneable;

enum State {
    A(Uncloneable),
    B(Uncloneable),
}

impl State {
    fn transition(mut self) -> Self {
        match self {
            State::A(u) => {
                if some_condition() {
                    if some_other_condition() {
                        State::B(u)
                    } else {
                        State::A(u)
                    }
                } else {
                    State::A(u)
                }
            }
            State::B(u) => State::A(u),
        }
    }
}

As you can see the expression value of the match statement provides us with the "new" state.
While this works, it becomes cumbersome when dealing with branching logic, as every execution path in the match arms needs to return a state.

I found that using labeled break statements can clean up the code significantly:

 fn transition(mut self) -> Self {
      match self {
          State::A(u) => 'label: {
              if some_condition() {
                  if some_other_condition() {
                      break 'label State::B(u);
                  }
              }
              State::A(u)
          }
          State::B(u) => State::A(u),
      }
  }

However, this solution falls short when closures are involved inside the match arms.
Ideally, I'd like the semantics of a mutable reference but with ownership, allowing for something like this:

 fn transition(mut self) -> Self {
      let this = &mut self;
      match this {
          State::A(u) => {
              if some_condition() {
                  if some_other_condition() {
                      *this = State::B(*u);
                  }
              }
          }
          State::B(u) => *this = State::A(*u),
      }
      *this
  }

Unfortunately, this is not feasible due to Rust's borrow checker and partial move issues.

As a workaround, I considered using an Option to hold a potential new state:

fn transition(self) -> Self {
      let mut new_state = None;
      match self {
          State::A(u) => {
              if some_condition() {
                  if some_other_condition() {
                      new_state = Some(State::B(u));
                  }
              }
          }
          State::B(u) => new_state = Some(State::A(u)),
      }
      if let Some(new_state) = new_state {
          new_state
      } else {
          self
      }
  }

This approach seems sound to me, but the borrow checker cannot ascertain this at compile time.

I'm seeking suggestions on achieving similar semantics to &mut but with ownership or any alternative patterns that might simplify this code. Any insights or advice would be greatly appreciated!

Thank you in advance!

P.S. If you need some more context on my use case, take a look at this:

I'm not satisfied with the current state of this code...

Thanks for posting code that shows the various problems.

Can you say a little more about why that is? Is it because you want to avoid another condition that checks the return value of the closure?

As far as I can tell, using return or break is the only way to do it.

Here's one other idea. Your conditional logic may be too complex for this, but perhaps it could work if you can write functions for the conditions. Sometimes writing small functions as closures works well.

fn transition(self) -> Self {
    match self {
        State::A(u) if some_condition() && some_other_condition() => {
            State::B(u)
        }
        State::B(u) => State::A(u),
        other => other,
    }
}

Thanks for responding.

This issue is about writing these state transitions neatly.
The challenge I'm facing is primarily due to the depth of nested closures and multiple conditional checks within them, which are inherent in my use case involving egui (e.g. Ui in egui - Rust). The clutter arises from having to propagate the return/break values through these nested layers. Here's an illustration of how complex this can become (this is not the particular code I have, but only an illustration):

impl State {
    fn transition(mut self) -> Self {
        match self {
            State::A(u) => 'label: {
                if some_condition0() {
                    if let Some(new_state) = (|| {
                        if some_condition1() {
                            return Some(State::C(u));
                        }

                        if let Some(new_state) = (|| {
                            if some_condition2() {
                                Some(State::D(u))
                            } else {
                                None
                            }
                        })() {
                            return new_state;
                        };
                        None
                    })() {
                        break 'label new_state;
                    }

                    if some_condition3() {
                        break 'label State::B(u);
                    }
                }
                State::A(u)
            }
            State::B(u) => State::A(u),
            State::C(u) => State::D(u),
            State::D(u) => State::B(u),
        }
    }
}

In this example, you can see the depth of nesting and the propagation of state transitions through multiple layers of closures. It becomes quite convoluted.

Contrastingly, the last approach I suggested simplifies this significantly (the code still looks cluttered due to the instant closure call, but this isn't present in the actual egui code):

fn transition(self) -> Self {
        let mut new_state = None;
        match self {
            State::A(u) => {
                if some_condition0() {
                    (|| {
                        if some_condition1() {
                            new_state = Some(State::C(u));
                        }

                        (|| {
                            if some_condition2() {
                                new_state = Some(State::D(u));
                            }
                        })();
                    })();

                    if some_condition3() {
                        new_state = Some(State::B(u));
                    }
                }
            }
            State::B(u) => new_state = Some(State::A(u)),
            State::C(u) => new_state = Some(State::D(u)),
            State::D(u) => new_state = Some(State::B(u)),
        }
        if let Some(new_state) = new_state {
            new_state
        } else {
            self
        }
    }

This approach is more streamlined and easier to manage, yet it’s not compatible with Rust's borrow checker constraints.

I appreciate your suggestion from the second comment. However, due to the complexity and depth of the conditions in my case, it doesn't quite address the core of my challenge.

I'm still looking for ideas or alternative patterns that could help simplify this kind of nested, conditional state transition logic, especially in the context of UI handling with egui. Any further insights or suggestions would be greatly appreciated.

I don't think using assignments to conditionally create the new state is possible, without using some sort of unsafe trick.

The problem is that matching a pattern will move the u in State::A(u) => out of self. Once that happens, self is no longer usable. So you can't use self as the default value in the fallthrough case. And of course you need the pattern match in order to use u in your conditionals.

With a ref pattern like State::A(ref u) => the move doesn't happen, but then you can't move u out of self when you do want to, in State::C(u). This is because you can't move a value out from behind a reference.

So as far as I know, every conditional branch needs to construct the new state as in the first approach you listed. I agree that this is not at all ideal.

1 Like

@LU15W1R7H Here is one approach that does what you asked, although it is somewhat prone to causing runtime errors. I used two Options for the new state and the (uncloneable) data of the old state, one of which is always None and the other Some. If you accidentally update the state twice, or try to reference the data after you've updated the state, you'll get a runtime error.

struct StateUpdater<T> {
    data: Option<T>,
    state: Option<State>,
}

impl<T> StateUpdater<T> {
    fn new(data: T) -> Self {
        StateUpdater { data: Some(data), state: None }
    }
    fn data(&self) -> &T {
        self.data.as_ref().expect("Cannot get data after updating state")
    }
    fn update<F: Fn(T) -> State>(&mut self, f: F) {
        self.state = Some(f(self.data.take().expect("Cannot update more than once")));
    }
    fn finish<F: Fn(T) -> State>(self, default: F) -> State {
        self.state.unwrap_or_else(|| default(self.data.unwrap()))
    }
}

impl State {
    fn transition(self) -> Self {
        match self {
            State::A(u) => {
                let mut updater = StateUpdater::new(u);
                if some_condition0(updater.data()) {
                    (|| {
                        if some_condition1(updater.data()) {
                            updater.update(|u| State::C(u));
                        }

                        (|| {
                            if some_condition2(updater.data()) {
                                updater.update(|u| State::D(u));
                            }
                        })();
                    })();

                    if some_condition3(updater.data()) {
                        updater.update(|u| State::B(u));
                    }
                }
                updater.finish(|u| State::A(u))
            }
            State::B(u) => State::A(u),
            State::C(u) => State::D(u),
            State::D(u) => State::B(u),
        }
    }
}

I realize this may still not fit your actual use case since that is more complex than the example. If so, hopefully it can be adapted.

EDIT: Also this still requires recreation of the original state, but at most once per enum variant.

And note that I did not test it, just got it to compile.

Here's a different version of StateUpdater that replaces the two Options with an enum. However, it depends on the replace_with crate, which has caveats about abort when a panic occurs -- because of that I prefer the version above with two Options.

enum StateUpdater<T> {
    Data(T),
    State(State),
}

impl<T> StateUpdater<T> {
    fn new(data: T) -> Self {
        StateUpdater::Data(data)
    }
    fn data(&self) -> &T {
        match self {
            Self::Data(data) => &data,
            Self::State(state) => {
                panic!("Cannot get data after updating state")
            }
        }
    }
    fn update<F: Fn(T) -> State>(&mut self, f: F) {
        replace_with::replace_with_or_abort(self, |self_| match self_ {
            Self::Data(data) => Self::State(f(data)),
            Self::State(state) => {
                panic!("Cannot update more than once")
            }
        });
    }
    fn finish<F: Fn(T) -> State>(self, default: F) -> State {
        match self {
            Self::Data(data) => default(data),
            Self::State(state) => state,
        }
    }
}

EDIT: Does your actual code always use the same data type in every enum variant, as in the example code? I assume the answer is no. But if the answer is yes, you can avoid the runtime errors in the StateUpdater and just update a variable with the enum constructor. The move of the data would be postponed until the end, when the enum constructor is applied.

fn transition(self) -> Self {
    match self {
        State::A(u) => {
            let mut new_state: Option<fn(Uncloneable) -> State> = None;
            if some_condition0(&u) {
                (|| {
                    if some_condition1(&u) {
                        new_state = Some(State::C);
                    }

                    (|| {
                        if some_condition2(&u) {
                            new_state = Some(State::D);
                        }
                    })();
                })();

                if some_condition3(&u) {
                    new_state = Some(State::B);
                }
            }
            new_state.unwrap_or(State::A)(u)
        }
        State::B(u) => State::A(u),
        State::C(u) => State::D(u),
        State::D(u) => State::B(u),
    }
}

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.