Molecule and retain

[argenkiwi]

This morning I learned that the retain API has been released as part of the Jetpack Compose December '25 release. I have been using ViewModels solely to retain a reference to my molecules and I am exploring replacing them with retain. I put together a simple composable function called retainMolecule:

@Composable
fun <T> retainMolecule(vararg keys: Any?, body: @Composable () -> T): StateFlow<T> {
    val context = AndroidUiDispatcher.Main
    val scope = retain { CoroutineScope(context) }
    return retain(*keys) {
        scope.launchMolecule(
            mode = RecompositionMode.ContextClock,
            context = context,
            body = body
        )
    }
}

Which I can use as follows in my views:

val stateFlow = retainMolecule { present() }
val state by stateFlow.collectAsStateWithLifecycle()

Is this approach good enough or may I be missing something?

[argenkiwi]

I ended up with something a bit simpler:

@Composable
context(scope: CoroutineScope)
fun <T> retainMolecule(
    vararg keys: Any?,
    mode: RecompositionMode = RecompositionMode.ContextClock,
    context: CoroutineContext = scope.coroutineContext,
    body: @Composable () -> T
) = retain(*keys) {
    scope.launchMolecule(mode, context, body =  body)
}

This way I can share the retained coroutine scope across views.

context(retain { CoroutineContext(AndroidUiDispatcher.Main) }) {
  val stateFlow = retainMolecule(someParameter) { present(someParameter) }
  val state by stateFlow.collectAsStateWithLifecycle()
  SomeView.Pane(state, ...)
  // More views or panes could go here, for example, when using nav3.
}