RememberRetained.kt

// A remember() that can survive recompositions and config changes but not process death.
@Composable
inline fun <T: Any> rememberRetained(
  key: String = currentCompositeKeyHash.toString(36), // Copied from rememberSaveable()
  crossinline init: @DisallowComposableCalls () -> T
): T {
  if (retainedStore[key] == null) {
    retainedStore[key] = init()
  }
  val activity = LocalContext.current.findActivity()
  remember {
    Cleaner(activity, key)
  }
  return retainedStore[key] as T
}

// This instance survives recomposition, but not Activity death.
// Should be safe to keep the Activity instance.
@PublishedApi
internal class Cleaner(private val activity: Activity, private val key: String): RememberObserver {
  override fun onAbandoned() {
    cleanUp()
  }

  override fun onForgotten() {
    if (!activity.isChangingConfigurations) {
      cleanUp()
    }
  }

  override fun onRemembered() {}

  private fun cleanUp() {
    retainedStore.remove(key)
  }
}

// Just a global variable as we don't need to survive process death
@PublishedApi
internal val retainedStore = mutableMapOf<String, Any>()

@PublishedApi
internal fun Context.findActivity(): Activity {
  var context = this
  while (context !is Activity) {
    if (context is ContextWrapper) {
      context = context.baseContext
    } else {
      error("Couldn't find Activity!")
    }
  }
  return context
}