Use ScatterMap for Compose ScopeMap

Changes underlying data structure to use `ScatterMap` without changing access patterns. Adds a special case for set containing a single value, which is much cheaper than allocating a new set.

Test: ScopeMapTest
Change-Id: I945cefd280dcbcdd24b1b447ef1eb899704811b7
diff --git a/collection/collection/api/current.txt b/collection/collection/api/current.txt
index 53cd0b5..2c3b879 100644
--- a/collection/collection/api/current.txt
+++ b/collection/collection/api/current.txt
@@ -1639,7 +1639,7 @@
     method public void putAll(kotlin.sequences.Sequence<? extends kotlin.Pair<? extends K,? extends V>> pairs);
     method public V? remove(K key);
     method public boolean remove(K key, V value);
-    method public void removeIf(kotlin.jvm.functions.Function2<? super K,? super V,java.lang.Boolean> predicate);
+    method public inline void removeIf(kotlin.jvm.functions.Function2<? super K,? super V,java.lang.Boolean> predicate);
     method public operator void set(K key, V value);
     method public int trim();
   }
diff --git a/collection/collection/api/restricted_current.txt b/collection/collection/api/restricted_current.txt
index b1ab5c9..15f1d60 100644
--- a/collection/collection/api/restricted_current.txt
+++ b/collection/collection/api/restricted_current.txt
@@ -1711,7 +1711,8 @@
     method public void putAll(kotlin.sequences.Sequence<? extends kotlin.Pair<? extends K,? extends V>> pairs);
     method public V? remove(K key);
     method public boolean remove(K key, V value);
-    method public void removeIf(kotlin.jvm.functions.Function2<? super K,? super V,java.lang.Boolean> predicate);
+    method public inline void removeIf(kotlin.jvm.functions.Function2<? super K,? super V,java.lang.Boolean> predicate);
+    method @kotlin.PublishedApi internal V? removeValueAt(int index);
     method public operator void set(K key, V value);
     method public int trim();
   }
diff --git a/collection/collection/src/commonMain/kotlin/androidx/collection/ScatterMap.kt b/collection/collection/src/commonMain/kotlin/androidx/collection/ScatterMap.kt
index 547f6a3..8bfc87c 100644
--- a/collection/collection/src/commonMain/kotlin/androidx/collection/ScatterMap.kt
+++ b/collection/collection/src/commonMain/kotlin/androidx/collection/ScatterMap.kt
@@ -987,7 +987,7 @@
     /**
      * Removes any mapping for which the specified [predicate] returns true.
      */
