Check how we can efficiently organize many different states by using type safe tree-like structures!

Android mascot handling nested states in a state machine, thanks to Gemini

The journey

This article is part of a series and in the previous articles, I explained my need to create a State Machine to manage complex workflows in the presentation layer of an application. In the second article, I showed how to create the state machine, the DSL to build the states and after that how to trigger side effects. For this final State Machine builder article, we will show how we can create nested states.

All of the articles in the series can be found here:

1. MVS: Model View State Machine
2. MVS series: Building the State Machine
3. MVS series: Implementing side effects in our State Machine
4. MVS series: Creating nested states (you are here)
5. MVS series: State Machine usage and reusable patterns (Coming soon)

All about organization

I can hear you thinking: Why do we need nested states? I thought you would keep it simple! Yes yes, but I also want to provide a small set of tools that can help us cut down on boilerplate. Nested states can help us a lot to manage complex workflows, by hierarchically organizing states. Parent states can define transitions and side effects which the children effectively inherit, making it easier to de-duplicate certain logic, which effectively means we can cut down on boilerplate! But let me show that with a simple example. Let’s say we want to load data with a specific search filter. We basically have a loading state and success and failure states. Every time we update the filter, we effectively reset the search process. Without nested states, we have to duplicate the event to handle new filter inputs.

Let’s define the following states and events:

sealed interface SearchState {

  data class Search(
    val filter: MyFilter,
  ) : SearchState

  data class Failure(
    val filter: MyFilter,
    val error: Throwable,
  ) : SearchState

  data class SearchResults(
    val data: MyData,
  ): SearchState

}

sealed interface SearchEvent {

  data class StartSearch(
    val filter: MyFilter,
  ) : SearchEvents

  data class Failure(
    val error: Throwable,
  ) : SearchEvents

  data object Retry : SearchEvent

  data class Success(
    val data: MyData,
  )
}

As mentioned, we would like an update to our filter to re-run the search, so with the current prototype of our State Machine it would look something like this:

val searchStateMachine = stateMachine<SearchState, SearchEvent>(
  scope = viewModelScope,
  initialState = SearchState.Search(filter = MyFilter(query = "")),
) {
  state<SearchState.Search> {
    // Allow to reset the search with start search event #1
    onEvent<SearchEvent.StartSearch> { _, event ->
      SearchState.Search(filter = event.filter)
    }
    onEvent<SearchEvent.Success> { _, event ->
      SearchState.SearchResults(data = event.data)
    }
    onEvent<SearchEvent.Failure> { state, event ->
      SearchState.Failure(
        filter = event.filter,
        error = event.error,
      )
    }
    sideEffect { state ->
      searchMyDataUseCase.run(state.filter)
        .fold(
          onSuccess = SearchState::SearchResults,
          onFailure = SearchState::Failure
        )
        .also(::onEvent)
    }
  }
  state<SearchState.Failure> {
    // Allow to reset the search with start search event #2
    onEvent<SearchEvent.StartSearch> { _, event ->
      SearchState.Search(filter = event.filter)
    }
    onEvent<SearchEvent.Retry> { state, _ ->
      SearchState.Search(filter = state.filter)
    }
  }
  state<SearchState.SearchResults> {
    // Allow to reset the search with start search event #3
    onEvent<SearchEvent.StartSearch> { _, event ->
      SearchState.Search(filter = event.filter)
    }
  }
}

You can see here that the StartSearch event is defined 3 times! How can we do better than this? Well, if our stateMachine can handle generic events, that would be great!

val searchStateMachine = stateMachine<SearchState, SearchEvent>(
  scope = viewModelScope,
  initialState = SearchState.Search(filter = MyFilter(query = "")),
) {
  // We simply handle the event in the parent, now all states
  // will inherit this behavior
  onEvent<SearchEvent.StartSearch> { _, event ->
    SearchState.Search(filter = event.filter)
  }
  state<SearchState.Search> {
    onEvent<SearchEvent.Success> { _, event ->
      SearchState.SearchResults(data = event.data)
    }
    onEvent<SearchEvent.Failure> { state, event ->
      SearchState.Failure(
        filter = event.filter,
        error = event.error,
      )
    }
    sideEffect { state ->
      searchMyDataUseCase.run(state.filter)
        .fold(
          onSuccess = SearchState::SearchResults,
          onFailure = SearchState::Failure
        )
        .also(::onEvent)
    }
  }
  state<SearchState.Failure> {
    onEvent<SearchEvent.Retry> { state, _ ->
      SearchState.Search(filter = state.filter)
    }
  }
  state<SearchState.SearchResults>()
}

