Sometimes you encounter problems that you feel are solvable on the one hand, but on the other hand you also feel that you only have this problem because you put yourself into the situation. So without further ado, lets get into the topic of scoping a Jetpack Architecture ViewModel to our Composables!

Scoping a ViewModel the normal way

Even though you could scope a ViewModel in Android Auto when you really want it, normally you scope the ViewModel to one of the following:

1. The Activity
2. The Fragment (even though with Compose you shouldn’t use Fragments no longer)
3. The navigation graph
4. The navigation destination

The ViewModel in Jetpack Compose usually relies on the LocalViewModelStoreOwner: Once you use your first composable it is already there, waiting for you to use!

class MainActivity: ComponentActivity {

  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    setContent {
      MyComposable()
    }
  }

}

@Composable
fun MyComposable() {
  val viewModelStoreOwner = LocalViewModelStoreOwner.current
  // viewModelStoreOwner == MainActivity
  
  // We can use the following helper function to simply get a ViewModel
  // Assuming it has no dependencies:
  val viewModel = viewModel<MyViewModel>()
}

Now if we would inject ViewModel in the simplest of ways, we can see that in the viewModel(..) helper function, we have a default value for the the ViewModelStoreOwner parameter: the LocalViewModelStoreOwner:

Nesting with NavGraph

I like the NavGraph and one of the reasons is: It has great support for ViewModels! Every destination in your graph will have its own ViewModel store, because every NavBackStackEntry is a ViewModelStoreOwner.

@Composable
fun MyComposable() {
  println(LocalViewModelStoreOwner?.current?.javaClass?.simpleName)
  // prints "MainActivity"
  
  val navController = rememberNavController()
  NavHost(
    navController = navController,
    startDestination = "home",
  ) {
    composable("home") { homeBackstackEntry ->
      println(LocalViewModelStoreOwner.current == homeBackstackEntry)
      // prints "true"
    }
    composable("other") { otherBackstackEntry ->
      println(LocalViewModelStoreOwner.current == otherBackstackEntry)
      // prints "true"
    }
  }
}

This also means that if you move from one destination to the other, the nav graph will be able to save your ViewModel and destroy your ViewModel whenever it is necessary. You will see that if you would go from home to other and then pop the navigation stack, any ViewModel in the home destination will magically reappear as it was saved. But as soon as you pop the stack on your other destination, any ViewModel there will be cleared directly!

Scope a ViewModel to a Composable function

What I want, is to scope a ViewModel to the Composable function it is inside of! I want the ViewModel to be created when the Composable function enters composition, and the ViewModel to be cleared when the Composable leaves the composition.

As you know, ViewModels are stored in the ViewModelStore inside a ViewModelStoreOwner. So what we need to do, is basically remember a ViewModelStore somehow inside our Composable, use that to spawn and a ViewModel and then we should be good to go! My requirements are simple:

1. I want to have a ViewModel scoped to my Composable function. It should be created when my Composable function becomes part of the composition and it should be cleared if the Composable function is leaving the composition.
2. I want the ViewModel to survive orientation changes.
3. I want to be able to pass a key, and if that key changes, my ViewModel should be re-initialized with the new key.

@Composable
fun Example() {
  // If condition is true, the code block enters the composition
  // Ifcondition turns to false, the following code leaves the composition
  if (condition) {
    // Due to the ViewModelStoreOwner being some kind of root component
    // like an Activity or a NavGraph destination, the viewModel will
    // remain the same when we flip the condition multiple times.

    // I want this to be different: As soon as this piece of code leaves
    // the composition, it should clear any ViewModel it contains.
    // If it reappears, I want a fresh ViewModel again.
    val viewModel = viewModel<ExampleViewModel>()
    val text by viewModel.text.collectAsState()
    Text(text)
  }
}

Use Cases

This is the hardest part to explain, as I need to show why it is worth it and how no other tool can work. We have some use cases where we would like to have components nested into other components.

