[JS IR] Improve InternalStringMap iterators & Extend InternalMap interface

^KT-59001
2023-06-19 16:39:40 +02:00
parent 6f0628cb8a
commit f202865080
4 changed files with 222 additions and 105 deletions
@@ -30,7 +30,7 @@ public actual open class HashMap<K, V> : AbstractMutableMap<K, V>, MutableMap<K,

        override fun containsEntry(element: Map.Entry<K, V>): Boolean = this@HashMap.containsEntry(element)

-        override operator fun iterator(): MutableIterator<MutableEntry<K, V>> = internalMap.iterator()
+        override operator fun iterator(): MutableIterator<MutableEntry<K, V>> = internalMap.entriesIterator()

        override fun removeEntry(element: Map.Entry<K, V>): Boolean {
            if (contains(element)) {
@@ -107,7 +107,7 @@ public actual open class HashMap<K, V> : AbstractMutableMap<K, V>, MutableMap<K,

    actual override fun containsKey(key: K): Boolean = internalMap.contains(key)

-    actual override fun containsValue(value: V): Boolean = internalMap.any { it.value == value }
+    actual override fun containsValue(value: V): Boolean = internalMap.containsValue(value)

    private var _entries: MutableSet<MutableMap.MutableEntry<K, V>>? = null
    actual override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
@@ -87,7 +87,7 @@ internal class InternalHashMap<K, V> private constructor(

    fun isEmpty(): Boolean = _size == 0
    fun containsKey(key: K): Boolean = findKey(key) >= 0
-    fun containsValue(value: V): Boolean = findValue(value) >= 0
+    override fun containsValue(value: V): Boolean = findValue(value) >= 0

    override operator fun get(key: K): V? {
        val index = findKey(key)
@@ -113,7 +113,7 @@ internal class InternalHashMap<K, V> private constructor(
        }
    }

-    fun putAll(from: Map<out K, V>) {
+    override fun putAll(from: Map<out K, V>) {
        checkIsMutable()
        putAllEntries(from.entries)
    }
@@ -143,10 +143,6 @@ internal class InternalHashMap<K, V> private constructor(
        length = 0
    }

-    override fun iterator(): MutableIterator<MutableMap.MutableEntry<K, V>> {
-        return entriesIterator()
-    }
-
    val keys: MutableSet<K>
        get() {
            val cur = keysView
@@ -211,7 +207,7 @@ internal class InternalHashMap<K, V> private constructor(
    private val capacity: Int get() = keysArray.size
    private val hashSize: Int get() = hashArray.size

-    internal fun checkIsMutable() {
+    override fun checkIsMutable() {
        if (isReadOnly) throw UnsupportedOperationException()
    }

@@ -425,13 +421,18 @@ internal class InternalHashMap<K, V> private constructor(
        }
    }

-    internal fun containsEntry(entry: Map.Entry<K, V>): Boolean {
+    override fun containsEntry(entry: Map.Entry<K, V>): Boolean {
        val index = findKey(entry.key)
        if (index < 0) return false
        return valuesArray!![index] == entry.value
    }

-    internal fun getEntry(entry: Map.Entry<K, V>): MutableMap.MutableEntry<K, V>? {
+    override fun containsOtherEntry(entry: Map.Entry<*, *>): Boolean {
+        @Suppress("UNCHECKED_CAST")
+        return containsEntry(entry as Map.Entry<K, V>)
+    }
+
+    override fun getEntry(entry: Map.Entry<K, V>): MutableMap.MutableEntry<K, V>? {
        val index = findKey(entry.key)
        return if (index < 0 || valuesArray!![index] != entry.value) {
            null
@@ -451,21 +452,6 @@ internal class InternalHashMap<K, V> private constructor(

    private fun contentEquals(other: Map<*, *>): Boolean = _size == other.size && containsAllEntries(other.entries)

-    internal fun containsAllEntries(m: Collection<*>): Boolean {
-        val it = m.iterator()
-        while (it.hasNext()) {
-            val entry = it.next()
-            try {
-                @Suppress("UNCHECKED_CAST") // todo: get rid of unchecked cast here somehow
-                if (entry == null || !containsEntry(entry as Map.Entry<K, V>))
-                    return false
-            } catch (e: ClassCastException) {
-                return false
-            }
-        }
-        return true
-    }
-
    private fun putEntry(entry: Map.Entry<K, V>): Boolean {
        val index = addKey(entry.key)
        val valuesArray = allocateValuesArray()
@@ -493,7 +479,7 @@ internal class InternalHashMap<K, V> private constructor(
        return updated
    }

-    internal fun removeEntry(entry: Map.Entry<K, V>): Boolean {
+    override fun removeEntry(entry: Map.Entry<K, V>): Boolean {
        checkIsMutable()
        val index = findKey(entry.key)
        if (index < 0) return false
@@ -502,7 +488,7 @@ internal class InternalHashMap<K, V> private constructor(
        return true
    }

-    internal fun removeValue(element: V): Boolean {
+    override fun removeValue(element: V): Boolean {
        checkIsMutable()
        val index = findValue(element)
        if (index < 0) return false
@@ -510,9 +496,9 @@ internal class InternalHashMap<K, V> private constructor(
        return true
    }

-    internal fun keysIterator() = KeysItr(this)
-    internal fun valuesIterator() = ValuesItr(this)
-    internal fun entriesIterator() = EntriesItr(this)
+    override fun keysIterator() = KeysItr(this)
+    override fun valuesIterator() = ValuesItr(this)
+    override fun entriesIterator() = EntriesItr(this)

    private companion object {
        private const val MAGIC = -1640531527 // 2654435769L.toInt(), golden ratio
@@ -638,16 +624,16 @@ internal class InternalHashMap<K, V> private constructor(
 }

 internal class HashMapKeys<E> internal constructor(
-    private val backing: InternalHashMap<E, *>,
+    private val backing: InternalMap<E, *>,
 ) : MutableSet<E>, AbstractMutableSet<E>() {

    override val size: Int get() = backing.size
-    override fun isEmpty(): Boolean = backing.isEmpty()
-    override fun contains(element: E): Boolean = backing.containsKey(element)
+    override fun isEmpty(): Boolean = backing.size == 0
+    override fun contains(element: E): Boolean = backing.contains(element)
    override fun clear() = backing.clear()
    override fun add(element: E): Boolean = throw UnsupportedOperationException()
    override fun addAll(elements: Collection<E>): Boolean = throw UnsupportedOperationException()
-    override fun remove(element: E): Boolean = backing.removeKey(element) >= 0
+    override fun remove(element: E): Boolean = backing.remove(element) != null
    override fun iterator(): MutableIterator<E> = backing.keysIterator()

    override fun removeAll(elements: Collection<E>): Boolean {
@@ -662,11 +648,11 @@ internal class HashMapKeys<E> internal constructor(
 }

 internal class HashMapValues<V> internal constructor(
-    val backing: InternalHashMap<*, V>,
+    private val backing: InternalMap<*, V>,
 ) : MutableCollection<V>, AbstractMutableCollection<V>() {

    override val size: Int get() = backing.size
-    override fun isEmpty(): Boolean = backing.isEmpty()
+    override fun isEmpty(): Boolean = backing.size == 0
    override fun contains(element: V): Boolean = backing.containsValue(element)
    override fun add(element: V): Boolean = throw UnsupportedOperationException()
    override fun addAll(elements: Collection<V>): Boolean = throw UnsupportedOperationException()
@@ -700,11 +686,11 @@ internal class HashMapValues<V> internal constructor(
 * See also [KT-42248](https://youtrack.jetbrains.com/issue/KT-42428).
 */
 internal abstract class HashMapEntrySetBase<K, V, E : Map.Entry<K, V>> internal constructor(
-    val backing: InternalHashMap<K, V>,
+    val backing: InternalMap<K, V>,
 ) : MutableSet<E>, AbstractMutableSet<E>() {

    override val size: Int get() = backing.size
-    override fun isEmpty(): Boolean = backing.isEmpty()
+    override fun isEmpty(): Boolean = backing.size == 0
    override fun contains(element: E): Boolean = backing.containsEntry(element)
    protected abstract fun getEntry(element: Map.Entry<K, V>): E?
    override fun clear() = backing.clear()
@@ -725,7 +711,7 @@ internal abstract class HashMapEntrySetBase<K, V, E : Map.Entry<K, V>> internal
 }

 internal class HashMapEntrySet<K, V> internal constructor(
-    backing: InternalHashMap<K, V>,
+    backing: InternalMap<K, V>,
 ) : HashMapEntrySetBase<K, V, MutableMap.MutableEntry<K, V>>(backing) {

    override fun getEntry(element: Map.Entry<K, V>): MutableMap.MutableEntry<K, V>? = backing.getEntry(element)
@@ -6,22 +6,35 @@
 package kotlin.collections

 /**
- * The common interface of [InternalStringMap] and [InternalHashCodeMap].
+ * The common interface of [InternalStringMap] and [InternalHashMap].
 */
-internal interface InternalMap<K, V> : MutableIterable<MutableMap.MutableEntry<K, V>> {
+internal interface InternalMap<K, V> {
    val size: Int
-    operator fun contains(key: K): Boolean
-    operator fun get(key: K): V?

    fun put(key: K, value: V): V?
-    fun remove(key: K): V?
-    fun clear(): Unit
+    fun putAll(from: Map<out K, V>)

-    fun createJsMap(): dynamic {
-        val result = js("Object.create(null)")
-        // force to switch object representation to dictionary mode
-        result["foo"] = 1
-        jsDeleteProperty(result, "foo")
-        return result
+    operator fun get(key: K): V?
+    fun getEntry(entry: Map.Entry<K, V>): MutableMap.MutableEntry<K, V>?
+
+    operator fun contains(key: K): Boolean
+    fun containsValue(value: V): Boolean
+    fun containsEntry(entry: Map.Entry<K, V>): Boolean
+    fun containsOtherEntry(entry: Map.Entry<*, *>): Boolean
+
+    fun remove(key: K): V?
+    fun removeValue(value: V): Boolean
+    fun removeEntry(entry: Map.Entry<K, V>): Boolean
+
+    fun clear()
+
+    fun keysIterator(): MutableIterator<K>
+    fun valuesIterator(): MutableIterator<V>
+    fun entriesIterator(): MutableIterator<MutableMap.MutableEntry<K, V>>
+
+    fun checkIsMutable() {}
+
+    fun containsAllEntries(m: Collection<Map.Entry<*, *>>): Boolean {
+        return m.all(this::containsOtherEntry)
    }
 }
@@ -10,6 +10,26 @@ package kotlin.collections

 import kotlin.collections.MutableMap.MutableEntry

+internal external interface JsRawArray<E> {
+    fun push(element: E)
+    fun pop(): E
+
+    var length: Int
+}
+
+private inline fun <E> JsRawArray<E>.getElement(index: Int): E {
+    return (this.asDynamic()[index]).unsafeCast<E>()
+}
+
+private inline fun <E> JsRawArray<E>.setElement(index: Int, element: E) {
+    this.asDynamic()[index] = element
+}
+
+private inline fun <E> JsRawArray<E>.replaceElementAtWithLast(index: Int) {
+    setElement(index, getElement(length - 1))
+    pop()
+}
+
 /**
 * A simple wrapper around JavaScript Map for key type is string.
 *
@@ -17,19 +37,25 @@ import kotlin.collections.MutableMap.MutableEntry
 * because we want to have it erased to Any? in order not to generate type-safe override bridges for
 * [get], [contains], [remove] etc, if they ever are generated.
 */
-internal class InternalStringMap<K, V> : InternalMap<K, V> {
+internal open class InternalStringMap<K, V> : InternalMap<K, V> {
+    private fun createJsMap(): dynamic {
+        val result = js("Object.create(null)")
+        // force to switch object representation to dictionary mode
+        result["foo"] = 1
+        jsDeleteProperty(result.unsafeCast<Any>(), "foo")
+        return result
+    }
+
+    private fun <E> createJsArray(): JsRawArray<E> {
+        return js("[]").unsafeCast<JsRawArray<E>>()
+    }

    private var backingMap: dynamic = createJsMap()
-    override var size: Int = 0
-        private set
+    private var values = createJsArray<V>()
+    private var keys = createJsArray<K>()

-//    /**
-//     * A mod count to track 'value' replacements in map to ensure that the 'value' that we have in the
-//     * iterator entry is guaranteed to be still correct.
-//     * This is to optimize for the common scenario where the values are not modified during
-//     * iterations where the entries are never stale.
-//     */
-//    private var valueMod: Int = 0
+    override val size: Int
+        get() = keys.length

    override operator fun contains(key: K): Boolean {
        if (key !is String) return false
@@ -38,77 +64,169 @@ internal class InternalStringMap<K, V> : InternalMap<K, V> {

    override operator fun get(key: K): V? {
        if (key !is String) return null
-        val value = backingMap[key]
-        return if (value !== undefined) value.unsafeCast<V>() else null
+        val index = backingMap[key]
+        return if (index !== undefined) values.getElement(index.unsafeCast<Int>()) else null
    }

+    override fun getEntry(entry: Map.Entry<K, V>): MutableEntry<K, V>? {
+        val key = entry.key as? String ?: return null
+        val index = backingMap[key]
+        if (index === undefined) {
+            return null
+        }
+        val value = values.getElement(index.unsafeCast<Int>())
+        if (value == entry.value) {
+            return EntryRef(entry.key, value, this)
+        }
+        return null
+    }
+
+    override fun containsValue(value: V): Boolean {
+        return values.unsafeCast<Array<V>>().contains(value)
+    }
+
+    override fun containsEntry(entry: Map.Entry<K, V>): Boolean {
+        return getEntry(entry) != null
+    }
+
+    override fun containsOtherEntry(entry: Map.Entry<*, *>): Boolean {
+        @Suppress("UNCHECKED_CAST")
+        return containsEntry(entry as Map.Entry<K, V>)
+    }
+
+    override fun removeEntry(entry: Map.Entry<K, V>): Boolean {
+        val key = getEntry(entry)?.key ?: return false
+        return remove(key) != null
+    }
+
+    override fun removeValue(value: V): Boolean {
+        val index = values.unsafeCast<Array<V>>().indexOf(value)
+        if (index < 0) {
+            return false
+        }
+        removeKeyIndex(keys.getElement(index), index)
+        return true
+    }

    override fun put(key: K, value: V): V? {
        require(key is String)
-        val oldValue = backingMap[key]
-        backingMap[key] = value
+        val index = backingMap[key]
+        if (index !== undefined) {
+            val i = index.unsafeCast<Int>()
+            val oldValue = values.getElement(i)
+            values.setElement(i, value)
+            return oldValue
+        }

-        if (oldValue === undefined) {
-            size++
-//            structureChanged(host)
-            return null
-        } else {
-//            valueMod++
-            return oldValue.unsafeCast<V>()
+        backingMap[key] = size
+        keys.push(key)
+        values.push(value)
+        return null
+    }
+
+    override fun putAll(from: Map<out K, V>) {
+        for ((key, value) in from) {
+            put(key, value)
        }
    }

    override fun remove(key: K): V? {
        if (key !is String) return null
-        val value = backingMap[key]
-        if (value !== undefined) {
-            jsDeleteProperty(backingMap, key)
-            size--
-//            structureChanged(host)
-            return value.unsafeCast<V>()
-        } else {
-//            valueMod++
+        val index = backingMap[key]
+        if (index === undefined) {
            return null
        }
+
+        val i = index.unsafeCast<Int>()
+        val removedValue = values.getElement(i)
+
+        removeKeyIndex(key, i)
+        return removedValue
    }

+    internal fun removeKeyIndex(key: K, index: Int) {
+        jsDeleteProperty(backingMap.unsafeCast<Any>(), key as Any)
+
+        if (index + 1 == size) {
+            keys.pop()
+            values.pop()
+        } else {
+            keys.replaceElementAtWithLast(index)
+            values.replaceElementAtWithLast(index)
+            backingMap[keys.getElement(index)] = index
+        }
+    }

    override fun clear() {
        backingMap = createJsMap()
-        size = 0
+        keys = createJsArray()
+        values = createJsArray()
    }

+    override fun keysIterator(): MutableIterator<K> = KeysItr(this)
+    override fun valuesIterator(): MutableIterator<V> = ValuesItr(this)
+    override fun entriesIterator(): MutableIterator<MutableEntry<K, V>> = EntriesItr(this)

-    override fun iterator(): MutableIterator<MutableEntry<K, V>> {
-        return object : MutableIterator<MutableEntry<K, V>> {
-            private val keys: Array<String> = js("Object").keys(backingMap)
-            private val iterator = keys.iterator()
-            private var lastKey: String? = null
+    private abstract class BaseItr<K, V>(protected val map: InternalStringMap<K, V>) {
+        protected var lastIndex = -1
+        protected var index = 0

-            override fun hasNext(): Boolean = iterator.hasNext()
-
-            override fun next(): MutableEntry<K, V> {
-                val key = iterator.next()
-                lastKey = key
-                @Suppress("UNCHECKED_CAST")
-                return newMapEntry(key as K)
+        protected fun goNext() {
+            if (index >= map.size) {
+                throw NoSuchElementException()
            }
+            lastIndex = index++
+        }

-            override fun remove() {
-                @Suppress("UNCHECKED_CAST")
-                this@InternalStringMap.remove(checkNotNull(lastKey) as K)
-            }
+        fun hasNext(): Boolean = index < map.size
+
+        fun remove() {
+            map.removeKeyIndex(map.keys.getElement(lastIndex), lastIndex)
+            lastIndex = -1
        }
    }

-    private fun newMapEntry(key: K): MutableEntry<K, V> = object : MutableEntry<K, V> {
-        override val key: K get() = key
-        override val value: V get() = this@InternalStringMap[key].unsafeCast<V>()
+    private abstract class Itr<I, K, V>(
+        private val iterableArray: JsRawArray<I>,
+        map: InternalStringMap<K, V>,
+    ) : MutableIterator<I>, BaseItr<K, V>(map) {
+        override fun next(): I {
+            goNext()
+            return iterableArray.getElement(lastIndex)
+        }
+    }

-        override fun setValue(newValue: V): V = this@InternalStringMap.put(key, newValue).unsafeCast<V>()
+    private class KeysItr<K, V>(map: InternalStringMap<K, V>) : Itr<K, K, V>(map.keys, map)
+    private class ValuesItr<K, V>(map: InternalStringMap<K, V>) : Itr<V, K, V>(map.values, map)

-        override fun hashCode(): Int = AbstractMap.entryHashCode(this)
-        override fun toString(): String = AbstractMap.entryToString(this)
-        override fun equals(other: Any?): Boolean = AbstractMap.entryEquals(this, other)
+    private class EntriesItr<K, V>(map: InternalStringMap<K, V>) : MutableIterator<MutableEntry<K, V>>, BaseItr<K, V>(map) {
+        override fun next(): MutableEntry<K, V> {
+            goNext()
+            val key = map.keys.getElement(lastIndex)
+            val value = map.values.getElement(lastIndex)
+            return EntryRef(key, value, map)
+        }
+    }
+
+    private class EntryRef<K, V>(
+        override val key: K,
+        override var value: V,
+        private val map: InternalStringMap<K, V>,
+    ) : MutableEntry<K, V> {
+        override fun setValue(newValue: V): V {
+            val prevValue = value
+            map.put(key, newValue)
+            value = newValue
+            return prevValue
+        }
+
+        override fun equals(other: Any?): Boolean =
+            other is Map.Entry<*, *> &&
+                    other.key == key &&
+                    other.value == value
+
+        override fun hashCode(): Int = key.hashCode() xor value.hashCode()
+
+        override fun toString(): String = "$key=$value"
    }
 }