You can see that we only define the StartSearch events once and the transition will always run when the event comes in. These kinds of optimizations will make our code less error prone, simply because we have less code to write.

Setting the rules

Just like last time, we can define a set of rules to define the behavior that we want:

• We should be able to define nested states. This is basically like building a state machine within a state machine, so we can inherit all the nice side effects, transitions and state definition builders.

Press enter or click to view image in full size

Nested State Machine diagram

• The children of a nested state will inherit from the nested state itself, which in term is a sub type of its parent. In this way we can keep everything nicely organized.
• Any state in the machine can transition to any other state -> we do not constrain children to only transition to children of the same parent, but we can transition to any other state.
• Our state machine should be able to handle events, just like a state. Inner states will get the first opportunity to transition on an event, and unhandled events will flow up to the parents. Only 1 state or one of its ancestors can handle an event.

Events flow up from state all the way to the root if necessary

• Our state machine should also be able to trigger side effects, just like a state. Side effects should be triggered ancestor first, and cancelled child first.

Getting started

So we have an existing State Machine and builder, we should make some amendments. There are 2 main things we have to achieve here:

• Adding generics for nested states
• Creating reusable builders.

Generics

We have to make some adjustments. We already have our type restricted, where we only want to register states inside our state machine that are somehow inheriting from the main state. We also use the same type to restrict the return type of the transition we trigger when we handle an event:

data object RandomObject

sealed interface ExampleState {
  data object GoodStuff: ExampleState
  data object OtherGoodStuff: ExampleState
}

sealed interface ExampleEvent {
  data object GoToOther: RandomEvent
}

// We see our stateMachine takes states of type ExampleState
// and events of type ExampleEvent
val stateMachine = stateMachine<ExampleState, ExampleEvent>(..) {
  // ERROR RandomObject doesn't inherit ExampleState!!
  state<RandomObject>()

  state<ExampleState.GoodStuff> {
    // ERROR RandomObject doesn't inherit ExampleEvent
    onEvent<RandomObject> { _, _ ->
      ExampleState.OtherGoodStuff
    }

    onEvent<ExampleEvent.GoToOther> { _, _ ->
      ExampleState.OtherGoodStuff
    }
  }

  state<ExampleState.OtherGoodStuff> {
    onEvent<ExampleEvent.GoToOther> { _, _ ->
      // ERROR RandomObject doesn't inherit ExampleState
      RandomObject
    }
  }
}

We will need to find a way to restrict defining nested states, but still allow transitioning to all the states we defined in our machine! In order to achieve this, our StateMachineBuilder and StateMachineDefinition should have an extra generic type: The CurrentState type. In this way, we can restrict the builder to only allow the proper nested types, but allow it to transition to any other state in the machine. Also, we are going to give it the same type of state definition as any other state, as we want it to contain all the side effects and transitions. Furthermore, we want to keep track of all the nested states:

// Old
data class StateMachineDefinition<State : Any, Event : Any>(
  val states: Map<KClass<out State>, StateDefinition<State, State, Event>>,
)

// New
data class StateMachineDefinition<State : Any, CurrentState : State, Event : Any>(
  val self: StateDefinition<State, CurrentState, Event>,
  val states: Map<KClass<out CurrentState>, StateDefinition<State, CurrentState, Event>>,
  val nestedStates: Map<KClass<out CurrentState>, StateMachineDefinition<State, CurrentState, Event>>,
)

Furthermore, our state machine builder should also get the new DSL methods to add side effects and transitions. We already have that logic in our StateBuilder, so we will extract it and put it in 2 separate builders: SideEffectBuilder and TransitionBuilder.

data class SideEffect<State : Any, in CurrentState : State, out Event : Any>(
  val key: (CurrentState) -> Any,
  val effect: suspend StateMachine<State, Event>.(CurrentState) -> Unit,
)

@StateDsl
interface SideEffectBuilder<State : Any, CurrentState : State, Event : Any> {

  fun sideEffect(
    key: (CurrentState) -> Any = { it },
    effect: suspend StateMachine<State, Event>.(CurrentState) -> Unit,
  )

  fun buildSideEffects(): List<SideEffect<State, CurrentState, Event>>

}

internal class DefaultSideEffectBuilder<State : Any, CurrentState : State, Event : Any> :
  SideEffectBuilder<State, CurrentState, Event> {

  private var sideEffects: List<SideEffect<State, CurrentState, Event>> = emptyList()

  override fun sideEffect(
    key: (CurrentState) -> Any,
    effect: suspend StateMachine<State, Event>.(CurrentState) -> Unit,
  ) {
    sideEffects += SideEffect(key, effect)
  }

  override fun buildSideEffects() = sideEffects

}

