Decompose retained component creation from Composable fun

RetainedComposeComponent.kt

import android.app.Activity
import android.os.Bundle
import androidx.activity.BackEventCompat
import androidx.activity.OnBackPressedCallback
import androidx.activity.OnBackPressedDispatcher
import androidx.activity.compose.LocalOnBackPressedDispatcherOwner
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.platform.LocalSavedStateRegistryOwner
import androidx.lifecycle.LifecycleOwner
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import androidx.savedstate.SavedStateRegistry
import androidx.savedstate.SavedStateRegistryOwner
import com.arkivanov.decompose.ComponentContext
import com.arkivanov.decompose.DefaultComponentContext
import com.arkivanov.essenty.backhandler.BackHandler
import com.arkivanov.essenty.instancekeeper.InstanceKeeper
import com.arkivanov.essenty.instancekeeper.InstanceKeeperDispatcher
import com.arkivanov.essenty.instancekeeper.getOrCreate
import com.arkivanov.essenty.instancekeeper.instanceKeeper
import com.arkivanov.essenty.lifecycle.LifecycleRegistry
import com.arkivanov.essenty.lifecycle.create
import com.arkivanov.essenty.lifecycle.destroy
import com.arkivanov.essenty.lifecycle.essentyLifecycle
import com.arkivanov.essenty.lifecycle.pause
import com.arkivanov.essenty.lifecycle.resume
import com.arkivanov.essenty.lifecycle.start
import com.arkivanov.essenty.lifecycle.stop
import com.arkivanov.essenty.lifecycle.subscribe
import com.arkivanov.essenty.parcelable.ParcelableContainer
import com.arkivanov.essenty.statekeeper.StateKeeper
import com.arkivanov.essenty.statekeeper.StateKeeperDispatcher
import com.arkivanov.essenty.statekeeper.consume

@Composable
fun <T> retainedComponent(
    key: String = "RootRetainedComponent",
    handleBackButton: Boolean = true,
    factory: (ComponentContext) -> T,
): T {
    val context = LocalContext.current
    val lifecycleOwner = LocalLifecycleOwner.current
    val savedStateRegistryOwner = LocalSavedStateRegistryOwner.current
    val viewModelStoreOwner = LocalViewModelStoreOwner.current
    val onBackPressedDispatcher = LocalOnBackPressedDispatcherOwner.current?.onBackPressedDispatcher

    return remember(
        context,
        lifecycleOwner,
        savedStateRegistryOwner,
        viewModelStoreOwner,
        onBackPressedDispatcher
    ) {
        retainedComponent(
            key = key,
            lifecycleOwner = lifecycleOwner,
            savedStateRegistryOwner = savedStateRegistryOwner,
            viewModelStoreOwner = requireNotNull(viewModelStoreOwner),
            onBackPressedDispatcher = if (handleBackButton) onBackPressedDispatcher else null,
            isChangingConfigurations = { (context as? Activity)?.isChangingConfigurations ?: false },
            factory = factory,
        )
    }
}

private fun <T> retainedComponent(
    key: String,
    lifecycleOwner: LifecycleOwner,
    savedStateRegistryOwner: SavedStateRegistryOwner,
    viewModelStoreOwner: ViewModelStoreOwner,
    onBackPressedDispatcher: OnBackPressedDispatcher?,
    isChangingConfigurations: () -> Boolean,
    factory: (ComponentContext) -> T,
): T {
    val lifecycle = lifecycleOwner.essentyLifecycle()
    val stateKeeper = SafeStateKeeper(savedStateRegistryOwner.savedStateRegistry)
    val instanceKeeper = viewModelStoreOwner.instanceKeeper()

    check(!stateKeeper.isRegistered(key = key)) { "Another retained component is already registered with the key: $key" }

    val holder = instanceKeeper.getOrCreate(key = key) {
        RetainedComponentHolder(
            savedState = stateKeeper.consume(key = key),
            factory = factory,
        )
    }

    lifecycle.subscribe(
        onCreate = { holder.lifecycle.create() },
        onStart = { holder.lifecycle.start() },
        onResume = { holder.lifecycle.resume() },
        onPause = {
            if (!isChangingConfigurations()) {
                holder.lifecycle.pause()
            }
        },
        onStop = {
            if (!isChangingConfigurations()) {
                holder.lifecycle.stop()
            }
        },
        onDestroy = {
            holder.onBackEnabledChangedListener = null
            if (!isChangingConfigurations()) {
                holder.lifecycle.destroy()
            }
        },
    )

    stateKeeper.register(key = key) { holder.stateKeeper.save() }

    if (onBackPressedDispatcher != null && holder.onBackEnabledChangedListener == null) {
        val onBackPressedCallback = DelegateOnBackPressedCallback(holder.onBackPressedDispatcher)
        holder.onBackEnabledChangedListener = { onBackPressedCallback.isEnabled = it }
        onBackPressedDispatcher.addCallback(lifecycleOwner, onBackPressedCallback)
    }

    return holder.component
}

private class DelegateOnBackPressedCallback(
    private val dispatcher: OnBackPressedDispatcher,
) : OnBackPressedCallback(enabled = dispatcher.hasEnabledCallbacks()) {
    override fun handleOnBackPressed() {
        dispatcher.onBackPressed()
    }

    override fun handleOnBackStarted(backEvent: BackEventCompat) {
        dispatcher.dispatchOnBackStarted(backEvent)
    }

    override fun handleOnBackProgressed(backEvent: BackEventCompat) {
        dispatcher.dispatchOnBackProgressed(backEvent)
    }

    override fun handleOnBackCancelled() {
        dispatcher.dispatchOnBackCancelled()
    }
}

private class RetainedComponentHolder<out T>(
    savedState: ParcelableContainer?,
    factory: (ComponentContext) -> T,
) : InstanceKeeper.Instance {
    val lifecycle: LifecycleRegistry = LifecycleRegistry()
    val stateKeeper: StateKeeperDispatcher = StateKeeperDispatcher(savedState = savedState)
    private val instanceKeeper: InstanceKeeperDispatcher = InstanceKeeperDispatcher()

    var onBackEnabledChangedListener: ((Boolean) -> Unit)? = null

    val onBackPressedDispatcher: OnBackPressedDispatcher =
        OnBackPressedDispatcher(
            fallbackOnBackPressed = null,
            onHasEnabledCallbacksChanged = {
                onBackEnabledChangedListener?.invoke(it)
            },
        )

    val component: T = factory(
        DefaultComponentContext(
            lifecycle = lifecycle,
            stateKeeper = stateKeeper,
            instanceKeeper = instanceKeeper,
            backHandler = BackHandler(onBackPressedDispatcher),
        )
    )

    override fun onDestroy() {
        instanceKeeper.destroy()
    }
}

@Suppress("FunctionName") // Factory function
private fun SafeStateKeeper(
    savedStateRegistry: SavedStateRegistry,
    isSavingAllowed: () -> Boolean = { true }
): StateKeeper {
    val dispatcher = StateKeeperDispatcher(savedStateRegistry.consumeRestoredStateForKey(KEY_STATE)?.getParcelable(KEY_STATE))

    if (savedStateRegistry.getSavedStateProvider(KEY_STATE) != null) {
        return dispatcher
    }

    savedStateRegistry.registerSavedStateProvider(KEY_STATE) {
        Bundle().apply {
            if (isSavingAllowed()) {
                putParcelable(KEY_STATE, dispatcher.save())
            }
        }
    }

    return dispatcher
}

private const val KEY_STATE = "STATE_KEEPER_STATE"