-    public fun removeIf(predicate: (K, V) -> Boolean) {
+    public inline fun removeIf(predicate: (K, V) -> Boolean) {
         forEachIndexed { index ->
             @Suppress("UNCHECKED_CAST")
             if (predicate(keys[index] as K, values[index] as V)) {
@@ -1048,7 +1048,8 @@
         }
     }
 
-    private fun removeValueAt(index: Int): V? {
+    @PublishedApi
+    internal fun removeValueAt(index: Int): V? {
         _size -= 1
 
         // TODO: We could just mark the entry as empty if there's a group
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 e806853..d5a32ab 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
@@ -20,7 +20,7 @@
 import androidx.compose.runtime.changelist.ChangeList
 import androidx.compose.runtime.collection.IdentityArrayMap
 import androidx.compose.runtime.collection.IdentityArraySet
-import androidx.compose.runtime.collection.IdentityScopeMap
+import androidx.compose.runtime.collection.ScopeMap
 import androidx.compose.runtime.collection.fastForEach
 import androidx.compose.runtime.snapshots.fastAll
 import androidx.compose.runtime.snapshots.fastAny
@@ -471,12 +471,13 @@
      * A map of observable objects to the [RecomposeScope]s that observe the object. If the key
      * object is modified the associated scopes should be invalidated.
      */
-    private val observations = IdentityScopeMap<RecomposeScopeImpl>()
+    private val observations = ScopeMap<RecomposeScopeImpl>()
 
     /**
      * Used for testing. Returns the objects that are observed
      */
-    internal val observedObjects get() = observations.values.filterNotNull()
+    internal val observedObjects
+        @TestOnly @Suppress("AsCollectionCall") get() = observations.map.asMap().keys
 
     /**
      * A set of scopes that were invalidated conditionally (that is they were invalidated by a
@@ -489,18 +490,19 @@
     /**
      * A map of object read during derived states to the corresponding derived state.
      */
-    private val derivedStates = IdentityScopeMap<DerivedState<*>>()
+    private val derivedStates = ScopeMap<DerivedState<*>>()
 
     /**
      * Used for testing. Returns dependencies of derived states that are currently observed.
      */
-    internal val derivedStateDependencies get() = derivedStates.values.filterNotNull()
+    internal val derivedStateDependencies
+        @TestOnly @Suppress("AsCollectionCall") get() = derivedStates.map.asMap().keys
 
     /**
      * Used for testing. Returns the conditional scopes being tracked by the composer
      */
-    internal val conditionalScopes: List<RecomposeScopeImpl> get() =
-        conditionallyInvalidatedScopes.toList()
+    internal val conditionalScopes: List<RecomposeScopeImpl>
+        @TestOnly get() = conditionallyInvalidatedScopes.toList()
 
     /**
      * A list of changes calculated by [Composer] to be applied to the [Applier] and the
@@ -526,7 +528,7 @@
      * scopes that were already dismissed by composition and should be ignored in the next call
      * to [recordModificationsOf].
      */
-    private val observationsProcessed = IdentityScopeMap<RecomposeScopeImpl>()
+    private val observationsProcessed = ScopeMap<RecomposeScopeImpl>()
 
     /**
      * A map of the invalid [RecomposeScope]s. If this map is non-empty the current state of
@@ -856,21 +858,21 @@
         }
 
         if (forgetConditionalScopes && conditionallyInvalidatedScopes.isNotEmpty()) {
-            observations.removeValueIf { scope ->
+            observations.removeScopeIf { scope ->
                 scope in conditionallyInvalidatedScopes || invalidated?.let { scope in it } == true
             }
             conditionallyInvalidatedScopes.clear()
             cleanUpDerivedStateObservations()
         } else {
             invalidated?.let {
-                observations.removeValueIf { scope -> scope in it }
+                observations.removeScopeIf { scope -> scope in it }
                 cleanUpDerivedStateObservations()
             }
         }
     }
 
     private fun cleanUpDerivedStateObservations() {
-        derivedStates.removeValueIf { derivedState -> derivedState !in observations }
+        derivedStates.removeScopeIf { derivedState -> derivedState !in observations }
         if (conditionallyInvalidatedScopes.isNotEmpty()) {
             conditionallyInvalidatedScopes.removeValueIf { scope -> !scope.isConditional }
         }
@@ -979,7 +981,7 @@
             if (pendingInvalidScopes) {
                 trace("Compose:unobserve") {
                     pendingInvalidScopes = false
-                    observations.removeValueIf { scope -> !scope.valid }
+                    observations.removeScopeIf { scope -> !scope.valid }
                     cleanUpDerivedStateObservations()
                 }
             }
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
deleted file mode 100644
index 3ee5190..0000000
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/IdentityScopeMap.kt
+++ /dev/null
@@ -1,331 +0,0 @@
-/*
- * Copyright 2020 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.collection
-
-import androidx.compose.runtime.identityHashCode
-import kotlin.contracts.ExperimentalContracts
-
-/**
- * Maps values to a set of scopes using the [identityHashCode] for both the value and the
- * scope for uniqueness.
- */
-@OptIn(ExperimentalContracts::class)
-internal class IdentityScopeMap<T : Any> {
-    /**
-     * The array of indices into [values] and [scopeSets], in the order that they are sorted
-     * in the [IdentityScopeMap]. The length of the used values is [size], and all remaining values
-     * are the unused indices in [values] and [scopeSets].
-     */
-    var valueOrder: IntArray = IntArray(50) { it }
-        private set
-
-    /**
-     * The [identityHashCode] for the keys in the collection. We never use the actual
-     * values
-     */
-    var values: Array<Any?> = arrayOfNulls(50)
-        private set
-
-    /**
-     * The [IdentityArraySet]s for values, in the same index order as [values], indexed
-     * by [valueOrder]. The consumed values may extend beyond [size] if a value has been removed.
-     */
-    var scopeSets: Array<IdentityArraySet<T>?> = arrayOfNulls(50)
-        private set
-
-    /**
-     * The number of values in the map.
-     */
-    var size = 0
-
-    /**
-     * Returns the [IdentityArraySet] for the value at the given [index] order in the map.
-     */
-    private fun scopeSetAt(index: Int): IdentityArraySet<T> {
-        return scopeSets[valueOrder[index]]!!
-    }
-
-    /**
-     * Adds a [value]/[scope] pair to the map and returns `true` if it was added or `false` if
-     * it already existed.
-     */
-    fun add(value: Any, scope: T): Boolean {
-        val valueSet = getOrCreateIdentitySet(value)
-        return valueSet.add(scope)
-    }
-
-    /**
-     * Returns true if any scopes are associated with [element]
-     */
-    operator fun contains(element: Any): Boolean = find(element) >= 0
-
-    /**
-     * Executes [block] for all scopes mapped to the given [value].
-     */
-    inline fun forEachScopeOf(value: Any, block: (scope: T) -> Unit) {
-        val index = find(value)
-        if (index >= 0) {
-            scopeSetAt(index).fastForEach(block)
-        }
-    }
-
-    /**
-     * Returns the existing [IdentityArraySet] for the given [value] or creates a new one
-     * and insertes it into the map and returns it.
-     */
-    private fun getOrCreateIdentitySet(value: Any): IdentityArraySet<T> {
-        val size = size
-        val valueOrder = valueOrder
-        val values = values
-        val scopeSets = scopeSets
-
-        val index: Int
-        if (size > 0) {
-            index = find(value)
-
-            if (index >= 0) {
-                return scopeSetAt(index)
-            }
-        } else {
-            index = -1
-        }
-
-        val insertIndex = -(index + 1)
-
-        if (size < valueOrder.size) {
-            val valueIndex = valueOrder[size]
-            values[valueIndex] = value
-            val scopeSet = scopeSets[valueIndex] ?: IdentityArraySet<T>().also {
-                scopeSets[valueIndex] = it
-            }
-
-            // insert into the right location in keyOrder
-            if (insertIndex < size) {
-                valueOrder.copyInto(
-                    destination = valueOrder,
-                    destinationOffset = insertIndex + 1,
-                    startIndex = insertIndex,
-                    endIndex = size
-                )
-            }
-            valueOrder[insertIndex] = valueIndex
-            this.size++
-            return scopeSet
-        }
-
-        // We have to increase the size of all arrays
-        val newSize = valueOrder.size * 2
-        val valueIndex = size
-        val newScopeSets = scopeSets.copyOf(newSize)
-        val scopeSet = IdentityArraySet<T>()
-        newScopeSets[valueIndex] = scopeSet
-        val newValues = values.copyOf(newSize)
-        newValues[valueIndex] = value
-
-        val newKeyOrder = IntArray(newSize)
-        for (i in size + 1 until newSize) {
-            newKeyOrder[i] = i
-        }
-
-        if (insertIndex < size) {
-            valueOrder.copyInto(
-                destination = newKeyOrder,
-                destinationOffset = insertIndex + 1,
-                startIndex = insertIndex,
-                endIndex = size
-            )
-        }
-        newKeyOrder[insertIndex] = valueIndex
-        if (insertIndex > 0) {
-            valueOrder.copyInto(
-                destination = newKeyOrder,
-                endIndex = insertIndex
-            )
-        }
-        this.scopeSets = newScopeSets
-        this.values = newValues
-        this.valueOrder = newKeyOrder
-        this.size++
-        return scopeSet
-    }
-
-    /**
-     * Removes all values and scopes from the map
-     */
-    fun clear() {
-        val scopeSets = scopeSets
-        val valueOrder = valueOrder
-        val values = values
-
-        for (i in scopeSets.indices) {
-            scopeSets[i]?.clear()
-            valueOrder[i] = i
-            values[i] = null
-        }
-
-        size = 0
-    }
-
-    /**
-     * 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)
-
-        val valueOrder = valueOrder
-        val scopeSets = scopeSets
-        val values = values
-        val size = size
-        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
-                    )
-                }
-                val newSize = size - 1
-                valueOrder[newSize] = valueOrderIndex
-                values[valueOrderIndex] = null
-                this.size = newSize
-            }
-            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: T) -> Boolean) {
-        removingScopes { scopeSet ->
-            scopeSet.removeValueIf(predicate)
-        }
-    }
-
-    /**
-     * Removes given scope from all sets. If all scopes for a given value are removed, that value
-     * is removed as well.
-     */
-    fun removeScope(scope: T) {
-        removingScopes { scopeSet ->
-            scopeSet.remove(scope)
-        }
-    }
-
-    private inline fun removingScopes(removalOperation: (IdentityArraySet<T>) -> Unit) {
-        val valueOrder = valueOrder
-        val scopeSets = scopeSets
-        val values = values
-        var destinationIndex = 0
-        for (i in 0 until size) {
-            val valueIndex = valueOrder[i]
-            val set = scopeSets[valueIndex]!!
-            removalOperation(set)
-            if (set.size > 0) {
-                if (destinationIndex != i) {
-                    // We'll bubble-up the now-free key-order by swapping the index with the one
-                    // we're copying from. This means that the set can be reused later.
-                    val destinationKeyOrder = valueOrder[destinationIndex]
-                    valueOrder[destinationIndex] = valueIndex
-                    valueOrder[i] = destinationKeyOrder
-                }
-                destinationIndex++
-            }
-        }
-        // Remove hard references to values that are no longer in the map
-        for (i in destinationIndex until size) {
-            values[valueOrder[i]] = null
-        }
-        size = destinationIndex
-    }
-
-    /**
-     * Returns the index into [valueOrder] of the found [value] of the
-     * value, or the negative index - 1 of the position in which it would be if it were found.
-     */
-    private fun find(value: Any?): Int {
-        val valueIdentity = identityHashCode(value)
-        var low = 0
-        var high = size - 1
-
-        val values = values
-        val valueOrder = valueOrder
-        while (low <= high) {
-            val mid = (low + high).ushr(1)
-            val midValue = values[valueOrder[mid]]
-            val midValHash = identityHashCode(midValue)
-            when {
-                midValHash < valueIdentity -> low = mid + 1
-                midValHash > valueIdentity -> high = mid - 1
-                value === midValue -> return mid
-                else -> return findExactIndex(mid, value, valueIdentity)
-            }
-        }
-        return -(low + 1)
-    }
-
-    /**
-     * When multiple items share the same [identityHashCode], then we must find the specific
-     * index of the target item. This method assumes that [midIndex] has already been checked
-     * for an exact match for [value], but will look at nearby values to find the exact item index.
-     * If no match is found, the negative index - 1 of the position in which it would be will
-     * be returned, which is always after the last item with the same [identityHashCode].
-     */
-    private fun findExactIndex(midIndex: Int, value: Any?, valueHash: Int): Int {
-        val values = values
-        val valueOrder = valueOrder
-
-        // hunt down first
-        for (i in midIndex - 1 downTo 0) {
-            val v = values[valueOrder[i]]
-            if (v === value) {
-                return i
-            }
-            if (identityHashCode(v) != valueHash) {
-                break // we've gone too far
-            }
-        }
-
-        for (i in midIndex + 1 until size) {
-            val v = values[valueOrder[i]]
-            if (v === value) {
-                return i
-            }
-            if (identityHashCode(v) != valueHash) {
-                // We've gone too far. We should insert here.
-                return -(i + 1)
-            }
-        }
-
-        // We should insert at the end
-        return -(size + 1)
-    }
-}
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/ScopeMap.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/ScopeMap.kt
new file mode 100644
index 0000000..44ee173
--- /dev/null
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/collection/ScopeMap.kt
@@ -0,0 +1,146 @@
+/*
+ * Copyright 2023 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.collection
+
+import androidx.collection.MutableScatterSet
+import androidx.collection.mutableScatterMapOf
+
+/**
+ * Maps values to a set of scopes.
+ */
+internal class ScopeMap<T : Any> {
+    val map = mutableScatterMapOf<Any, Any>()
+
+    /**
+     * The number of values in the map.
+     */
+    val size get() = map.size
+
+    /**
+     * Adds a [key]/[scope] pair to the map and returns `true` if it was added or `false` if
+     * it already existed.
+     */
+    fun add(key: Any, scope: T): Boolean =
+        when (val value = map[key]) {
+            null -> {
+                map[key] = scope
+                true
+            }
+            is MutableScatterSet<*> -> {
+                @Suppress("UNCHECKED_CAST")
+                (value as MutableScatterSet<T>).add(scope)
+            }
+            else -> {
+                if (value !== scope) {
+                    val set = MutableScatterSet<T>()
+                    map[key] = set
+                    @Suppress("UNCHECKED_CAST")
+                    set.add(value as T)
+                    set.add(scope)
+                } else {
+                    false
+                }
+            }
+        }
+
+    /**
+     * Returns true if any scopes are associated with [element]
+     */
+    operator fun contains(element: Any): Boolean = map.containsKey(element)
+
+    /**
+     * Executes [block] for all scopes mapped to the given [key].
+     */
+    inline fun forEachScopeOf(key: Any, block: (scope: T) -> Unit) {
+        when (val value = map[key]) {
+            null -> { /* do nothing */ }
+            is MutableScatterSet<*> -> {
+                @Suppress("UNCHECKED_CAST")
+                (value as MutableScatterSet<T>).forEach(block)
+            }
+            else -> {
+                @Suppress("UNCHECKED_CAST")
+                block(value as T)
+            }
+        }
+    }
+
+    /**
+     * Removes all values and scopes from the map
+     */
+    fun clear() {
+        map.clear()
+    }
+
+    /**
+     * Remove [scope] from the scope set for [key]. If the scope set is empty after [scope] has
+     * been remove the reference to [key] is removed as well.
+     *
+     * @param key the key of the scope map
+     * @param scope the scope being removed
+     * @return true if the value was removed from the scope
+     */
+    fun remove(key: Any, scope: T): Boolean {
+        val value = map[key] ?: return false
+        return when (value) {
+            is MutableScatterSet<*> -> {
+                @Suppress("UNCHECKED_CAST")
+                val set = value as MutableScatterSet<T>
+
+                val removed = set.remove(scope)
+                if (removed && set.isEmpty()) {
+                    map.remove(key)
+                }
+                return removed
+            }
+            scope -> {
+                map.remove(key)
+                true
+            }
+            else -> false
+        }
+    }
+
+    /**
+     * Removes all scopes that match [predicate]. If all scopes for a given value have been
+     * removed, that value is removed also.
+     */
+    inline fun removeScopeIf(crossinline predicate: (scope: T) -> Boolean) {
+        map.removeIf { _, value ->
+            when (value) {
+                is MutableScatterSet<*> -> {
+                    @Suppress("UNCHECKED_CAST")
+                    val set = value as MutableScatterSet<T>
+                    set.removeIf(predicate)
+                    set.isEmpty()
+                }
+                else -> {
+                    @Suppress("UNCHECKED_CAST")
+                    predicate(value as T)
+                }
+            }
+        }
+    }
+
+    /**
+     * Removes given scope from all sets. If all scopes for a given value are removed, that value
+     * is removed as well.
+     */
+    fun removeScope(scope: T) {
+        removeScopeIf { it === scope }
+    }
+}
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotStateObserver.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotStateObserver.kt
index c346941..80c43ea 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotStateObserver.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotStateObserver.kt
@@ -23,7 +23,7 @@
 import androidx.compose.runtime.collection.IdentityArrayIntMap
 import androidx.compose.runtime.collection.IdentityArrayMap
 import androidx.compose.runtime.collection.IdentityArraySet
-import androidx.compose.runtime.collection.IdentityScopeMap
+import androidx.compose.runtime.collection.ScopeMap
 import androidx.compose.runtime.collection.fastForEach
 import androidx.compose.runtime.collection.mutableVectorOf
 import androidx.compose.runtime.composeRuntimeError
@@ -378,7 +378,7 @@
         /**
          * Values that have been read during the scope's [SnapshotStateObserver.observeReads].
          */
-        private val valueToScopes = IdentityScopeMap<Any>()
+        private val valueToScopes = ScopeMap<Any>()
 
         /**
          * Reverse index (scope -> values) for faster scope invalidation.
@@ -422,7 +422,7 @@
         /**
          * Invalidation index from state objects to derived states reading them.
          */
-        private val dependencyToDerivedStates = IdentityScopeMap<DerivedState<*>>()
+        private val dependencyToDerivedStates = ScopeMap<DerivedState<*>>()
 
         /**
          * Last derived state value recorded during read.
diff --git a/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt b/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/ScopeMapTest.kt
similarity index 66%
rename from compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt
rename to compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/ScopeMapTest.kt
index 01ef6a8..9c3887b 100644
--- a/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/IdentityScopeMapTest.kt
+++ b/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/collection/ScopeMapTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2023 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.
@@ -16,24 +16,21 @@
 
 package androidx.compose.runtime.collection
 
-import kotlin.test.AfterTest
 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
 
-class IdentityScopeMapTest {
-    private val map = IdentityScopeMap<Scope>()
+class ScopeMapTest {
+    private val map = ScopeMap<Scope>()
 
     private val scopeList = listOf(Scope(10), Scope(12), Scope(1), Scope(30), Scope(10))
     private val valueList = listOf(Value("A"), Value("B"))
 
     @Test
     fun emptyConstruction() {
-        val m = IdentityScopeMap<Test>()
+        val m = ScopeMap<Scope>()
         assertEquals(0, m.size)
     }
 
@@ -83,8 +80,7 @@
         map.add(valueList[1], scopeList[1])
         map.clear()
         assertEquals(0, map.size)
-        assertEquals(0, map.scopeSets[0]!!.size)
-        assertEquals(0, map.scopeSets[1]!!.size)
+        assertEquals(0, map.map.size)
     }
 
     @Test
@@ -115,17 +111,16 @@
         map.add(valueC, scopeList[3])
 
         // remove a scope that won't cause any values to be removed:
-        map.removeValueIf { scope ->
+        map.removeScopeIf { scope ->
             scope === scopeList[1]
         }
         assertEquals(3, map.size)
 
         // remove the last scope in a set:
-        map.removeValueIf { scope ->
+        map.removeScopeIf { scope ->
             scope === scopeList[2]
         }
         assertEquals(2, map.size)
-        assertEquals(0, map.scopeSets[map.valueOrder[2]]!!.size)
 
         map.forEachScopeOf(valueList[1]) {
             fail("There shouldn't be any scopes for this value")
@@ -146,39 +141,6 @@
         assertFalse(Value("D") in map)
     }
 
-    /**
-     * Validate the test maintains the internal assumptions of the map.
-     */
-    @AfterTest
-    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)
-    data class Value(val s: String)
+    class Scope(val item: Int)
+    class Value(val s: String)
 }