data class StateTransition<State : Any, CurrentState : State, Event : Any>(
  val transition: (CurrentState, Event) -> State,
)

@StateDsl
interface TransitionBuilder<State : Any, CurrentState : State, Event : Any> {

  fun <E : Event> onEvent(eventClass: KClass<E>, transition: (CurrentState, E) -> State)
  fun buildTransitions(): Map<KClass<out Event>, StateTransition<State, CurrentState, Event>>

}

internal class DefaultTransitionBuilder<State : Any, CurrentState : State, Event : Any> :
  TransitionBuilder<State, CurrentState, Event> {

  private var transitions = mapOf<KClass<out Event>, StateTransition<State, CurrentState, Event>>()

  override fun <E : Event> onEvent(
    eventClass: KClass<E>,
    transition: (CurrentState, E) -> State,
  ) {
    @Suppress("UNCHECKED_CAST")
    val stateTransition = StateTransition(transition) as StateTransition<State, CurrentState, Event>
    transitions = transitions
      .plus(eventClass to stateTransition)
  }

  override fun buildTransitions() = transitions

}

From here on, the StateBuilder and StateMachineBuilder can easily reuse the logic. I will show the direct changes to the StateBuilder first and you should just assume the same for the StateMachineBuilder:

 class StateBuilder<State : Any, CurrentState : State, Event : Any>:
   SideEffectBuilder<State, CurrentState, Event> by DefaultSideEffectBuilder(),
   TransitionBuilder<State, CurrentState, Event> by DefaultTransitionBuilder() {

  inline fun <reified E : Event> onEvent(
    noinline transitionBuilder: (CurrentState, E) -> State,
  ) {
    onEvent(E::class, transitionBuilder)
  }

  fun build(): StateDefinition<State, CurrentState, Event> {
    return StateDefinition(
      transitions = buildTransitions(),
      sideEffects = buildSideEffects(),
    )
  }

}

Here you can see we implement the SideEffectBuilder and TransitionBuilder, but delegate the actual implementation to their respective builders. Also: we keep the inline function here because we cannot define inline functions from the interfaces. So our StateMachineBuilder will have the same code, including the convenience inline function to help us build everything more efficiently.

After this, our StateDefinition needs to keep track of its parent as we can have nested states. Also, we can keep track of its own class, which will be helpful later when we will deal with the side effects:

data class StateDefinition<State : Any, CurrentState : State, Event : Any>(
  val clazz: KClass<out State>, // Added for convenience
  val parent: KClass<out State>?, // To keep track of our parent
  val sideEffects: List<SideEffect<State, CurrentState, Event>>,
  val transitions: Map<KClass<out Event>, StateTransition<State, CurrentState, Event>>,
)

Our StateBuilder will finally also need to add these 2 properties to its constructor:

 @StateDsl
