Fix bug when rememberSaveable with input is restored in the wrong order

We need to not restart DisposableEffect when the value changes. Instead we use rememberUpdatedState in order to update the value to be used in the captured lambda.

Fixes: 214396074
Test: new test in RememberSaveableTest
Change-Id: I153fc8a719135f7dd4d7ed5407f2553bed49ce68
diff --git a/compose/runtime/runtime-saveable/src/androidAndroidTest/kotlin/androidx/compose/runtime/saveable/RememberSaveableTest.kt b/compose/runtime/runtime-saveable/src/androidAndroidTest/kotlin/androidx/compose/runtime/saveable/RememberSaveableTest.kt
index 033b693..1d15b6c 100644
--- a/compose/runtime/runtime-saveable/src/androidAndroidTest/kotlin/androidx/compose/runtime/saveable/RememberSaveableTest.kt
+++ b/compose/runtime/runtime-saveable/src/androidAndroidTest/kotlin/androidx/compose/runtime/saveable/RememberSaveableTest.kt
@@ -410,6 +410,44 @@
             assertThat(composedValue).isEqualTo(1)
         }
     }
+
+    @Test
+    fun changingInputIsNotAffectingOrderOfRestoration() {
+        var counter = 0
+        var input by mutableStateOf(0)
+        var withInput: Int? = null
+        var withoutInput: String? = null
+
+        restorationTester.setContent {
+            withInput = rememberSaveable(input) { counter++ }
+            withoutInput = rememberSaveable { (counter++).toString() }
+        }
+
+        rule.runOnIdle {
+            assertThat(withInput).isNotNull()
+            withInput = null
+            input++
+        }
+
+        var expectedWithInput: Int? = null
+        var expectedWithoutInput: String? = null
+
+        rule.runOnIdle {
+            assertThat(withInput).isNotNull()
+            assertThat(withoutInput).isNotNull()
+            expectedWithInput = withInput
+            expectedWithoutInput = withoutInput
+            withInput = null
+            withoutInput = null
+        }
+
+        restorationTester.emulateSavedInstanceStateRestore()
+
+        rule.runOnIdle {
+            assertThat(withInput).isEqualTo(expectedWithInput)
+            assertThat(withoutInput).isEqualTo(expectedWithoutInput)
+        }
+    }
 }
 
 @Composable
diff --git a/compose/runtime/runtime-saveable/src/commonMain/kotlin/androidx/compose/runtime/saveable/RememberSaveable.kt b/compose/runtime/runtime-saveable/src/commonMain/kotlin/androidx/compose/runtime/saveable/RememberSaveable.kt
index 44c0b1f..68f4cee 100644
--- a/compose/runtime/runtime-saveable/src/commonMain/kotlin/androidx/compose/runtime/saveable/RememberSaveable.kt
+++ b/compose/runtime/runtime-saveable/src/commonMain/kotlin/androidx/compose/runtime/saveable/RememberSaveable.kt
@@ -25,6 +25,7 @@
 import androidx.compose.runtime.neverEqualPolicy
 import androidx.compose.runtime.referentialEqualityPolicy
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.snapshots.SnapshotMutableState
 import androidx.compose.runtime.structuralEqualityPolicy
 
@@ -87,17 +88,19 @@
         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) {
+        // we want to use the latest instances of saver and value in the valueProvider lambda
+        // without restarting DisposableEffect as it would cause re-registering the provider in
+        // the different order. so we use rememberUpdatedState.
+        val saverState = rememberUpdatedState(saver)
+        val valueState = rememberUpdatedState(value)
+
+        DisposableEffect(registry, finalKey) {
             val valueProvider = {
-                with(saverHolder.value) { SaverScope { registry.canBeSaved(it) }.save(value) }
+                with(saverState.value) {
+                    SaverScope { registry.canBeSaved(it) }.save(valueState.value)
+                }
             }
             registry.requireCanBeSaved(valueProvider())
             val entry = registry.registerProvider(finalKey, valueProvider)