Update Activity reference (from Compose) on configuration changes (Activity recreation)

ActivityLifecycleComposable.kt

import android.app.Activity
import android.content.Context
import android.content.ContextWrapper
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.core.app.ComponentActivity
import androidx.lifecycle.ViewModelStoreOwner


@Composable
fun ActivityLifecycle() {
    // Access container Activity
    val context = LocalContext.current
    val compositionActivity = context.findActivity()

    // Observe state of configuration changes when disposing
    LaunchedEffect(compositionActivity) {
        compositionActivity.application.registerActivityLifecycleCallbacks(object : android.app.Application.ActivityLifecycleCallbacks {
            // Reference to Activity between creation and destruction to be updated on new creation
            var creationActivity: Activity? = compositionActivity
            // When an Activity is recreatedk, the ViewModelStore stays the same, even if the Activity object changes
            var creationActivityViewModelStore: ViewModelStore? = null
          
            override fun onActivityCreated(activity: Activity, savedInstanceState: android.os.Bundle?) {
                val sameViewModelStore = creationActivityViewModelStore == (activity as? ViewModelStoreOwner)?.viewModelStore // Detect new Activity instance
                if (sameViewModelStore) { // We are on the same Activity the ViewModel loaded on
                    creationActivity = activity // Update it
                }
            }

            override fun onActivityStarted(activity: Activity) {
                // no-op
            }

            override fun onActivityResumed(activity: Activity) {
                // no-op
            }

            override fun onActivityPaused(activity: Activity) {
                // no-op
            }

            override fun onActivityStopped(activity: Activity) {
                // no-op
            }

            override fun onActivitySaveInstanceState(activity: Activity, outState: android.os.Bundle) {
                // no-op
            }

            override fun onActivityDestroyed(activity: Activity) {
                if (creationActivity != activity) return // Ignore other activities
                if (activity.isChangingConfigurations) { // Activity is being recreated, update the ViewModelStoreOwner to find new version of this Activity
                    creationActivityViewModelStore = (activity as? ViewModelStoreOwner)?.viewModelStore
                } else { // Activity is being destroyed forever, therefore unregister the callback
                    activity.application.unregisterActivityLifecycleCallbacks(this)
                }
            }
        })
    }
}

internal fun Context.findActivity(): Activity {
    var ctx = this
    while (ctx is ContextWrapper) {
        if (ctx is Activity) {
            return ctx
        }
        ctx = ctx.baseContext
    }
    throw IllegalStateException(
        "Expected an activity context for detecting configuration changes for a NavBackStackEntry but instead found: $ctx"
    )
}