`rememberSaveable` but in-memory instead of on-disk (like `viewModel()`)

In Jetpack Compose, is there something like rememberSaveable but does not save to disk so it works like viewModel()? Currently I'm having this:

class MyViewModel : ViewModel() {
	var initialized: Boolean = false

	var needsToBeInitialized = mutableStateOf<Boolean?>(null)
}

@Composable
func HelloWorld() {
	val model: MyViewModel = viewModel()

	// this needs to survive device rotation
	var needsToBeInitialized by model.needsToBeInitialized

	if (!model.initialized) {
		model.needsToBeInitialized.value = getBooleanOrNullSomehow()

		model.initialized = true
	}
}

With rememberSaveable it's just a nice one-liner, so I was thinking if it was possible to achieve something like viewModel() in just a one line as well, so it stores data in memory instead.

You can achieve what you are describing by creating a custom remember function that is using a "generic" ViewModel class that keeps any state assigned to it in memory.

The "generic" ViewModel and the custom remember function (I called it rememberInMemory)

import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisallowComposableCalls
import androidx.compose.runtime.remember
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewmodel.compose.viewModel

/**
 * A generic [ViewModel] class that will keep any state assigned to it in memory.
 */
class StateMapViewModel : ViewModel() {
    val states: MutableMap<String, Pair<Array<out Any?>, Any?>> = mutableMapOf()
}

/**
 * Remember the value produced by [init].
 *
 * It behaves similarly to [remember], but the stored value will survive configuration changes
 * (for example when the screen is rotated in the Android application).
 * The value will not survive process recreation/death.
 *
 * @param dataKey A string to be used as a key for the saved value.
 * @param inputs A set of inputs such that, when any of them have changed, will cause the state to
 * reset and [init] to be rerun
 * @param init A factory function to create the initial value of this state
 * @return The state that is stored in memory.
 */
@Composable
inline fun <reified T> rememberInMemory(
    dataKey: String,
    vararg inputs: Any?,
    crossinline init: @DisallowComposableCalls () -> T,
): T {
    val vm: StateMapViewModel = viewModel()
    return remember(*inputs) {
        val value = vm.states[dataKey]
            ?.takeIf { it.first.contentEquals(inputs) }
            ?: Pair(inputs, init()).also {
                vm.states[dataKey] = it
            }
        value.second as T
    }
}

When calling rememberInMemory a string has to be provided as the dataKey argument. This key will be used to get the current stored value. If the given dataKey does not have a stored value yet, the init lambda will be called to obtain the initial value, which will be stored and returned. The inputs are optional variable length arguments and when any of them have changed it will cause the init lambda to be called again to obtain the new value, which will be stored and returned.

The scope of the data being kept in memory will be the same as the default ViewModel scope, which means this function will work the same as a ViewModel approach would.

When using navigation the StateMapViewModel will be scoped to each NavBackStackEntry. This means that every time when navigating to a new destination, a new StateMapViewModel instance will be created and when navigating back or up, the existing StateMapViewModel instance from the previous NavBackStackEntry will be used. So even with navigation this approach will work the same as would using a ViewModel directly.

If you are using any dependency injection frameworks replace the viewModel() call inside the rememberInMemory function with the correct viewModel()/getViewModel() call from the framework.

<hr>

Below is a demo app showing how this works with data scoped to 3 different navigation destinations and some data scoped to their parent.

Since the demo is using Compose Navigation, you have to add a dependency to your app/build.gradle to try the demo.

implementation "androidx.navigation:navigation-compose:2.5.3"

MainActivity.kt

import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent

class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)

        setContent {
            Demo()
        }
    }
}

Demo.kt

import androidx.compose.foundation.background
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.ColumnScope
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.KeyboardArrowDown
import androidx.compose.material.icons.filled.KeyboardArrowUp
import androidx.compose.material3.Button
import androidx.compose.material3.Divider
import androidx.compose.material3.Icon
import androidx.compose.material3.IconButton
import androidx.compose.material3.IconButtonDefaults
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Surface
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.MutableState
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.tooling.preview.Preview
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.dp
import androidx.navigation.compose.NavHost
import androidx.navigation.compose.composable
import androidx.navigation.compose.currentBackStackEntryAsState
import androidx.navigation.compose.rememberNavController
import kotlinx.coroutines.delay

