So I needed to manage states in my app once again. I decided it was time to finally address this issue properly and create a package that I could reuse whenever I wanted.
As always, it all starts with simple experiments. I analyze the problems and try to push the solution to its extreme in order to achieve the best possible syntax.
Requirements
- States that are independent and can lead to other states.
- The ability to know, in a state, when it is active or inactive.
- A state machine that encapsulates the current state.
Let's take a simple enough example: 2 states, on and off, that automatically switch from one to the other every n seconds. A first trivial approach allows us to imagine the state machine like this:
@MainActor
protocol StateMachineState: Sendable {
func enter(stateMachine: StateMachine) async
func exit(stateMachine: StateMachine) async
}
extension StateMachineState {
func enter(stateMachine: StateMachine) async { }
func exit(stateMachine: StateMachine) async { }
}
@MainActor
final class StateMachine {
var currentState: StateMachineState
init(initialState: StateMachineState) {
currentState = initialState
}
func start() async {
await currentState.enter(stateMachine: self)
}
func transition(to newState: StateMachineState) async {
await currentState.exit()
currentState = newState
await currentState.enter(stateMachine: self)
}
}From there, we can build our states:
struct OnState: StateMachineState {
let delay: TimeInterval
func enter(stateMachine: StateMachine) async {
print("ON")
Task {
try? await Task.sleep(for: .seconds(delay))
await stateMachine.transition(to: OffState(delay: delay))
}
}
}
struct OffState: StateMachineState {
let delay: TimeInterval
func enter(stateMachine: StateMachine) async {
print("OFF")
Task {
try? await Task.sleep(for: .seconds(delay))
await stateMachine.transition(to: OnState(delay: delay))
}
}
}And finally, use our state machine like this:
let stateMachine = StateMachine(initialState: OffState(delay: 1))
await stateMachine.start()This implementation meets the requirements. The states are independent, the state machine encapsulates the current state, and the enter() and exit() functions allow states to be notified when they are active.
So what's wrong?
Although this approach is completely functional, it is actually prone to several possible human errors.
First, the enter() and exit() functions have the entire state machine as a parameter. It's theoretically possible for a developer to accidentally call functions other than transition(to:) and thus break the functionality.
struct OnState: StateMachineState {
func enter(stateMachine: StateMachine) async {
// This is technically useless, but it's possible.
await stateMachine.start()
}
}Next, the parameters of the states are always independent. This is a good thing in some cases. But in others, like here, we would like to declare only one delay. Once again, it's not out of the question that the developer might change the delay of the on state without modifying that of the off state.
Finally, if the developer keeps a reference to the machine in a state, they must remember to declare the machine as weak. And you know as well as I do that "they must remember" eventually translates to "they forgot."
A SwiftUI approach
The @Environment property wrapper can be our source of inspiration to improve our component. In the same way that the SwiftUI environment can propagate to child views, it would be interesting for an environment to propagate to our states.
Let's try to imagine the most optimal syntax and see what we can do to achieve it.
struct Context: StateMachineContext {
let delay: TimeInterval
}
struct OffState: StateMachineState {
@StateTransition private var transition
@StateContext(Context.self) private var context
func enter() async {
Task {
try? await Task.sleep(for: .seconds(context.delay))
await transition(to: OnState())
}
}
}
let stateMachine = StateMachine(initialState: OffState())
.context(Context(delay: 1))There are quite a few changes. Don't panic: let's break down this difficult problem into several small, easy-to-solve problems.
It’s all about injection
The first question that probably comes to your mind when seeing this: how can we retrieve data via a property wrapper without going through an initializer?
The trick is to wrap our state and inject the necessary data using Mirror. Mirrors in Swift allow us to dynamically explore the internal structure of an object, even if its properties are private. Thus, we can access the property wrapper and inject the necessary data.
struct StateMachineStateWrapper<Context: StateMachineContext>: Sendable {
let state: StateMachineState
let context: Context
@MainActor
func makeInjection() {
let mirror: Mirror = .init(reflecting: state)
for child in mirror.children {
if let propertyWrapper = child.value as? StateContext<Context> {
propertyWrapper.inject(context: context)
}
}
}
}Here's what our @StateContext looks like:
@MainActor
@propertyWrapper
public final class StateContext<Context: StateMachineContext> {
public var wrappedValue: Context {
guard let context = context else {
print("Accessing @StateContext within a state that is not managed by a StateMachine.")
return Context.defaultValue
}
return context
}
private var context: Context?
public nonisolated init(_ context: Context.Type) { }
func inject(context: Context) {
self.context = context
}
}The message displayed if the context has not yet been injected should remind you of something. If you've ever tried to modify a property annotated with @State in SwiftUI in a view that isn't installed, you must have received this warning:
Accessing State's value outside of being installed on a View. This will result in a constant Binding of the initial value and will not update.
It's exactly the same concept here. If the context hasn't been injected yet, it means our state isn't encapsulated in a StateMachineStateWrapper and therefore isn't managed by the state machine. This is similar to how SwiftUI is able to detect within a View if it is installed in the hierarchy. Clever!
Regarding the @StateTransition property wrapper, we can use the same injection mechanism. This time we inject a StateTransitionHandler that will handle the link between the state and the state machine.
@MainActor
@propertyWrapper
public final class StateTransition {
public var wrappedValue: StateTransitionHandler {
guard let transitionHandler = transitionHandler else {
assertionFailure("Accessing @StateTransition within a state that is not managed by a StateMachine.")
return .init()
}
return transitionHandler
}
private var transitionHandler: StateTransitionHandler?
public nonisolated init() { }
func inject(transitionHandler: StateTransitionHandler) {
self.transitionHandler = transitionHandler
}
}@MainActor
public final class StateTransitionHandler {
typealias Handler = @Sendable (_ newState: StateMachineState) async -> Void
private var transitionHandler: Handler?
public func callAsFunction(to newState: StateMachineState) async {
await transitionHandler?(newState)
}
func onTransition(_ handler: @escaping Handler) {
self.transitionHandler = handler
}
}Bonus point: callAsFunction allows us to directly use the transition property and call it as a function. The syntax is even simpler!
@StateTransition private var transition
// Without callAsFunction
await transition.transition(to: OnState())
// With callAsFunction
await transition(to: OnState())All that's left is to coordinate everything in the state machine.
protocol StateMachineContext: Sendable {
static var defaultValue: Self { get }
}
struct EmptyStateMachineContext: StateMachineContext {
static let defaultValue: Self = .init()
}@MainActor
public final class StateMachine<Context: StateMachineContext> {
public private(set) var currentState: StateMachineState
private let initialState: StateMachineState
private let context: Context
private let transitionHandler: StateTransitionHandler = .init()
private var currentStateWrapper: StateMachineStateWrapper<Context>
private nonisolated init(
initialState: StateMachineState,
context: Context
) {
self.initialState = initialState
self.context = context
self.currentStateWrapper = .init(
state: initialState,
transitionHandler: transitionHandler,
context: context
)
}
public func start() async {
transitionHandler.onTransition { [weak self] newState in
guard let self else { return }
await self.transition(to: newState)
}
currentStateWrapper.makeInjection()
await initialState.enter()
}
public func transition(to newState: StateMachineState) async {
await currentState.exit()
currentState = newState
currentStateWrapper.makeInjection()
await currentState.enter()
}
}To get our syntax that allows us to inject our environment, we can use conditional extensions.
public extension StateMachine where Context == EmptyStateMachineContext {
convenience nonisolated init(initialState: StateMachineState) {
self.init(
initialState: initialState,
context: .init()
)
}
func context<C: StateMachineContext>(_ context: C) -> StateMachine<C> {
.init(
initialState: initialState,
context: context
)
}
}In summary
Here is the final syntax of our state machine.
struct Context: StateMachineContext {
let delay: TimeInterval
}
struct OffState: StateMachineState {
@StateTransition private var transition
@StateContext(Context.self) private var context
func enter() async {
Task {
try? await Task.sleep(for: .seconds(context.delay))
await transition(to: OnState())
}
}
}
let stateMachine = StateMachine(initialState: OffState())
.context(Context(delay: 1))In addition to being pretty cool, this syntax is much less prone to errors.
- The state machine is no longer directly referenced in the
enter()function. It's no longer possible to break the functionality of the machine from a state or to forget aweak. Goodbye memory leaks by forgetfulness. - A context can be declared globally in all states. Changing a parameter applies it to all states automatically.
The icing on the cake is that this syntax helps us better understand the black magic of SwiftUI.
To conclude, StateMachine is available as a Swift package on my GitHub. Feel free to take a look to explore the subject further!
