Prototype interoperability layer between `SavedStateHandle` and the Compose saveable API

SaveableSavedStateHandle.kt

/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Example usage
 */
@OptIn(ExperimentalPagerApi::class)
class PagerViewModel(savedStateHandle: SavedStateHandle) : ViewModel() {
    val pagerState: PagerState = savedStateHandle.saveable(
        key = "pagerState",
        saver = PagerState.Saver
    ) {
        PagerState(
            pageCount = 10,
            offscreenLimit = 2,
        )
    }
}

/**
 * A basic interop between [SavedStateHandle] and [Saver], so the latter can be used to save
 * state holders into the [SavedStateHandle].
 *
 * This implementation is based on [rememberSaveable], [SaveableStateRegistry] and
 * [DisposableSaveableStateRegistry], with some simplifications since there will be exactly one
 * state provider storing exactly one value.
 */
fun <T : Any> SavedStateHandle.saveable(
    key: String,
    saver: Saver<T, out Any> = autoSaver(),
    init: () -> T,
): T {
    val value = get<Bundle?>(key)?.let {
        @Suppress("UNCHECKED_CAST")
        (saver as Saver<T, Any>).restore(it.getParcelableArrayList<Parcelable>("state")!![0])
    } ?: init()

    setSavedStateProvider(key) {
        val saved = with(saver) { SaverScope(::canBeSavedToBundle).save(value) }
        Bundle().apply {
            @Suppress("UNCHECKED_CAST")
            putParcelableArrayList("state", arrayListOf(saved) as ArrayList<out Parcelable>)
        }
    }

    return value
}

/**
 * A basic interop between [SavedStateHandle] and [Saver], so the latter can be used to save
 * state holders into the [SavedStateHandle].
 *
 * This implementation is based on [rememberSaveable], [SaveableStateRegistry] and
 * [DisposableSaveableStateRegistry], with some simplifications since there will be exactly one
 * state provider storing exactly one value.
 *
 * Use this overload if you remember a mutable state with a type which can't be stored in the
 * Bundle so you have to provide a custom saver object.
 */
fun <T> SavedStateHandle.saveable(
    key: String,
    stateSaver: Saver<T, out Any>,
    init: () -> MutableState<T>
): MutableState<T> = saveable(
    saver = mutableStateSaver(stateSaver),
    key = key,
    init = init
)

/**
 * Copied from RememberSaveable.kt
 */
@Suppress("UNCHECKED_CAST")
private fun <T> mutableStateSaver(inner: Saver<T, out Any>) = with(inner as Saver<T, Any>) {
    Saver<MutableState<T>, MutableState<Any?>>(
        save = { state ->
            require(state is SnapshotMutableState<T>) {
                "If you use a custom MutableState implementation you have to write a custom " +
                    "Saver and pass it as a saver param to saveable()"
            }
            mutableStateOf(save(state.value), state.policy as SnapshotMutationPolicy<Any?>)
        },
        restore = @Suppress("UNCHECKED_CAST") {
            require(it is SnapshotMutableState<Any?>)
            mutableStateOf(
                if (it.value != null) restore(it.value!!) else null,
                it.policy as SnapshotMutationPolicy<T?>
            ) as MutableState<T>
        }
    )
}

/**
 * Checks that [value] can be stored inside [Bundle].
 *
 * Copied from DisposableSaveableStateRegistry.android.kt
 */
private fun canBeSavedToBundle(value: Any): Boolean {
    // SnapshotMutableStateImpl is Parcelable, but we do extra checks
    if (value is SnapshotMutableState<*>) {
        if (value.policy === neverEqualPolicy<Any?>() ||
            value.policy === structuralEqualityPolicy<Any?>() ||
            value.policy === referentialEqualityPolicy<Any?>()
        ) {
            val stateValue = value.value
            return if (stateValue == null) true else canBeSavedToBundle(stateValue)
        } else {
            return false
        }
    }
    for (cl in AcceptableClasses) {
        if (cl.isInstance(value)) {
            return true
        }
    }
    return false
}

/**
 * Contains Classes which can be stored inside [Bundle].
 *
 * Some of the classes are not added separately because:
 *
 * This classes implement Serializable:
 * - Arrays (DoubleArray, BooleanArray, IntArray, LongArray, ByteArray, FloatArray, ShortArray,
 * CharArray, Array<Parcelable, Array<String>)
 * - ArrayList
 * - Primitives (Boolean, Int, Long, Double, Float, Byte, Short, Char) will be boxed when casted
 * to Any, and all the boxed classes implements Serializable.
 * This class implements Parcelable:
 * - Bundle
 *
 * Note: it is simplified copy of the array from SavedStateHandle (lifecycle-viewmodel-savedstate).
 *
 * Copied from DisposableSaveableStateRegistry.android.kt
 */
private val AcceptableClasses = arrayOf(
    Serializable::class.java,
    Parcelable::class.java,
    String::class.java,
    SparseArray::class.java,
    Binder::class.java,
    Size::class.java,
    SizeF::class.java
)