• Reusable building blocks: If would be nice to create Composables as building blocks and use those in your code that can have their own state and ViewModel inside of them. This can avoid the problem of having God-like ViewModels that control way too much content on the screen.
• Bottom sheets: In the application I currently work on, we have one screen supported by multiple bottom sheets that pop up if you push buttons. These bottom sheets help you to change certain pieces of the data, and they live inside a ModalBottomSheet component. These separate bottom sheets benefit from having their own ViewModel to encapsulate the logic inside, and it would be good to clear the inner ViewModel when the bottom sheet closes or changes to a different one.
• Encapsulated views: We have a screen that shows a planner. On top of the planner is a date picker. Every time you pick a date (it offers tomorrow/yesterday), the screen slides like to left or right depending on what you pressed and we load the data from the backend what we display on the new screen. It is nice to have a parent ViewModel that handles the date picking, and a child ViewModel that handles the data for that date, encapsulated within the main component.
• Complex ViewPagers: Like above, if you have a ViewModel that can feed a list of id’s to a ViewPager, then every page can essentially host itself if it has its own ViewModel. It would be nice to discard unused ViewModels after scrolling to the next pages.
• Dialog: Here you can probably use the NavHost to handle your dialogs, but I had the dialog inside my normal component and I didn’t want to show old data after dismissing it and reopening it. I also want any dialog to be self sustainable: If it offers certain actions, it should have its own ViewModel to take care of business.

First iteration

So we can use the good old remember function to remember a value, then upon key changes clear the old value and create the new!

@Composable
fun rememberViewModelStoreOwner(key: Any?): ViewModelStoreOwner() {
  val viewModelStoreOwner = remember(key) { 
    object: ViewModelStoreOwner {
      override val viewModelStore = ViewModelStore()
    }
  }
  DisposableEffect(viewModelStoreOwner) {
    onDispose {
      viewModelStoreOwner.clear()
    }
  }
  return viewModelStoreOwner
}

@Composable
fun WithViewModelStoreOwner(
  key: Any?,
  content: @Composable () -> Unit,
) {
  CompositionLocalProvider(
    value = LocalViewModelStoreOwner provides rememberViewModelStoreOwner(key),
    content = content,
  )
}

This is what you see here: We remember the ViewModelStoreOwner based on a key. Then with a DisposableEffect we always clear the ViewModelStoreOwner whenever it changes or the composable is finished. It is effective. But what it doesn’t do, is survive orientation changes: The main thing we want a ViewModel to do! Also, if you navigate away from the screen it triggers the onClear, so if you pop back to the screen the ViewModel is gone as well.

A closer look at the NavHost

I was wondering: How do they do it from the Jetpack Navigation library? How can they maintain state while navigating away from one destination to another? The answer is quite simple: They use…. a ViewModel! That is right! They have a ViewModelStore right inside a ViewModel inside the NavHost (somewhere). So basically, any ViewModel you have inside the navigation graph is stored inside another ViewModel and the navigation component is smart enough to keep track of it all, as you can see in the following snippets of the internal code:

You can see that every backStackEntryId will have a ViewModelStore linked to it, and upon popping and navigating it can be cleared or created, depending on the navigation commands you give it!

Updating our code

Knowing this, I can also incorporate the power of a ViewModel to store my ViewModels! In this case, I will make my ViewModel responsible of returning the right ViewModelStoreOwner based on the key it is given. If the key changes, we clear the ViewModelStore before returning it.

@Composable
fun rememberViewModelStoreOwner(
  key: Any?,
): ViewModelStoreOwner {
  val viewModelStoreOwnerViewModel = viewModel<ViewModelStoreOwnerViewModel>() {
    ViewModelStoreOwnerViewModel(key)
  }
  viewModelStoreOwnerViewModel.updateKey(key)

  DisposableEffect(Unit) {
    onDispose {
      viewModelStoreOwnerViewModel.resetViewModelStore()
    }
  }
  return remember(viewModelStoreOwnerViewModel) { 
    viewModelStoreOwnerViewModel.viewModelStore
  }
}

