[flexiglass] Introduce rememberSession {} API

This API allows for composables to associate a `remember` with a session
object, so that the state is maintained even when the composable exits
the composition.

For flexiglass, this is useful because it allows state to be preserved
across scenes; this allows for notification scroll position to be
preserved after going to-then-from Bouncer, for example.

Flag: ACONFIG com.android.systemui.scene_container DEVELOPMENT
Test: (with follow-up CL)
  1. Open Shade
  2. Scroll notifications
  3. Go to Bouncer
  4. Return to Shade, observe scroll state preserved

Change-Id: I37ed03d0dd95d6e01d064a693a46f73cebb2ff0b

2 files changed, 314 insertions, 0 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/shared/SessionStorage.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/shared/SessionStorage.kt
new file mode 100644
index 000000000000..dc5891915bfc
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/shared/SessionStorage.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 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
+ *
+ *      http://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.
+ */
+
+package com.android.systemui.scene.session.shared
+
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+
+/** Data store for [Session][com.android.systemui.scene.session.ui.composable.Session]. */
+class SessionStorage {
+    private var _storage by mutableStateOf(hashMapOf<String, StorageEntry>())
+
+    /**
+     * Data store containing all state retained for invocations of
+     * [rememberSession][com.android.systemui.scene.session.ui.composable.Session.rememberSession]
+     */
+    val storage: MutableMap<String, StorageEntry>
+        get() = _storage
+
+    /**
+     * Storage for an individual invocation of
+     * [rememberSession][com.android.systemui.scene.session.ui.composable.Session.rememberSession]
+     */
+    class StorageEntry(val keys: Array<out Any?>, var stored: Any?)
+
+    /** Clears the data store; any downstream usage within `@Composable`s will be recomposed. */
+    fun clear() {
+        _storage = hashMapOf()
+    }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/ui/composable/Session.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/ui/composable/Session.kt
new file mode 100644
index 000000000000..924aa540aa7f
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/session/ui/composable/Session.kt
@@ -0,0 +1,270 @@
+/*
+ * Copyright (C) 2024 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
+ *
+ *      http://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.
+ */
+
+package com.android.systemui.scene.session.ui.composable
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.SideEffect
+import androidx.compose.runtime.currentCompositeKeyHash
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.saveable.Saver
+import androidx.compose.runtime.saveable.SaverScope
+import androidx.compose.runtime.saveable.mapSaver
+import androidx.compose.runtime.saveable.rememberSaveable
+import androidx.compose.runtime.setValue
+import com.android.systemui.scene.session.shared.SessionStorage
+import com.android.systemui.util.kotlin.mapValuesNotNullTo
+
+/**
+ * An explicit storage for remembering composable state outside of the lifetime of a composition.
+ *
+ * Specifically, this allows easy conversion of standard
+ * [remember][androidx.compose.runtime.remember] invocations to ones that are preserved beyond the
+ * callsite's existence in the composition.
+ *
+ * ```kotlin
+ * @Composable
+ * fun Parent() {
+ *   val session = remember { Session() }
+ *   ...
+ *   if (someCondition) {
+ *     Child(session)
+ *   }
+ * }
+ *
+ * @Composable
+ * fun Child(session: Session) {
+ *   val state by session.rememberSession { mutableStateOf(0f) }
+ *   ...
+ * }
+ * ```
+ */
+interface Session {
+    /**
+     * Remember the value returned by [init] if all [inputs] are equal (`==`) to the values they had
+     * in the previous composition, otherwise produce and remember a new value by calling [init].
+     *
+     * @param inputs A set of inputs such that, when any of them have changed, will cause the state
+     *   to reset and [init] to be rerun
+     * @param key An optional key to be used as a key for the saved value. If `null`, we use the one
+     *   automatically generated by the Compose runtime which is unique for the every exact code
+     *   location in the composition tree
+     * @param init A factory function to create the initial value of this state
+     * @see androidx.compose.runtime.remember
+     */
+    @Composable fun <T> rememberSession(key: String?, vararg inputs: Any?, init: () -> T): T
+}
+
+/** Returns a new [Session], optionally backed by the provided [SessionStorage]. */
+fun Session(storage: SessionStorage = SessionStorage()): Session = SessionImpl(storage)
+
+/**
+ * Remember the value returned by [init] if all [inputs] are equal (`==`) to the values they had in
+ * the previous composition, otherwise produce and remember a new value by calling [init].
+ *
+ * @param inputs A set of inputs such that, when any of them have changed, will cause the state to
+ *   reset and [init] to be rerun
+ * @param key An optional key to be used as a key for the saved value. If not provided we use the
+ *   one automatically generated by the Compose runtime which is unique for the every exact code
+ *   location in the composition tree
+ * @param init A factory function to create the initial value of this state
+ * @see androidx.compose.runtime.remember
+ */
+@Composable
+fun <T> Session.rememberSession(vararg inputs: Any?, key: String? = null, init: () -> T): T =
+    rememberSession(key, inputs, init = init)
+
+/**
+ * An explicit storage for remembering composable state outside of the lifetime of a composition.
+ *
+ * Specifically, this allows easy conversion of standard [rememberSession] invocations to ones that
+ * are preserved beyond the callsite's existence in the composition.
+ *
+ * ```kotlin
+ * @Composable
+ * fun Parent() {
+ *   val session = rememberSaveableSession()
+ *   ...
+ *   if (someCondition) {
+ *     Child(session)
+ *   }
+ * }
+ *
+ * @Composable
+ * fun Child(session: SaveableSession) {
+ *   val state by session.rememberSaveableSession { mutableStateOf(0f) }
+ *   ...
+ * }
+ * ```
+ */
+interface SaveableSession : Session {
+    /**
+     * Remember the value produced by [init].
+     *
+     * It behaves similarly to [rememberSession], but the stored value will survive the activity or
+     * process recreation using the saved instance state mechanism (for example it happens when the
+     * screen is rotated in the Android application).
+     *
+     * @param inputs A set of inputs such that, when any of them have changed, will cause the state
+     *   to reset and [init] to be rerun
+     * @param saver The [Saver] object which defines how the state is saved and restored.
+     * @param key An optional key to be used as a key for the saved value. If not provided we use
+     *   the automatically generated by the Compose runtime which is unique for the every exact code
+     *   location in the composition tree
+     * @param init A factory function to create the initial value of this state
+     * @see rememberSaveable
+     */
+    @Composable
+    fun <T : Any> rememberSaveableSession(
+        vararg inputs: Any?,
+        saver: Saver<T, out Any>,
+        key: String?,
+        init: () -> T,
+    ): T
+}
+
+/**
+ * Returns a new [SaveableSession] that is preserved across configuration changes.
+ *
+ * @param inputs A set of inputs such that, when any of them have changed, will cause the state to
+ *   reset.
+ * @param key An optional key to be used as a key for the saved value. If not provided we use the
+ *   automatically generated by the Compose runtime which is unique for the every exact code
+ *   location in the composition tree.
+ */
+@Composable
+fun rememberSaveableSession(
+    vararg inputs: Any?,
+    key: String? = null,
+): SaveableSession =
+    rememberSaveable(inputs, SaveableSessionImpl.SessionSaver, key) { SaveableSessionImpl() }
+
+private class SessionImpl(
+    private val storage: SessionStorage = SessionStorage(),
+) : Session {
+    @Composable
+    override fun <T> rememberSession(key: String?, vararg inputs: Any?, init: () -> T): T {
+        val storage = storage.storage
+        val compositeKey = currentCompositeKeyHash
+        // key is the one provided by the user or the one generated by the compose runtime
+        val finalKey =
+            if (!key.isNullOrEmpty()) {
+                key
+            } else {
+                compositeKey.toString(MAX_SUPPORTED_RADIX)
+            }
+        if (finalKey !in storage) {
+            val value = init()
+            SideEffect { storage[finalKey] = SessionStorage.StorageEntry(inputs, value) }
+            return value
+        }
+        val entry = storage[finalKey]!!
+        if (!inputs.contentEquals(entry.keys)) {
+            val value = init()
+            SideEffect { entry.stored = value }
+            return value
+        }
+        @Suppress("UNCHECKED_CAST") return entry.stored as T
+    }
+}
+
+private class SaveableSessionImpl(
+    saveableStorage: MutableMap<String, StorageEntry> = mutableMapOf(),
+    sessionStorage: SessionStorage = SessionStorage(),
+) : SaveableSession, Session by Session(sessionStorage) {
+
+    var saveableStorage: MutableMap<String, StorageEntry> by mutableStateOf(saveableStorage)
+
+    @Composable
+    override fun <T : Any> rememberSaveableSession(
+        vararg inputs: Any?,
+        saver: Saver<T, out Any>,
+        key: String?,
+        init: () -> T,
+    ): T {
+        val compositeKey = currentCompositeKeyHash
+        // key is the one provided by the user or the one generated by the compose runtime
+        val finalKey =
+            if (!key.isNullOrEmpty()) {
+                key
+            } else {
+                compositeKey.toString(MAX_SUPPORTED_RADIX)
+            }
+
+        @Suppress("UNCHECKED_CAST") (saver as Saver<T, Any>)
+
+        if (finalKey !in saveableStorage) {
+            val value = init()
+            SideEffect { saveableStorage[finalKey] = StorageEntry.Restored(inputs, value, saver) }
+            return value
+        }
+        when (val entry = saveableStorage[finalKey]!!) {
+            is StorageEntry.Unrestored -> {
+                val value = saver.restore(entry.unrestored) ?: init()
+                SideEffect {
+                    saveableStorage[finalKey] = StorageEntry.Restored(inputs, value, saver)
+                }
+                return value
+            }
+            is StorageEntry.Restored<*> -> {
+                if (!inputs.contentEquals(entry.inputs)) {
+                    val value = init()
+                    SideEffect {
+                        saveableStorage[finalKey] = StorageEntry.Restored(inputs, value, saver)
+                    }
+                    return value
+                }
+                @Suppress("UNCHECKED_CAST") return entry.stored as T
+            }
+        }
+    }
+
+    sealed class StorageEntry {
+        class Unrestored(val unrestored: Any) : StorageEntry()
+
+        class Restored<T>(val inputs: Array<out Any?>, var stored: T, val saver: Saver<T, Any>) :
+            StorageEntry() {
+            fun SaverScope.saveEntry() {
+                with(saver) { stored?.let { save(it) } }
+            }
+        }
+    }
+
+    object SessionSaver :
+        Saver<SaveableSessionImpl, Any> by mapSaver(
+            save = { sessionScope: SaveableSessionImpl ->
+                sessionScope.saveableStorage.mapValues { (k, v) ->
+                    when (v) {
+                        is StorageEntry.Unrestored -> v.unrestored
+                        is StorageEntry.Restored<*> -> {
+                            with(v) { saveEntry() }
+                        }
+                    }
+                }
+            },
+            restore = { savedMap: Map<String, Any?> ->
+                SaveableSessionImpl(
+                    saveableStorage =
+                        savedMap.mapValuesNotNullTo(mutableMapOf()) { (k, v) ->
+                            v?.let { StorageEntry.Unrestored(v) }
+                        }
+                )
+            }
+        )
+}
+
+private const val MAX_SUPPORTED_RADIX = 36