Remove tracking information from release recompose scopes

Recompose scopes can be tracked by the recomposer past when they
are removed from composition. Since it is expensive to accurately
track the lifetime of scopes they are cleaned up lazily. This
lazy cleanup might not trigger in some cases causing the references
tracked by the scope to leak.

This change first removes tracked references from recmpose scopes
that are removed from composition and, second, reduces the chance
that such scopes are tracked longer then necessary by cleaning up
the conditionally invalid scopes table when scopes are removed.

Fixes: 230168389
Test: ./gradlew :compose:r:r:tDUT
Change-Id: Idfe3e3d5d6d65bac2cf06ee19fd00a5d9253d8a1
(cherry picked from commit 11dbe1f459907ec1c23e730c77e33de39fc894cf)
Merged-In: Idfe3e3d5d6d65bac2cf06ee19fd00a5d9253d8a1
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
index f8ace7f..3435fd2 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
@@ -1776,7 +1776,7 @@
                     is RecomposeScopeImpl -> {
                         val composition = previous.composition
                         if (composition != null) {
-                            previous.composition = null
+                            previous.release()
                             composition.pendingInvalidScopes = true
                         }
                     }
@@ -2671,7 +2671,7 @@
                             val composition = data.composition
                             if (composition != null) {
                                 composition.pendingInvalidScopes = true
-                                data.composition = null
+                                data.release()
                             }
                             reader.reposition(group)
                             recordSlotTableOperation { _, slots, _ ->
@@ -2980,7 +2980,7 @@
                                 // The recompose scope is always at slot 0 of a restart group.
                                 val recomposeScope = slots.slot(anchor, 0) as? RecomposeScopeImpl
                                 // Check for null as the anchor might not be for a recompose scope
-                                recomposeScope?.let { it.composition = toComposition }
+                                recomposeScope?.adoptedBy(toComposition)
                             }
                         }
                     }
@@ -3975,7 +3975,7 @@
                 val composition = slot.composition
                 if (composition != null) {
                     composition.pendingInvalidScopes = true
-                    slot.composition = null
+                    slot.release()
                 }
             }
         }
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 1dc664c..1b652e8 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
@@ -405,6 +405,12 @@
     internal val derivedStateDependencies get() = derivedStates.values.filterNotNull()
 
     /**
+     * Used for testing. Returns the conditional scopes being tracked by the composer
+     */
+    internal val conditionalScopes: List<RecomposeScopeImpl> get() =
+        conditionallyInvalidatedScopes.toList()
+
+    /**
      * A list of changes calculated by [Composer] to be applied to the [Applier] and the
      * [SlotTable] to reflect the result of composition. This is a list of lambdas that need to
      * be invoked in order to produce the desired effects.
@@ -691,6 +697,7 @@
 
     private fun cleanUpDerivedStateObservations() {
         derivedStates.removeValueIf { derivedValue -> derivedValue !in observations }
+        conditionallyInvalidatedScopes.removeValueIf { scope -> !scope.isConditional }
     }
 
     override fun recordReadOf(value: Any) {
@@ -1110,3 +1117,16 @@
         this[key] = IdentityArraySet<V>().also { it.add(value) }
     }
 }
+
+/**
+ * This is provided natively in API 26 and this should be removed if 26 is made the lowest API
+ * level supported
+ */
+private inline fun <E> HashSet<E>.removeValueIf(predicate: (E) -> Boolean) {
+    val iter = iterator()
+    while (iter.hasNext()) {
+        if (predicate(iter.next())) {
+            iter.remove()
+        }
+    }
+}
\ No newline at end of file
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/RecomposeScopeImpl.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/RecomposeScopeImpl.kt
index 36271aa..4914cce 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/RecomposeScopeImpl.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/RecomposeScopeImpl.kt
@@ -47,11 +47,14 @@
  * [Composer.startRestartGroup] and is used to track how to restart the group.
  */
 internal class RecomposeScopeImpl(
-    var composition: CompositionImpl?
+    composition: CompositionImpl?
 ) : ScopeUpdateScope, RecomposeScope {
 
     private var flags: Int = 0
 
+    var composition: CompositionImpl? = composition
+        private set
+
     /**
      * An anchor to the location in the slot table that start the group associated with this
      * recompose scope.
@@ -150,6 +153,24 @@
         composition?.invalidate(this, value) ?: InvalidationResult.IGNORED
 
     /**
+     * Release the recompose scope. This is called when the recompose scope has been removed by the
+     * compostion because the part of the composition it was tracking was removed.
+     */
+    fun release() {
+        composition = null
+        trackedInstances = null
+        trackedDependencies = null
+    }
+
+    /**
+     * Called when the data tracked by this recompose scope moves to a different composition when
+     * for example, the movable content it is part of has moved.
+     */
+    fun adoptedBy(composition: CompositionImpl) {
+        this.composition = composition
+    }
+
+    /**
      * Invalidate the group which will cause [composition] to request this scope be recomposed.
      *
      * Unlike [invalidateForResult], this method is thread safe and calls the thread safe
diff --git a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionAndDerivedStateTests.kt b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionAndDerivedStateTests.kt
index 125ac7d..78cf889 100644
--- a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionAndDerivedStateTests.kt
+++ b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionAndDerivedStateTests.kt
@@ -503,9 +503,36 @@
 
         // Validate there are only 2 observed dependencies, one for intermediateState, one for itemValue
         val observed = (composition as? CompositionImpl)?.derivedStateDependencies ?: emptyList()
-        println(observed)
         assertEquals(2, observed.count())
     }
+
+    @Test
+    fun changingDerivedStateShouldNotAccumulateConditionalScopes() = compositionTest {
+
+        var reload by mutableStateOf(0)
+
+        compose {
+            val derivedState = remember {
+                derivedStateOf {
+                    List(reload) { it }
+                }
+            }
+
+            if (reload % 2 == 0) {
+                Wrap {
+                    Text("${derivedState.value.size}")
+                }
+            }
+        }
+
+        reload++
+
+        advance()
+
+        val conditionalScopes = (composition as? CompositionImpl)?.conditionalScopes ?: emptyList()
+
+        assertEquals(0, conditionalScopes.count { it.isConditional })
+    }
 }
 
 @Composable