@Composable
fun WithViewModelStoreOwner(
  key: Any?,
  content: @Composable () -> Unit,
) {
  CompositionLocalProvider(
    value = LocalViewModelStoreOwner provides rememberViewModelStoreOwner(key),
    content = content,
  )
}

private class ViewModelStoreOwnerViewModel(
  private var key: Any?,
) : ViewModel() {

  private var viewModelStore: ViewModelStoreOwner = object : ViewModelStoreOwner {
    override val viewModelStore = ViewModelStore()
  }

  fun get(key: Any?): ViewModelStoreOwner {
    if (key != this.key) {
      this.key = key
      resetViewModelStore()
    }
    return viewModelStore
  }

  fun resetViewModelStore() {
    viewModelStore.viewModelStore.clear()
  }

  override fun onCleared() {
    resetViewModelStore()
  }

}

However, we are still not there. We still depend on the DisposableEffect to clear the ViewModel, and onDispose {} is definitely called on configuration changes! Also, if I would use rememberViewModelStoreOwner multiple times, it would give me the same instance. This could mean if Composable A is disposed of, we could lose the ViewModel in Composable B :(. Also, if we would navigate back and forth, we will lose our data, because of onDispose.

Unique key for every instance

Lets first solve the problem of uniqueness: what we wish to do is to create a single ViewModelStoreOwner per Composable function. Luckily, Compose keeps track of the position of your Composable with a certain hash: currentCompositeKeyHash. It actually uses this under the hood of rememberSaveable. When you call rememberSaveable, the current hash of the composable is used as a key to store the data when we save state. If you recreate a composable, it will generate the same hashes and then you can restore any value based on that. Knowing that our composables will have pretty stable keys by relying on this field, we can create a unique ViewModelStoreOwnerViewModel for every key:

@Composable
fun rememberViewModelStoreOwner(
  key: Any?,
): ViewModelStoreOwner {
  val viewModelKey = "rememberViewModelStoreOwner#" + currentCompositeKeyHash.toString(36)
  val viewModelStoreOwnerViewModel = viewModel<ViewModelStoreOwnerViewModel>(
    key = viewModelKey,
  ) {
    ViewModelStoreOwnerViewModel(key)
  }
  viewModelStoreOwnerViewModel.updateKey(key)

  DisposableEffect(Unit) {
    onDispose {
      viewModelStoreOwnerViewModel.resetViewModelStore()
    }
  }
  return remember(viewModelStoreOwnerViewModel) { 
    viewModelStoreOwnerViewModel.viewModelStore
  }
}

@Composable
fun Example() {
  // All these will have their own unique ViewModelStoreOwner
  val viewModelStoreOwner1 = rememberViewModelStoreOwner(Unit)
  val viewModelStoreOwner2 = rememberViewModelStoreOwner(Unit)
  if (condition) {
    // This means that whenever condition turns into false, only
    // viewModelStoreOwner3 will be cleared
    val viewModelStoreOwner3 = rememberViewModelStoreOwner(Unit)
  }

Save me!

So now lets tackle surviving configuration changes. When there are configuration changes, the Activity will recreate itself. Just before it does this, it will trigger a onSaveInstanceState. This in turn will trigger the save methods in all its children Fragments, Views and Composables! We are very interested in this because if we can detect when we are saving changes, we know that we either get:

• Navigated away with the option to restore state later
• App destroyed by system in background to save resources
• App is recreated due to config changes, like rotation

I started to write a whole solution to detect the moment we save using LocalSavedStateRegisterOwner but I figured out it was not reliable. The problem is that saving state is triggered not when we navigate away in all cases, and also when the app is simply in background it is triggered. We need something more reliable!

Lifecycle to the rescue (As always…)

So after throwing away a bunch of code and a bunch of time, I came up with the following conditions to keep or destroy a ViewModel, based on the lifecycle where the Composable is in. The lifecycle is interesting, as it can be the lifecycle of your Activity or the lifecycle of your navigation back stack entry. The nifty thing is that if you navigate away from your current destination within the NavHost, the back stack entry will not be destroyed. It will return to a CREATED state. Only when the back stack entry is disposed of will the lifecycle become DESTROYED.

1. If our Composable is removed and we are in a resumed state, we know the component holding the Lifecycle is still alive but our Composable was simply removed from it. So in this case we clear the ViewModelStoreOwner. The next time our Composable is added again, it will get a fresh ViewModel!
2. If our lifecycle is not resumed and our Composable is removed, we keep the data. It can be that we navigate to another location (so our current location will be Stopped but not Destroyed), or that we are experiencing a configuration change and the lifecycle will be recreated. If our navigation destination was actually destroyed, its internal ViewModelStore will be cleared, and our own ViewModelStoreOwnerViewModel will be cleared like magic as well.
3. If our Composable is removed and our lifecycle gets to RESUMED and our Composable is not re-created in the mean time, it means our Composable was actually removed during recreation or while the app was in the background. In this case we clear the ViewModel upon resuming.

Keeping the right states

We should store the last known lifecycle inside our ViewModel, as well as whether we are attached or not. Now our ViewModel will decide whether it clears the ViewModelStore or not. Therefore, our rememberViewModelStoreOwner turns into a simpler implementation:

@Composable
fun rememberViewModelStoreOwner(
  key: Any?,
): ViewModelStoreOwner {
  val viewModelKey = "rememberViewModelStoreOwner#" + currentCompositeKeyHash.toString(36)
  val localLifecycle = LocalLifecycleOwner.current.lifecycle
  val viewModelStoreOwnerViewModel = viewModel<ViewModelStoreOwnerViewModel>(key = viewModelKey) {
    ViewModelStoreOwnerViewModel(key, localLifecycle)
  }
  viewModelStoreOwnerViewModel.update(key, localLifecycle)

  DisposableEffect(Unit) {
    viewModelStoreOwnerViewModel.attachComposable()
    onDispose {
      viewModelStoreOwnerViewModel.detachComposable()
    }
  }
  return remember(viewModelStoreOwnerViewModel) {
    viewModelStoreOwnerViewModel.viewModelStore
  }
}

You can see that we pass 3 things to our ViewModel:

1. The key with which we decide to return a new ViewModelStoreOwner
2. The current lifecycle the Composable is in (this may outlive the Composable)
3. Whether our Composable is attached or not, by checking when it is being disposed of.

Updating the ViewModel

Now we just have to implement the 3 rules we stated earlier inside our ViewModel. So here we clear the stored view models for the 3 cases:

1. Our key changed
2. We detached and lifecycle state is RESUMED
3. We reach a resumed state later on but we are no longer attached

If we are attached or our lifecycle is destroyed (could happen during recreation on a configuration change) then we do nothing. So we can simply combine some states and collect these with our viewModelScope.

@Stable
private class ViewModelStoreOwnerViewModel(
  private var key: Any?,
  initialLifecycle: Lifecycle,
) : ViewModel() {

  private val attachedLifecycle = MutableStateFlow<Lifecycle?>(initialLifecycle)
  private val isAttachedToComposable = MutableSharedFlow<Boolean>()

  val viewModelStore: ViewModelStoreOwner = object : ViewModelStoreOwner {
    override val viewModelStore = ViewModelStore()
  }

  init {
    viewModelScope.launch {
      attachedLifecycle
        .flatMapLatest { lifecycle -> lifecycle?.eventFlow ?: emptyFlow() }
        .collectLatest {
          // Make sure we release the lifecycle once it is destroyed
          if (it == Lifecycle.Event.ON_DESTROY) {
            attachedLifecycle.update { null }
          }
        }
    }
    viewModelScope.launch {
      isAttachedToComposable.collectLatest { isAttached ->
        when {
          // If we are attached or we are destroyed, we do not need to do anything
          isAttached || attachedLifecycle.value == null -> return@collectLatest
          // If we are detached and the lifecycle state is resumed, we should reset the view model store
          attachedLifecycle.value?.currentState == Lifecycle.State.RESUMED -> resetViewModelStore()
          else -> {
            // We wait for the lifecycle event ON_RESUME to be triggered before resetting the ViewModelStore
            // If in the mean time we are attached again, this work is cancelled
            attachedLifecycle
              .flatMapLatest { lifecycle -> lifecycle?.eventFlow ?: emptyFlow() }
              // Wait for first event that matches ON_RESUME.
              .firstOrNull { it == Lifecycle.Event.ON_RESUME } ?: return@collectLatest
            resetViewModelStore()
          }
        }
      }
    }
  }

  fun update(key: Any?, lifecycle: Lifecycle) {
    if (key != this.key) {
      this.key = key
      resetViewModelStore()
    }
    attachedLifecycle.update { lifecycle }
  }

  fun attachComposable() {
    viewModelScope.launch { isAttachedToComposable.emit(true) }
  }

  fun detachComposable() {
    viewModelScope.launch { isAttachedToComposable.emit(false) }
  }

  override fun onCleared() {
    super.onCleared()
    resetViewModelStore()
  }

  private fun resetViewModelStore() {
    viewModelStore.viewModelStore.clear()
  }

}

This seems to work perfectly fine! However, if we have a lot of composables coming and going, it would be nice to not have a million ViewModelStoreOwnerViewModels spawned, each with a unique key. We could wrap it all inside one and only one ViewModel instance and keep track of the different ViewModelStore for the different hash keys in there, stored inside a map. In that way, any time a Composable leaves the screen, we can clear the view model store and remove everything from memory.

The gist of it

The full implementation that uses the map to store the ViewModelStoreOwners, together with the Composables, can be found in this gist.

Example time!

It is time to show off the results. In the following you can see in what ways we can use the view model store owner in our ContrivedExample:

@Composable
fun ContrivedExample(
  state: State,
  someCondition: Boolean,
  dialogState: SomeDialogState?,
) {
  when(state) {
    State.Something -> {
      WithViewModelStoreOwner(state.originalData) {
      // Cleared when state turns into State.Something, 
      // Recreated when originalData is changing
      val viewModel = viewModel {
        SomethingViewModel(initialValue = state.originalData)
      }
        SomethingComponent(state.originalData)
      }
    }
    State.SomethingElse -> {
      WithViewModelStoreOwner(state.otherData) {
        // Cleared when state turns into State.Something, 
        // recreated when otherData is changing
        val viewModel = viewModel {
          SomethingElseViewModel(initialValue = state.otherData)
        }
        SomethingElseComponent()
      }
    }
  }
  if (someCondition) {
    WithViewModelStoreOwner(Unit) {
      // If someCondition turns false, this viewModel is cleared magically
      val viewModel = viewModel<SpecialViewModel>()
      ComposableShownWithSomeCondition(viewModel = viewModel)
    }
  }
  dialogState?.let {
    WithViewModelStoreOwner(it) {
      // If our dialog state changes, view model will reinitialize
      // If our dialog state turns to null, view model is cleared
      val viewModel = viewModel<SpecialViewModel>()
      ComposableShownWithSomeCondition(viewModel = viewModel)
    }
  }
}

Okay those were great examples! But now is the time for some real moving pictures. I will refer to my earlier example project where I have a LoadingResult and a standard component to handle loading/failure/success states, more explained here. I will spare you the ViewModel implementations, the following should give a rough idea how we can now structure our screens which is following:

• We have a container called MyScreen.
• At the top is a date picker, it has 2 icon buttons: < and > to change dates to either the previous or next day.
• The rest of the screen is a component that shows the weekday based on the picked date. It pretends to load this weekday from somewhere, and it has a ViewModel scoped to the date that is picked.

@Composable
fun MyScreen() {
  val viewModel = viewModel<MyViewModel>()
  val selectedDate by viewModel.selectedDate.collectAsState()
  Column {
    DateSelector(
      selectedDate = selectedDate,
      onDateChanged = viewModel::onDateChanged
    )

    HorizontalDivider()

    AnimatedContent(
      targetState = selectedDate,
      transitionSpec = TransitionSpec,
    ) { date ->
      WithViewModelStoreOwner(date) {
        val viewModel = viewModel() {
          WeekDayForDateViewModel(date)
        }
        WeekDayForDate(viewModel)
      }
      // Also possible:
      // val viewModel = viewModel(rememberViewModelStoreOwner(date) {
      //   WeekDayForDateViewModel(date)
      // }
      // WeekDayForDate(viewModel)
    }
  }
}

/**
  * Note that this component is not aware of the special ViewModel handling.
  * This is because we simply wrap it in WithViewModelStoreOwner {}.
  * It is always good for children components to not be aware of where
  * they are positioned and how to manage their scoping.
  */
@Composable
fun WeekDayForDate(viewModel: WeekDayForDateViewModel) {
  val viewState by viewModel.viewState
  LoadingResultScreen(
    onRefresh = viewModel::refresh,
    loadingResult = viewState,
  ) { state, _ ->
    Text(state.weekDay)
  }
}

private val TransitionSpec: AnimatedContentTransitionScope<LocalDate>.() -> ContentTransform = {
  val direction = initialState.compareTo(targetState)
  (fadeIn() + slideIn(initialOffset = { IntOffset(it.width * direction, 0) })) togetherWith 
    (fadeOut() + slideOut(targetOffset = { IntOffset(-it.width * direction, 0) }))
}

The idea is simple: If you select a date, we slide-animate to the appropriate date screen and “load” some data from elsewhere and display it. The date screens within the AnimatedContent have their own ViewModel and these will be cleared once a Composable is sliding out of the view. Also, I put in some randomization to show the screens can do proper error handling. All the error handling etc. is handled by the inner ViewModel. The following GIF was taken:

I can assure you that every time you switch dates the old screen slides out and its ViewModel is cleared. You trust me right? To convince you that I am not lying: the following GIF shows a snack bar message every time a view model is cleared! (and yes, these are triggered by the ViewModel’s onCleared method and not faked).

And what about configuration changes? Here you will see that when we rotate, the original Composable will keep its state and not turn into a loading spinner again.

Last but not least: We can navigate away if we dump it all in a NavHost, navigate back and keep our state, even if we rotate the screen!

The verdict

Actually pretty okay

I found out that it is hard to detect whether a Composable is actually leaving the composition because of recreation or simply because it is no longer needed. By using the power of ViewModel and Lifecycle I am happy to have found a non-hacky way to achieve the desired result. It is important to realize the constraint that stopping the lifecycle and detaching the Composable doesn’t mean we should get rid of our data. We will only be sure that we can clear any Composable by checking if its lifecycle is in a Resumed state and check whether we are attached or not. This is the only thing that disturbs me, because that synchronization relies on the Composable reaching a Resumed state, which could take a while. However, other ways to solve this more precisely would require hooking into parent Composables somehow to keep a tighter grip, and that would defeat the purpose of just simply dropping this solution into place. This solution can be dropped anywhere, in a ViewPager, or even in a LazyColumn where every item will clear its ViewModel when it is scrolled out from the View.

Usage in production

I actually will use this piece of code in production. The reason is simple: I already have a bad solution in place (see the “First iteration” mentioned in this article 🙂 ), it cannot get a lot worse. I already potentially lose ViewModel state on configuration changes, so this is an improvement. Also, in the end everything is tied to the LocalViewModelStoreOwner that is already there. So whenever that one is cleared, it will clean up any mess left behind by my implementation. I am perfectly fine taking any responsibility if some uncovered edge case would decide to pop up later and cause issues.

Final words

Writing this article was quite a journey. Just because you can, doesn’t mean you should. I went in multiple directions, and in the end found a solution that doesn’t require hooking into any framework in a weird or hacky way. It builds on my understand of the inner workings of Android Lifecycle and manages to take advantage of that. Please let me know what you think about this wonderful way of scoping ViewModel to your Composable! Also, if you spot any mistakes, I would be very happy to know, as this all started as an experiment.

And as always: If you like what you see, put those digital hands together! Joost out.

This article is previously published on proandroiddev.com