@Preview(showBackground = true)
@Composable
fun Demo() {
    @Composable
    fun Surface(content: @Composable ColumnScope.() -> Unit) {
        Surface(modifier = Modifier.background(MaterialTheme.colorScheme.surface)) {
            Column(modifier = Modifier.padding(8.dp)) {
                content()
            }
        }
    }

    @Composable
    fun Counter(text: String, counterState: MutableState<Int>) {
        var counter by counterState
        Row(
            horizontalArrangement = Arrangement.spacedBy(4.dp),
            verticalAlignment = Alignment.CenterVertically
        ) {
            val modifier = Modifier
                .padding(2.dp)
                .background(
                    color = Color.Unspecified,
                    shape = RoundedCornerShape(100)
                )
            val colors = IconButtonDefaults.iconButtonColors(
                containerColor = MaterialTheme.colorScheme.primary,
                contentColor = MaterialTheme.colorScheme.onPrimary
            )
            IconButton(onClick = { counter-- }, modifier, colors = colors) {
                Icon(Icons.Filled.KeyboardArrowDown, contentDescription = "Decrement")
            }
            IconButton(onClick = { counter++ }, modifier, colors = colors) {
                Icon(Icons.Filled.KeyboardArrowUp, contentDescription = "Increment")
            }
            Text(text = "$text: $counter")
        }
    }

    val navController = rememberNavController()

    var countdown by rememberInMemory("Countdown") { mutableStateOf(0) }
    LaunchedEffect(Unit) {
        while (true) {
            countdown = 9
            while (countdown > 0) {
                delay(1000)
                countdown--
            }
            delay(1000)
        }
    }

    val parentCounter = rememberInMemory("Counter", countdown > 0) { mutableStateOf(0) }

    var previousCounter by rememberInMemory("Previous Counter") {
        mutableStateOf(parentCounter)
    }

    Column {
        Surface {
            Text("Countdown: $countdown seconds until parent counter resets")

            Counter("Parent counter", parentCounter)

            Row(horizontalArrangement = Arrangement.spacedBy(4.dp)) {
                val route = navController.currentBackStackEntryAsState().value?.destination?.route

                Button(onClick = { navController.navigateUp() }, enabled = route != "countersA") {
                    Text(text = "Go back")
                }

                when (route) {
                    "countersA" -> Button(onClick = { navController.navigate("countersB") }) {
                        Text(text = "Show Counters B")
                    }
                    "countersB" -> Button(onClick = { navController.navigate("countersC") }) {
                        Text(text = "Show Counters C")
                    }
                }
            }
        }

        Divider(thickness = Dp.Hairline)

        NavHost(navController = navController, startDestination = "countersA") {
            composable("countersA") {
                Surface {
                    val counterA = rememberInMemory("Counter") { mutableStateOf(-1) }
                    LaunchedEffect(Unit) {
                        previousCounter = counterA
                    }

                    Counter("Counter A", counterA)

                    Counter("Parent Counter", parentCounter)
                }
            }
            composable("countersB") {
                Surface {
                    val counterB = rememberInMemory("Counter") { mutableStateOf(-2) }
                    LaunchedEffect(Unit) {
                        previousCounter = counterB
                    }

                    Counter("Counter B", counterB)

                    Counter("Previous Counter", rememberInMemory("Previous Counter") { previousCounter })

                    Counter("Parent Counter", parentCounter)
                }
            }
            composable("countersC") {
                Surface {
                    val counterC = rememberInMemory("Counter") { mutableStateOf(-3) }
                    LaunchedEffect(Unit) {
                        previousCounter = counterC
                    }

                    Counter("Counter C", counterC)

                    Counter("Previous Counter", rememberInMemory("Previous Counter") { previousCounter })

                    Counter("Parent Counter", parentCounter)
                }
            }
        }
    }
}

I have improved @Ma3x's answer. This uses DisposableEffect as in rememberSaveable, so dataKey should be unnecessary now.

package com.example.stackoverflow

import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisallowComposableCalls
import androidx.compose.runtime.DisposableEffect
import androidx.compose.runtime.currentCompositeKeyHash
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberUpdatedState
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewmodel.compose.viewModel

/**
 * A generic [ViewModel] class that will keep any state assigned to it in memory.
 */
class StateMapViewModel : ViewModel() {
    val states: MutableMap<Int, ArrayDeque<Any>> = mutableMapOf()
}

/**
 * Remember the value produced by [init].
 *
 * It behaves similarly to [remember], but the stored value will survive configuration changes
 * (for example when the screen is rotated in the Android application).
 * The value will not survive process recreation.
 *
 * @param inputs A set of inputs such that, when any of them have changed, will cause the state to
 * reset and [init] to be rerun
 * @param init A factory function to create the initial value of this state
 */
@Composable
inline fun <reified T : Any> rememberInMemory(
    vararg inputs: Any?,
    crossinline init: @DisallowComposableCalls () -> T,
): T {
    val vm: StateMapViewModel = viewModel()

    val key = currentCompositeKeyHash

    val value = remember(*inputs) {
        val states = vm.states[key] ?: ArrayDeque<Any>().also { vm.states[key] = it }

        states.removeFirstOrNull() as T? ?: init()
    }

    val valueState = rememberUpdatedState(value)

    DisposableEffect(key) {
        onDispose {
            vm.states[key]?.addFirst(valueState.value)
        }
    }

    return value
}