class StateBuilder<State : Any, CurrentState : State, Event : Any>(
  private val clazz: KClass<CurrentState>,
  private val parent: KClass<out State>,
) :
  SideEffectBuilder<State, CurrentState, Event> by DefaultSideEffectBuilder(),
  TransitionBuilder<State, CurrentState, Event> by DefaultTransitionBuilder() {
// ....

Then we have to go back to our StateMachineBuilder itself, and add the functionality to build nested states:

 @StateDsl
class StateMachineBuilder<State : Any, CurrentState : State, Event : Any>(
  private val clazz: KClass<out State>,
  private val parent: KClass<out State>?,
) :
  SideEffectBuilder<State, CurrentState, Event> by DefaultSideEffectBuilder(),
  TransitionBuilder<State, CurrentState, Event> by DefaultTransitionBuilder() {

  // other code...

  var nestedStates = emptyMap<KClass<out CurrentState>, StateMachineDefinition<State, CurrentState, Event>>()
    private set

  // other code...

  inline fun <reified S : CurrentState> nestedState(
    noinline nestedStateBuilder: StateMachineBuilder<State, S, Event>.() -> Unit,
  ) {
    nestedState(S::class, nestedStateBuilder)
  }

  // other code...

  fun <S : CurrentState> nestedState(
    stateClass: KClass<out S>,
    stateBuilder: StateMachineBuilder<State, S, Event>.() -> Unit,
  ) {
    @Suppress("UNCHECKED_CAST")
    val nestedState = StateMachineBuilder<State, S, Event>(stateClass, clazz)
      .apply(stateBuilder)
      .build() as StateMachineDefinition<State, CurrentState, Event>
    this.nestedStates += nestedState.nestedStates.plus(stateClass to nestedState)
  }

  // other code...

}

Here you can see that we have added very similar code as how we were originally adding states to the machine, just now we not just collect a map with <KClass, StateDefinition> but also a map with <KClass, StateMachineDefinition>. Finally, we should adjust our build function

fun build(): StateMachineDefinition<State, CurrentState, Event> {
  return StateMachineDefinition(
    self = StateDefinition(
      clazz = clazz,
      parent = parent,
      sideEffects = buildSideEffects(),
      transitions = buildTransitions(),
    ),
    states = states + buildMap {
      nestedStates.values.forEach { nestedState ->
        putAll(nestedState.states)
      }
    },
    nestedStates = nestedStates + buildMap {
      nestedStates.values.forEach { nestedState ->
        putAll(nestedState.nestedStates)
      }
    },
  )
}

The whole class can be found here.

Adjusting the StateMachine

The status quo is that we can react to events and trigger side effects for a single state. However, with the current changes we also need to take the ancestors of the current state into account. In order to do this, we make a function that helps us in collecting the ancestors for a State

internal class DefaultStateMachine<out State : Any, in Event : Any>(
  private val scope: CoroutineScope,
  private val initialState: State,
  private val started: SharingStarted = SharingStarted.WhileSubscribed(5000),
  private val definition: StateMachineDefinition<State, State, Event>,
) : StateMachine<State, Event> { 

  // Other code...

  /**
   * Collect the definition of the state and all nested parents
   */
  private fun State.getSelfAndAncestors(): List<StateDefinition<State, in State, Event>> = buildList {
    definition.states[this@getSelfAndAncestors::class]?.let { state ->
      add(state)
      var clazz = state.parent
      while (clazz != null) {
        clazz = definition.nestedStates[clazz]?.let { nestedState ->
          add(nestedState.self)
          nestedState.self.parent
        }
      }
    }

    // Finally add the base definition of the state machine to the list
    add(definition.self)
  }

  // Other code...

}

With this, we can make an adjustment to our handleEvents extension function:

@Suppress("AssignedValueIsNeverRead")
private fun Flow<Event>.handleEvents(): Flow<State> {
  var lastState: State = initialState
  return runningFold(initial = { lastState }) { state, event ->
    state.getSelfAndAncestors()
      // SelfAndAncestors tree is already sorted (self first, root last)
      // Therefore, we can just find the first one that has a transition
      // defined for the current event
      .firstNotNullOfOrNull { definition -> definition.transitions[event::class] }
      ?.transition(state, event)
      ?: state
  }
    .onEach { lastState = it }
}

After this, we can do a similar thing with our side effects!

private inner class SideEffectsHandler(
  private val upstream: Flow<State>,
) : Flow<State> {

  @Suppress("AssignedValueIsNeverRead")
  override suspend fun collect(collector: FlowCollector<State>) {
    coroutineScope {
      var sideEffectJobs = mapOf<JobKey<State>, Job>()

      upstream.collect { state ->

        // Collect all side effects we wish to run right now
        // We will make sure these are ordered root first and child last
        val sideEffects = buildMap {
          state.getSelfAndAncestors().forEach { stateDefinition ->
            stateDefinition.sideEffects.forEachIndexed { index, sideEffect ->
              put(
                JobKey(stateDefinition.clazz, index, sideEffect.key(state)),
                sideEffect.effect,
              )
            }
          }
        }

        // Cancel any old nested side effects, youngest jobs first
        // In this case we need to reverse the entries of the existing
        // Side effect jobs because younger side effects are added to the
        // end of the map
        sideEffectJobs.minus(sideEffects.keys)
          .values
          .reversed()
          .forEach { it.cancelAndJoin() }

        // Start non-existing nested side effects, outer (parents) first
        // We store these in the sideEffectJobs. Our sideEffects
        // Are already correctly ordered, so here we can use a simple
        // mapValues operation to either return an existing job or 
        // run a new one in the right order
        sideEffectJobs = sideEffects.mapValues { (jobKey, sideEffect) ->
          sideEffectJobs.getOrElse(jobKey) {
            launch { sideEffect(this@DefaultStateMachine, state) }
          }
        }

        // Continue the flow like nothing happened :)
        collector.emit(state)
      }
    }
  }

}

All the code of the whole StateMachine class can be found here, and perhaps some smaller changes didn’t make the cut to be in this article, but they should be visible in the commit history.

The result

We can now utilize nesting to create clear and logically grouped states, create reusable event handling and trigger side effects from different levels. I have updated the form example to utilize this, but obviously in such a simple case, the benefits are limited. However, I am afraid that my idea to build a fake webshop with a basket, fake checkout and fake payment, would interfere with my plans to actually have a family life, as this is my own free time :). So I did my best to boost up the form example with these changes, even if nested states and such will have much better use in more complex workflows:

