Unify Composer observer map with Layout's observer map

Removes the custom observer map that was used in the Composer
in favor of the observer map produced for use by Layout. The
Layout observer map allocates more on initial use but has
fewer allocations during recomposition and has a performance
advantage especially for a large number of invalidations.

Adds a result to IdentityArraySet.remove() to indicate when an
item was removed from the set.

Adds `remove()` to IdentityScopeMap.

Test: :compose:runtime:runtime:tDUT
Change-Id: If88b740d9a91729a65bdc71eaa6d85dbd651044d
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 b7b078a..8435ff1 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
@@ -21,6 +21,7 @@
 )
 package androidx.compose.runtime
 
+import androidx.compose.runtime.collection.IdentityScopeMap
 import androidx.compose.runtime.tooling.InspectionTables
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentHashMapOf
@@ -959,8 +960,8 @@
     private var collectKeySources = false
     private var collectParameterInformation = false
     private var nodeExpected = false
-    private val observations: MutableList<Any> = mutableListOf()
-    private val observationsProcessed: MutableList<Any> = mutableListOf()
+    private val observations = IdentityScopeMap<RecomposeScopeImpl>()
+    private val observationsProcessed = IdentityScopeMap<RecomposeScopeImpl>()
     private val invalidations: MutableList<Invalidation> = mutableListOf()
     internal var pendingInvalidScopes = false
     private val entersStack = IntStack()
@@ -1231,7 +1232,7 @@
         if (childrenComposing == 0) {
             currentRecomposeScope?.let {
                 it.used = true
-                observations.insertIfMissing(value, it)
+                observations.add(value, it)
             }
         }
     }
@@ -1248,7 +1249,7 @@
         observations.forEachScopeOf(value) { scope ->
             if (scope.invalidateForResult() == InvalidationResult.IMMINENT) {
                 // If we process this during recordWriteOf, ignore it when recording modifications
-                observationsProcessed.insertIfMissing(value, scope)
+                observationsProcessed.add(value, scope)
             }
         }
     }
