Prevent pausable composition from dispatching redundant remembers

If a pausable composition is paused the composition, the state may
have changed to cause a remembered disposable effect to be replaced
before `apply()` is called. This can cause the disposable object
being remembered and forgotten in the same apply.

This change prevents the same apply from both remembering and
forgetting the same instance in the same apply.

This change also will ensure that paused compositions are no
longer recomposed outside of calls to resume. This change makes
it dangerous to call apply without first checking isComplete as
noted in the change method documentation.

Fixes: 404645679
Fixes: 407931790
Test: ./gradlew :compose:r:r:i-t:cC
Relnote: """Fix dispatching of remember observers in pausable
composition to avoid dispatching remembered/forgotten in the
same apply"""

Change-Id: I570b2436b95c7f8fb88a6f9824dbb9b8bccbc0fa
diff --git a/compose/runtime/runtime/integration-tests/src/androidInstrumentedTest/kotlin/androidx/compose/runtime/PausableCompositionInstrumentedTests.kt b/compose/runtime/runtime/integration-tests/src/androidInstrumentedTest/kotlin/androidx/compose/runtime/PausableCompositionInstrumentedTests.kt
new file mode 100644
index 0000000..aa25765
--- /dev/null
+++ b/compose/runtime/runtime/integration-tests/src/androidInstrumentedTest/kotlin/androidx/compose/runtime/PausableCompositionInstrumentedTests.kt
@@ -0,0 +1,119 @@
+/*
+ * Copyright 2025 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 androidx.compose.runtime
+
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.saveable.SaveableStateHolder
+import androidx.compose.runtime.saveable.rememberSaveableStateHolder
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.layout.SubcomposeLayout
+import androidx.compose.ui.layout.SubcomposeLayoutState
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.MediumTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@MediumTest
+@RunWith(AndroidJUnit4::class)
+class PausableCompositionInstrumentedTests {
+    @get:Rule val rule = createComposeRule()
+
+    @Test
+    fun changeTheKeyUsedInPrecomposition() {
+        val state = SubcomposeLayoutState()
+        var precompositionKey by mutableStateOf("1")
+        var addSlot by mutableStateOf(false)
+        lateinit var holder: SaveableStateHolder
+
+        rule.setContent {
+            holder = rememberSaveableStateHolder()
+            SubcomposeLayout(state) {
+                if (addSlot) {
+                    subcompose("measurepass") { holder.SaveableStateProvider("1") {} }
+                }
+                layout(10, 10) {}
+            }
+        }
+
+        val precomposition =
+            rule.runOnIdle {
+                val precomposition =
+                    state.createPausedPrecomposition("precomposition") {
+                        holder.SaveableStateProvider(precompositionKey) {}
+                    }
+
+                while (!precomposition.isComplete) {
+                    precomposition.resume { false }
+                }
+
+                precompositionKey = "2"
+
+                addSlot = true
+
+                precomposition
+            }
+
+        rule.runOnIdle {
+            while (!precomposition.isComplete) {
+                precomposition.resume { false }
+            }
+            precomposition.apply()
+        }
+    }
+
+    @Test
+    fun pausedPrecompositionIsNotRecomposedOnItsOwn() {
+        val state = SubcomposeLayoutState()
+        var key by mutableStateOf("1")
+        val composed = mutableListOf<String>()
+
+        rule.setContent { SubcomposeLayout(state) { layout(10, 10) {} } }
+
+        val precomposition =
+            rule.runOnIdle {
+                val precomposition = state.createPausedPrecomposition(Unit) { composed.add(key) }
+                while (!precomposition.isComplete) {
+                    precomposition.resume { false }
+                }
+
+                assertThat(composed).isEqualTo(listOf("1"))
+                composed.clear()
+
+                // recompose just composed composable (which wasn't yet applied)
+                key = "2"
+
+                precomposition
+            }
+
+        rule.runOnIdle {
+            // check that recomposition didn't happen on its own.
+            // we don't expect that for non applied compositions.
+            assertThat(composed).isEqualTo(emptyList<String>())
+
+            while (!precomposition.isComplete) {
+                precomposition.resume { false }
+            }
+            precomposition.apply()
+
+            // recomposition happened
+            assertThat(composed).isEqualTo(listOf("2"))
+        }
+    }
+}
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composition.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composition.kt
index a358a52..0b4b59c 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composition.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composition.kt
@@ -995,6 +995,16 @@
     override fun recompose(): Boolean =
         synchronized(lock) {
             drainPendingModificationsForCompositionLocked()
+            val pendingPausedComposition = pendingPausedComposition
+            if (pendingPausedComposition != null && !pendingPausedComposition.isRecomposing) {
+                // If the composition is pending do not recompose it now as the recomposition
+                // is in the control of the pausable composition and is supposed to happen when
+                // the resume is called. However, this may cause the pausable composition to go
+                // revert to an incomplete state. If isRecomposing is true then this is being
+                // called in resume()
+                pendingPausedComposition.markIncomplete()
+                return false
+            }
             guardChanges {
                 guardInvalidationsLocked { invalidations ->
                     val observer = observer()
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/PausableComposition.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/PausableComposition.kt
index 92e8d36..bf15fac 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/PausableComposition.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/PausableComposition.kt
@@ -64,7 +64,7 @@
     fun setPausableContent(content: @Composable () -> Unit): PausedComposition
 
     /**
-     * Set the content of a resuable composition. A [PausedComposition] that is currently paused. No
+     * Set the content of a reusable composition. A [PausedComposition] that is currently paused. No
      * composition is performed until [PausedComposition.resume] is called.
      * [PausedComposition.resume] should be called until [PausedComposition.isComplete] is `true`.
      * The composition should not be used until [PausedComposition.isComplete] is `true` and
@@ -100,7 +100,10 @@
     /**
      * Returns `true` when the [PausedComposition] is complete. [isComplete] matches the last value
      * returned from [resume]. Once a [PausedComposition] is [isComplete] the [apply] method should
-     * be called.
+     * be called. If the [apply] method is not called synchronously and immediately after [resume]
+     * returns `true` then this [isComplete] can return `false` as any state changes read by the
+     * paused composition while it is paused will cause the composition to require the paused
+     * composition to need to be resumed before it is used.
      */
     val isComplete: Boolean
 
