@Composable
fun <T> rememberSaveable(
    vararg inputs: Any?,
    stateSaver: Saver<T, out Any>,
    key: String? = null,
    init: () -> MutableState<T>
): MutableState<T> = rememberSaveable(
    *inputs,
    saver = mutableStateSaver(stateSaver),
    key = key,
    init = init
)

@Composable
fun <T : Any> rememberSaveable(
    vararg inputs: Any?,
    saver: Saver<T, out Any> = autoSaver(),
    key: String? = null,
    init: () -> T
): T {
    // key is the one provided by the user or the one generated by the compose runtime
    val finalKey = if (!key.isNullOrEmpty()) {
        key
    } else {
        currentCompositeKeyHash.toString(MaxSupportedRadix)
    }
    @Suppress("UNCHECKED_CAST")
    (saver as Saver<T, Any>)

    val registry = LocalSaveableStateRegistry.current
    // value is restored using the registry or created via [init] lambda
    val value = remember(*inputs) {
        // TODO not restore when the input values changed (use hashKeys?) b/152014032
        val restored = registry?.consumeRestored(finalKey)?.let {
            saver.restore(it)
        }
        restored ?: init()
    }

    // save the latest passed saver object into a state object to be able to use it when we will
    // be saving the value. keeping value in mutableStateOf() allows us to properly handle
    // possible compose transactions cancellations
    val saverHolder = remember { mutableStateOf(saver) }
    saverHolder.value = saver

    // re-register if the registry or key has been changed
    if (registry != null) {
        DisposableEffect(registry, finalKey, value) {
            val valueProvider = {
                with(saverHolder.value) { SaverScope { registry.canBeSaved(it) }.save(value) }
            }
            registry.requireCanBeSaved(valueProvider())
            val entry = registry.registerProvider(finalKey, valueProvider)
            onDispose {
                entry.unregister()
            }
        }
    }
    return value
}

/**
 * CompositionLocal with a current [SaveableStateRegistry] instance.
 */
val LocalSaveableStateRegistry = staticCompositionLocalOf<SaveableStateRegistry?> { null }

@Composable
@OptIn(ExperimentalComposeUiApi::class)
internal fun ProvideAndroidCompositionLocals(
    owner: AndroidComposeView,
    content: @Composable () -> Unit
) {
    val view = owner
		...
    val viewTreeOwners = owner.viewTreeOwners ?: throw IllegalStateException(
        "Called when the ViewTreeOwnersAvailability is not yet in Available state"
    )

    val saveableStateRegistry = remember {
        DisposableSaveableStateRegistry(view, viewTreeOwners.savedStateRegistryOwner)
    }
    DisposableEffect(Unit) {
        onDispose {
            saveableStateRegistry.dispose()
        }
    }

    CompositionLocalProvider(
        ...
        LocalSaveableStateRegistry provides saveableStateRegistry,
				...
    ) {
        ....
    }
}

/**
 * Creates [DisposableSaveableStateRegistry] associated with these [view] and [owner].
 */
internal fun DisposableSaveableStateRegistry(
    view: View,
    owner: SavedStateRegistryOwner
): DisposableSaveableStateRegistry {
    // The view id of AbstractComposeView is used as a key for SavedStateRegistryOwner. If there
    // are multiple AbstractComposeViews in the same Activity/Fragment with the same id(or with
    // no id) this means only the first view will restore its state. There is also an internal
    // mechanism to provide such id not as an Int to avoid ids collisions via view's tag. This
    // api is currently internal to compose:ui, we will see in the future if we need to make a
    // new public api for that use case.
    val composeView = (view.parent as View)
    val idFromTag = composeView.getTag(R.id.compose_view_saveable_id_tag) as? String
    val id = idFromTag ?: composeView.id.toString()
    return DisposableSaveableStateRegistry(id, owner)
}

/**
 * Creates [DisposableSaveableStateRegistry] with the restored values using [SavedStateRegistry] and
 * saves the values when [SavedStateRegistry] performs save.
 *
 * To provide a namespace we require unique [id]. We can't use the default way of doing it when we
 * have unique id on [AbstractComposeView] because we dynamically create [AbstractComposeView]s and
 * there is no way to have a unique id given there are could be any number of
 * [AbstractComposeView]s inside the same Activity. If we use [View.generateViewId]
 * this id will not survive Activity recreation. But it is reasonable to ask our users to have an
 * unique id on [AbstractComposeView].
 */