@@ -1287,7 +1288,7 @@
         var invalidated: HashSet<RecomposeScopeImpl>? = null
         for (value in values) {
             observations.forEachScopeOf(value) { scope ->
-                if (!observationsProcessed.removeValueScope(value, scope) &&
+                if (!observationsProcessed.remove(value, scope) &&
                     scope.invalidateForResult() != InvalidationResult.IGNORED
                 ) {
                     (
@@ -1301,7 +1302,7 @@
             }
         }
         invalidated?.let {
-            observations.removeValueIf { _, scope -> scope in it }
+            observations.removeValueIf { scope -> scope in it }
         }
     }
 
@@ -1429,7 +1430,7 @@
 
                 if (pendingInvalidScopes) {
                     pendingInvalidScopes = false
-                    observations.removeValueIf { _, scope -> !scope.valid }
+                    observations.removeValueIf { scope -> !scope.valid }
                 }
             } finally {
                 manager.dispatchAbandons()
@@ -3377,55 +3378,6 @@
     return false
 }
 
-private inline fun MutableList<Any>.removeValueIf(
-    predicate: (value: Any, scope: RecomposeScopeImpl) -> Boolean
-) {
-    var copyLocation = 0
-    for (index in 0 until size / 2) {
-        val slot = index * 2
-        val value = get(slot)
-        val scope = get(slot + 1) as RecomposeScopeImpl
-        if (!predicate(value, scope)) {
-            if (copyLocation != slot) {
-                // Keep the value by copying over a value that has been moved or removed.
-                set(copyLocation++, value)
-                set(copyLocation++, scope)
-            } else {
-                // No slots have been removed yet, just update the copy location
-                copyLocation += 2
-            }
-        }
-    }
-    if (copyLocation < size) {
-        // Delete any left-over slots.
-        subList(copyLocation, size).clear()
-    }
-}
-
-/**
- * Iterate through all the scopes associated with [value]. Returns `false` if [value] has no scopes
- * associated with it.
- */
-private inline fun MutableList<Any>.forEachScopeOf(
-    value: Any,
-    block: (scope: RecomposeScopeImpl) -> Unit
-): Boolean {
-    val valueHash = identityHashCode(value)
-    var index = findFirst(valueHash)
-    var result = false
-    while (index < size) {
-        val storedValue = get(index)
-        if (identityHashCode(storedValue) != valueHash) break
-        if (storedValue === value) {
-            val storedScope = get(index + 1) as RecomposeScopeImpl
-            block(storedScope)
-            result = true
-        }
-        index += 2
-    }
-    return result
-}
-
 private fun MutableList<Any>.find(value: Any, scope: RecomposeScopeImpl): Int {
     val valueHash = identityHashCode(value)
     val scopeHash = identityHashCode(scope)
@@ -3445,9 +3397,6 @@
     return -(index + 1)
 }
 
-private fun MutableList<Any>.findFirst(valueHash: Int) =
-    find(valueHash, 0).let { if (it < 0) -(it + 1) else it }
-
 private fun MutableList<Any>.find(valueHash: Int, scopeHash: Int): Int {
     var low = 0
     var high = (size / 2) - 1
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityArraySet.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityArraySet.kt
index 5e564eb..eca6421 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityArraySet.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityArraySet.kt
@@ -107,7 +107,7 @@
     /**
      * Remove [value] from the set.
      */
-    fun remove(value: T) {
+    fun remove(value: T): Boolean {
         val index = find(value)
         if (index >= 0) {
             if (index < size - 1) {
@@ -120,7 +120,9 @@
             }
             size--
             values[size] = null
+            return true
         }
+        return false
     }
 
     /**
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityScopeMap.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityScopeMap.kt
index 6c8ec16..8a9cec8 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityScopeMap.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityScopeMap.kt
@@ -174,10 +174,44 @@
     }
 
     /**
+     * Remove [scope] from the scope set for [value]. If the scope set is empty after [scope] has
+     * been remove the reference to [value] is removed as well.
+     *
+     * @param value the key of the scope map
+     * @param scope the scope being removed
+     * @return true if the value was removed from the scope
+     */
+    fun remove(value: Any, scope: T): Boolean {
+        val index = find(value)
+        if (index >= 0) {
+            val valueOrderIndex = valueOrder[index]
+            val set = scopeSets[valueOrderIndex] ?: return false
+            val removed = set.remove(scope)
+            if (set.size == 0) {
+                val startIndex = index + 1
+                val endIndex = size
+                if (startIndex < endIndex) {
+                    valueOrder.copyInto(
+                        destination = valueOrder,
+                        destinationOffset = index,
+                        startIndex = startIndex,
+                        endIndex = endIndex
+                    )
+                }
+                valueOrder[size - 1] = valueOrderIndex
+                values[valueOrderIndex] = null
+                size--
+            }
+            return removed
+        }
+        return false
+    }
+
+    /**
      * Removes all scopes that match [predicate]. If all scopes for a given value have been
      * removed, that value is removed also.
      */
-    inline fun removeValueIf(predicate: (scope: Any) -> Boolean) {
+    inline fun removeValueIf(predicate: (scope: T) -> Boolean) {
         var destinationIndex = 0
         for (i in 0 until size) {
             val valueIndex = valueOrder[i]
diff --git a/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityArraySetTest.kt b/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityArraySetTest.kt
index 6984a2d..8397eda 100644
--- a/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityArraySetTest.kt
+++ b/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityArraySetTest.kt
@@ -19,6 +19,7 @@
 import androidx.compose.runtime.identityHashCode
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFalse
 import kotlin.test.assertNotEquals
 import kotlin.test.assertNotSame
 import kotlin.test.assertNull
@@ -89,8 +90,9 @@
         list.forEach { set.add(it) }
 
         // remove a value that doesn't exist:
-        set.remove(Stuff(10))
+        val removed = set.remove(Stuff(10))
         assertEquals(list.size, set.size)
+        assertFalse(removed)
 
         // remove a value in the middle:
         testRemoveValueAtIndex(set.size / 2)
@@ -131,8 +133,9 @@
     private fun testRemoveValueAtIndex(index: Int) {
         val value = set[index]
         val initialSize = set.size
-        set.remove(value)
+        val removed = set.remove(value)
         assertEquals(initialSize - 1, set.size)
+        assertTrue(removed)
         assertNull(set.values[set.size])
         set.forEach { assertNotSame(value, it) }
     }
diff --git a/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt b/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt
index 4e451ad..391a571 100644
--- a/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt
+++ b/compose/runtime/runtime/src/test/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt
@@ -16,8 +16,12 @@
 
 package androidx.compose.runtime.collection
 
+import org.junit.After
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
 import kotlin.test.assertTrue
 import kotlin.test.fail
 
@@ -84,6 +88,25 @@
     }
 
     @Test
+    fun removeScope() {
+        val valueC = "C"
+        map.add(valueList[0], scopeList[0])
+        map.add(valueList[0], scopeList[1])
+        map.add(valueList[1], scopeList[2])
+        map.add(valueC, scopeList[3])
+
+        // remove a non existent value, should leave the map unmodified
+        val removed1 = map.remove(valueC, scopeList[0])
+        assertFalse(removed1)
+        assertEquals(3, map.size)
+
+        // Remove a reference which should remove the value
+        val removed2 = map.remove(valueList[1], scopeList[2])
+        assertTrue(removed2)
+        assertEquals(2, map.size)
+    }
+
+    @Test
     fun removeValueIf() {
         val valueC = "C"
         map.add(valueList[0], scopeList[0])
@@ -109,5 +132,38 @@
         }
     }
 
+    /**
+     * Validate the test maintains the internal assumptions of the map.
+     */
+    @After
+    fun validateMap() {
+        // Ensure that no duplicates exist in value-order and all indexes are represented
+        val pendingRepresentation = mutableSetOf(*map.values.indices.toList().toTypedArray())
+        map.valueOrder.forEach {
+            assertTrue(it in pendingRepresentation, "Index $it was duplicated")
+            pendingRepresentation.remove(it)
+        }
+        assertTrue(pendingRepresentation.isEmpty(), "Not all indexes are in the valueOrder map")
+
+        // Ensure values are non-null and sets are not empty for index < size and values are
+        // null and sets are empty or missing for >= size
+        val size = map.size
+        map.valueOrder.forEachIndexed { index, order ->
+            val value = map.values[order]
+            val set = map.scopeSets[order]
+            if (index < size) {
+                assertNotNull(value, "A value was unexpectedly null")
+                assertNotNull(set, "A set was unexpectedly null")
+                assertTrue(set.size > 0, "An empty set wasn't collected")
+            } else {
+                assertNull(value, "A reference to a removed value was retained")
+                assertTrue(
+                    actual = set == null || set.size == 0,
+                    message = "A non-empty set was dropped"
+                )
+            }
+        }
+    }
+
     data class Scope(val item: Int)
 }
\ No newline at end of file