@@ -127,6 +130,12 @@
     /**
      * Apply the composition. This is the last step of a paused composition and is required to be
      * called prior to the composition is usable.
+     *
+     * Calling [apply] should always be proceeded with a check of [isComplete] before it is called
+     * and potentially calling [resume] in a loop until [isComplete] returns `true`. This can happen
+     * if [resume] returned `true` but [apply] was not synchronously called immediately afterwords.
+     * Any state that was read that changed between when [resume] being called and [apply] being
+     * called may require the paused composition to be resumed before applied.
      */
     fun apply()
 
@@ -155,6 +164,7 @@
     Cancelled,
     InitialPending,
     RecomposePending,
+    Recomposing,
     ApplyPending,
     Applied,
 }
@@ -174,6 +184,8 @@
     internal val rememberManager =
         RememberEventDispatcher().apply { prepare(abandonSet, composer.errorContext) }
     internal val pausableApplier = RecordingApplier(applier.current)
+    internal val isRecomposing
+        get() = state == PausedCompositionState.Recomposing
 
     override val isComplete: Boolean
         get() = state >= PausedCompositionState.ApplyPending
@@ -193,9 +205,18 @@
                     if (invalidScopes.isEmpty()) markComplete()
                 }
                 PausedCompositionState.RecomposePending -> {
-                    invalidScopes = context.recomposePaused(composition, shouldPause, invalidScopes)
+                    state = PausedCompositionState.Recomposing
+                    try {
+                        invalidScopes =
+                            context.recomposePaused(composition, shouldPause, invalidScopes)
+                    } finally {
+                        state = PausedCompositionState.RecomposePending
+                    }
                     if (invalidScopes.isEmpty()) markComplete()
                 }
