[FIR] Implement util Multimap classes

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentImplicitReceiverStack.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.util.PersistentSetMultimap
 import org.jetbrains.kotlin.name.Name
 
 class PersistentImplicitReceiverStack private constructor(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirLocalScope.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.FirVariable
-import org.jetbrains.kotlin.fir.resolve.PersistentMultimap
+import org.jetbrains.kotlin.fir.util.PersistentMultimap
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.scopes.FirScope

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/util/ChainedIterator.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.util
+
+class ChainedIterator<T>(delegates: Collection<Iterator<T>>) : Iterator<T> {
+    private var metaIterator = delegates.iterator()
+    private var currentIterator: Iterator<T>? = null
+
+    private fun promote() {
+        if (currentIterator?.hasNext() == true) return
+        while (metaIterator.hasNext()) {
+            currentIterator = metaIterator.next()
+            if (currentIterator!!.hasNext()) return
+        }
+    }
+
+    override fun hasNext(): Boolean {
+        promote()
+        return currentIterator?.hasNext() == true
+    }
+
+    override fun next(): T {
+        promote()
+        return currentIterator?.next() ?: throw NoSuchElementException()
+    }
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/util/Multimap.kt

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.util
+
+interface Multimap<K, out V, out C : Collection<V>> {
+    operator fun get(key: K): C
+    operator fun contains(key: K): Boolean
+    val keys: Set<K>
+    val values: Collection<V>
+}
+
+interface MutableMultimap<K, V, C : Collection<V>> : Multimap<K, V, C> {
+    fun put(key: K, value: V)
+    fun putAll(key: K, values: Collection<V>) {
+        values.forEach { put(key, it) }
+    }
+
+    fun remove(key: K, value: V)
+    fun removeKey(key: K)
+
+    fun clear()
+}
+
+abstract class BaseMultimap<K, V, C : Collection<V>, MC : MutableCollection<V>> : MutableMultimap<K, V, C> {
+    private val map: MutableMap<K, MC> = mutableMapOf()
+    protected abstract fun createContainer(): MC
+    protected abstract fun createEmptyContainer(): C
+
+    override fun get(key: K): C {
+        @Suppress("UNCHECKED_CAST")
+        return map[key] as C? ?: createEmptyContainer()
+    }
+
+    override operator fun contains(key: K): Boolean {
+        return key in map
+    }
+
+    override val keys: Set<K>
+        get() = map.keys
+
+    override val values: Collection<V>
+        get() = object : AbstractCollection<V>() {
+            override val size: Int
+                get() = map.values.sumBy { it.size }
+
+            override fun iterator(): Iterator<V> {
+                return ChainedIterator(map.values.map { it.iterator() })
+            }
+        }
+
+    override fun put(key: K, value: V) {
+        val container = map.getOrPut(key) { createContainer() }
+        container.add(value)
+    }
+
+    override fun remove(key: K, value: V) {
+        val collection = map[key] ?: return
+        collection.remove(value)
+        if (collection.isEmpty()) {
+            map.remove(key)
+        }
+    }
+
+    override fun removeKey(key: K) {
+        map.remove(key)
+    }
+
+    override fun clear() {
+        map.clear()
+    }
+}
+
+class SetMultimap<K, V> : BaseMultimap<K, V, Set<V>, MutableSet<V>>() {
+    override fun createContainer(): MutableSet<V> {
+        return mutableSetOf()
+    }
+
+    override fun createEmptyContainer(): Set<V> {
+        return emptySet()
+    }
+}
+
+class ListMultimap<K, V> : BaseMultimap<K, V, List<V>, MutableList<V>>() {
+    override fun createContainer(): MutableList<V> {
+        return mutableListOf()
+    }
+
+    override fun createEmptyContainer(): List<V> {
+        return emptyList()
+    }
+}
+
+fun <K, V> setMultimapOf(): SetMultimap<K, V> = SetMultimap()
+fun <K, V> listMultimapOf(): ListMultimap<K, V> = ListMultimap()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/util/PersistentMultimap.kt

@@ -3,7 +3,7 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.fir.resolve
+package org.jetbrains.kotlin.fir.util
 
 import kotlinx.collections.immutable.*