internal fun DisposableSaveableStateRegistry(
    id: String,
    savedStateRegistryOwner: SavedStateRegistryOwner
): DisposableSaveableStateRegistry {
    val key = "${SaveableStateRegistry::class.java.simpleName}:$id"

    val androidxRegistry = savedStateRegistryOwner.savedStateRegistry
    val bundle = androidxRegistry.consumeRestoredStateForKey(key)
    val restored: Map<String, List<Any?>>? = bundle?.toMap()

    val saveableStateRegistry = SaveableStateRegistry(restored) {
        canBeSavedToBundle(it)
    }
    val registered = try {
        androidxRegistry.registerSavedStateProvider(key) {
            saveableStateRegistry.performSave().toBundle()
        }
        true
    } catch (ignore: IllegalArgumentException) {
        // this means there are two AndroidComposeViews composed into different parents with the
        // same view id. currently we will just not save/restore state for the second
        // AndroidComposeView.
        // TODO: we should verify our strategy for such cases and improve it. b/162397322
        false

    }
    return DisposableSaveableStateRegistry(saveableStateRegistry) {
        if (registered) {
            androidxRegistry.unregisterSavedStateProvider(key)
        }
    }
}

/**
 * [SaveableStateRegistry] which can be disposed using [dispose].
 */
internal class DisposableSaveableStateRegistry(
    saveableStateRegistry: SaveableStateRegistry,
    private val onDispose: () -> Unit
) : SaveableStateRegistry by saveableStateRegistry {

    fun dispose() {
        onDispose()
    }
}

interface SaveableStateRegistry {
    fun consumeRestored(key: String): Any?
    fun registerProvider(key: String, valueProvider: () -> Any?): Entry
    fun canBeSaved(value: Any): Boolean
    fun performSave(): Map<String, List<Any?>>
    interface Entry {
        fun unregister()
    }
}

/**
 * Creates [SaveableStateRegistry].
 *
 * @param restoredValues The map of the restored values
 * @param canBeSaved Function which returns true if the given value can be saved by the registry
 */
fun SaveableStateRegistry(
    restoredValues: Map<String, List<Any?>>?,
    canBeSaved: (Any) -> Boolean
): SaveableStateRegistry = SaveableStateRegistryImpl(restoredValues, canBeSaved)

private class SaveableStateRegistryImpl(
    restored: Map<String, List<Any?>>?,
    private val canBeSaved: (Any) -> Boolean
) : SaveableStateRegistry {

    private val restored: MutableMap<String, List<Any?>> =
        restored?.toMutableMap() ?: mutableMapOf()
    private val valueProviders = mutableMapOf<String, MutableList<() -> Any?>>()

    override fun canBeSaved(value: Any): Boolean = canBeSaved.invoke(value)

    override fun consumeRestored(key: String): Any? {
        val list = restored.remove(key)
        return if (list != null && list.isNotEmpty()) {
            if (list.size > 1) {
                restored[key] = list.subList(1, list.size)
            }
            list[0]
        } else {
            null
        }
    }

    override fun registerProvider(key: String, valueProvider: () -> Any?): Entry {
        require(key.isNotBlank()) { "Registered key is empty or blank" }
        @Suppress("UNCHECKED_CAST")
        valueProviders.getOrPut(key) { mutableListOf() }.add(valueProvider)
        return object : Entry {
            override fun unregister() {
                val list = valueProviders.remove(key)
                list?.remove(valueProvider)
                if (list != null && list.isNotEmpty()) {
                    // if there are other providers for this key return list back to the map
                    valueProviders[key] = list
                }
            }
        }
    }

    override fun performSave(): Map<String, List<Any?>> {
        val map = restored.toMutableMap()
        valueProviders.forEach { (key, list) ->
            if (list.size == 1) {
                val value = list[0].invoke()
                if (value != null) {
                    check(canBeSaved(value))
                    map[key] = arrayListOf<Any?>(value)
                }
            } else {
                // if we have multiple providers we should store null values as well to preserve
                // the order in which providers were registered. say there were two providers.
                // the first provider returned null(nothing to save) and the second one returned
                // "1". when we will be restoring the first provider would restore null (it is the
                // same as to have nothing to restore) and the second one restore "1".
                map[key] = List(list.size) { index ->
                    val value = list[index].invoke()
                    if (value != null) {
                        check(canBeSaved(value))
                    }
                    value
                }
            }
        }
        return map
    }
}