+                PausedCompositionState.Recomposing -> {
+                    composeRuntimeError("Recursive call to resume()")
+                }
                 PausedCompositionState.ApplyPending ->
                     error("Pausable composition is complete and apply() should be applied")
                 PausedCompositionState.Applied -> error("The paused composition has been applied")
@@ -215,7 +236,8 @@
         try {
             when (state) {
                 PausedCompositionState.InitialPending,
-                PausedCompositionState.RecomposePending ->
+                PausedCompositionState.RecomposePending,
+                PausedCompositionState.Recomposing ->
                     error("The paused composition has not completed yet")
                 PausedCompositionState.ApplyPending -> {
                     applyChanges()
@@ -240,6 +262,11 @@
         composition.pausedCompositionFinished()
     }
 
+    internal fun markIncomplete() {
+        if (state == PausedCompositionState.ApplyPending)
+            state = PausedCompositionState.RecomposePending
+    }
+
     private fun markComplete() {
         state = PausedCompositionState.ApplyPending
     }
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/internal/RememberEventDispatcher.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/internal/RememberEventDispatcher.kt
index fedee35..308eac4 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/internal/RememberEventDispatcher.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/internal/RememberEventDispatcher.kt
@@ -61,7 +61,9 @@
     private var abandoning: MutableSet<RememberObserver>? = null
     private var traceContext: CompositionErrorContext? = null
     private val remembering = mutableVectorOf<RememberObserverHolder>()
+    private val rememberSet = mutableScatterSetOf<RememberObserverHolder>()
     private var currentRememberingList = remembering
+    private var currentRememberSet = rememberSet
     private val leaving = mutableVectorOf<Any>()
     private val sideEffects = mutableVectorOf<() -> Unit>()
     private var releasing: MutableScatterSet<ComposeNodeLifecycleCallback>? = null
@@ -102,7 +104,9 @@
         this.abandoning = null
         this.traceContext = null
         this.remembering.clear()
+        this.rememberSet.clear()
         this.currentRememberingList = remembering
+        this.currentRememberSet = rememberSet
         this.leaving.clear()
         this.sideEffects.clear()
         this.releasing = null
@@ -115,6 +119,7 @@
 
     override fun remembering(instance: RememberObserverHolder) {
         currentRememberingList.add(instance)
+        currentRememberSet.add(instance)
     }
 
     override fun forgetting(
@@ -123,6 +128,12 @@
         priority: Int,
         endRelativeAfter: Int
     ) {
+        if (instance in currentRememberSet) {
+            currentRememberSet.remove(instance)
+            currentRememberingList.remove(instance)
+            val abandoning = abandoning ?: return
+            abandoning.add(instance.wrapped)
+        }
         recordLeaving(instance, endRelativeOrder, priority, endRelativeAfter)
     }
 
diff --git a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
index b903787..b69f4d9 100644
--- a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
+++ b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
@@ -268,7 +268,10 @@
         /**
          * Returns `true` when the [PausedPrecomposition] is complete. [isComplete] matches the last
          * value returned from [resume]. Once a [PausedPrecomposition] is [isComplete] the [apply]
-         * method should be called.
+         * method should be called. If the [apply] method is not called synchronously and
+         * immediately after [resume] returns `true` then this [isComplete] can return `false` as
+         * any state changes read by the paused composition while it is paused will cause the
+         * composition to require the paused composition to need to be resumed before it is used.
          */
         val isComplete: Boolean
 
@@ -296,6 +299,13 @@
          * Apply the composition. This is the last step of a paused composition and is required to
          * be called prior to the composition is usable.
          *
+         * Calling [apply] should always be proceeded with a check of [isComplete] before it is
+         * called and potentially calling [resume] in a loop until [isComplete] returns `true`. This
+         * can happen if [resume] returned `true` but [apply] was not synchronously called
+         * immediately afterwords. Any state that was read that changed between when [resume] being
+         * called and [apply] being called may require the paused composition to be resumed before
+         * applied.
+         *
          * @return [PrecomposedSlotHandle] you can use to premeasure the slot as well, or to dispose
          *   the composed content.
          */