sealed interface FormState {

  data class LoadingFormData(
    val simulateLoadingFailure: Boolean = false,
  ) : FormState

  data class FormLoadingFailure(val error: Throwable) : FormState

  sealed interface WithData : FormState {
    val value: String

    data class PendingInput(override val value: String) : WithData
    data class SavingForm(override val value: String) : WithData
    data class SavingFailure(
      override val value: String,
      val error: Throwable,
    ) : WithData

  }

  data object Success : FormState

}

sealed interface FormEvent {

  data class Failure(val error: Throwable) : FormEvent
  data class LoadingSuccess(val value: String) : FormEvent
  data object Retry : FormEvent
  data class Update(val value: String) : FormEvent
  data object SavingSuccess : FormEvent
  data object Save : FormEvent
  data object Reset : FormEvent

}

And then the view model can change to:

class FormViewModel(
  private val fetchFormDataUseCase: FetchFormDataUseCase,
  private val saveFormDataUseCase: SaveFormDataUseCase,
) : ViewModel() {

  private val formStateMachine = stateMachine<FormState, FormEvent>(
    scope = viewModelScope,
    initialState = FormState.LoadingFormData(simulateLoadingFailure = true),
  ) {
    sideEffect { state ->
      println("Current state: ${state::class.simpleName}")
    }
    onEvent<FormEvent.Reset> { _, _ ->
      FormState.LoadingFormData(simulateLoadingFailure = true)
    }
    state<FormState.LoadingFormData> {
      sideEffect { state ->
        fetchFormDataUseCase.run(simulateFailure = state.simulateLoadingFailure)
          .fold(
            onSuccess = FormEvent::LoadingSuccess,
            onFailure = FormEvent::Failure,
         )
         .also(::onEvent)
      }
      onEvent<FormEvent.LoadingSuccess> { _, event ->
        FormState.WithData.PendingInput(event.value)
      }
      onEvent<FormEvent.Failure> { _, event ->
        FormState.FormLoadingFailure(event.error)
      }
    }
    state<FormState.FormLoadingFailure> {
      onEvent<FormEvent.Retry> { _, _ ->
        FormState.LoadingFormData()
      }
    }
    nestedState<FormState.WithData> {
      sideEffect(key = { Unit }) { 
        println("Started editing form data")
        try {
          awaitCancellation()
        } finally {
          println("Finished editing form data!")
        }
      }
      state<FormState.WithData.PendingInput> {
        onEvent<FormEvent.Update> { state, event ->
          state.copy(value = event.value)
        }
        onEvent<FormEvent.Save> { state, _ ->
          FormState.WithData.SavingForm(value = state.value)
        }
      }
      state<FormState.WithData.SavingForm> {
        sideEffect { state ->
          saveFormDataUseCase.run(state.value)
            .fold(
              onSuccess = { FormEvent.SavingSuccess },
              onFailure = FormEvent::Failure,
            )
            .also(::onEvent)
        }
        onEvent<FormEvent.SavingSuccess> { _, _ ->
          FormState.Success
        }
        onEvent<FormEvent.Failure> { state, event ->
         FormState.WithData.SavingFailure(state.value, event.error)
        }
      }
      state<FormState.WithData.SavingFailure> {
        onEvent<FormEvent.Save> { state, _ ->
          FormState.WithData.SavingForm(state.value)
        }
        onEvent<FormEvent.Update> { _, event ->
          FormState.WithData.PendingInput(event.value)
        }
      }
    }

    state<FormState.Success>()
  }

  val viewState = formStateMachine.state.map { state ->
    when (state) {
      is FormState.LoadingFormData -> FormViewState.Loading
      is FormState.FormLoadingFailure -> FormViewState.Failure
      is FormState.WithData -> FormViewState.FormInput(
        value = state.value,
        inputErrorMessage = (state as? FormState.WithData.SavingFailure)?.error?.message,
        showLoading = state is FormState.WithData.SavingForm,
      )

      FormState.Success -> FormViewState.Success
    }
  }.stateIn(
    scope = viewModelScope,
    started = SharingStarted.WhileSubscribed(5000),
    initialValue = FormViewState.Loading,
  )

  fun onEvent(formEvent: FormEvent) {
    formStateMachine.onEvent(formEvent)
  }

}

This article was previously published on proandroiddev.com