Fix memory leak when observing derived state objects

When observing derived state objects, if the recompose snapshot changed
id it caused the derived state objects no longer observed by composition
to be retained by the recompose scope, causing them to leak. The
snapshot id would change if the read is the first time a derived state
object is read in the snapshot.

To detect when a derived state object is no longer necessary a token
value is used to record when the derived state was added into the
composition. This token was the current snapshot id. However, since
the id is not static through-out composition, this could cause the
notification to be skipped.

This change fixes the token to be the snapshot id at the beginning of
composition and, after that, ignores the snapshot id. This means the
token is consistent through-out the same composition allowing unused
derived object to be detected.

Fixes 230168389
Test: new test, ./gradlew :compose:r:r:tDUT

Change-Id: I877c826764592a7e424667d3e2b677d385cf0b7f
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 b80e009..28da12b 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
@@ -1206,6 +1206,7 @@
     private var reusingGroup = -1
     private var childrenComposing: Int = 0
     private var snapshot = currentSnapshot()
+    private var compositionToken: Int = 0
 
     private val invalidateStack = Stack<RecomposeScopeImpl>()
 
@@ -2700,7 +2701,7 @@
             val scope = RecomposeScopeImpl(composition as CompositionImpl)
             invalidateStack.push(scope)
             updateValue(scope)
-            scope.start(snapshot.id)
+            scope.start(compositionToken)
         } else {
             val invalidation = invalidations.removeLocation(reader.parent)
             val slot = reader.next()
@@ -2713,7 +2714,7 @@
             } else slot as RecomposeScopeImpl
             scope.requiresRecompose = invalidation != null
             invalidateStack.push(scope)
-            scope.start(snapshot.id)
+            scope.start(compositionToken)
         }
     }
 
@@ -2731,7 +2732,7 @@
         val scope = if (invalidateStack.isNotEmpty()) invalidateStack.pop()
         else null
         scope?.requiresRecompose = false
-        scope?.end(snapshot.id)?.let {
+        scope?.end(compositionToken)?.let {
             record { _, _, _ -> it(composition) }
         }
         val result = if (scope != null &&
@@ -3150,6 +3151,7 @@
         runtimeCheck(!isComposing) { "Reentrant composition is not supported" }
         trace("Compose:recompose") {
             snapshot = currentSnapshot()
+            compositionToken = snapshot.id
             providerUpdates.clear()
             invalidationsRequested.forEach { scope, set ->
                 val location = scope.anchor?.location ?: return
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 9b9d49a..edaf552 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
@@ -382,6 +382,11 @@
     private val observations = IdentityScopeMap<RecomposeScopeImpl>()
 
     /**
+     * Used for testing. Returns the objects that are observed
+     */
+    internal val observedObjects get() = observations.values.filterNotNull()
+
+    /**
      * A set of scopes that were invalidated conditionally (that is they were invalidated by a
      * [derivedStateOf] object) by a call from [recordModificationsOf]. They need to be held in the
      * [observations] map until invalidations are drained for composition as a later call to
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 db77fc2..98e83fb 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
@@ -354,6 +354,37 @@
         advance()
         revalidate()
     }
+
+    @Test
+    fun changingTheDerivedStateInstanceShouldRelease() = compositionTest {
+        var reload by mutableStateOf(0)
+
+        compose {
+            val items = remember(reload) {
+                derivedStateOf {
+                    List(10) { it }
+                }
+            }
+
+            Text("List of size ${items.value.size}")
+        }
+
+        validate {
+            Text("List of size 10")
+        }
+
+        repeat(10) {
+            reload++
+            advance()
+        }
+
+        revalidate()
+
+        // Validate there are only 2 observed objectt which should be `reload` and the last
+        // created derivedStateOf instance
+        val observed = (composition as? CompositionImpl)?.observedObjects ?: emptyList()
+        assertEquals(2, observed.count())
+    }
 }
 
 @Composable
diff --git a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/mock/CompositionTest.kt b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/mock/CompositionTest.kt
index 774cc09..0cd8df6 100644
--- a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/mock/CompositionTest.kt
+++ b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/mock/CompositionTest.kt
@@ -42,7 +42,7 @@
         // Create a test scope for the test using the test scope passed in by runTest
         val scope = object : CompositionTestScope, CoroutineScope by this@runTest {
             var composed = false
-            var composition: Composition? = null
+            override var composition: Composition? = null
 
             override lateinit var root: View
 
@@ -126,6 +126,11 @@
      * The last validator used.
      */
     var validator: (MockViewValidator.() -> Unit)?
+
+    /**
+     * Access to the composition created for the call to [compose]
+     */
+    val composition: Composition?
 }
 
 /**