[FIR] Make ImplicitReceiverStack persistent

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

@@ -5,8 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve
 
-import com.google.common.collect.LinkedHashMultimap
+import kotlinx.collections.immutable.*
-import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -21,34 +20,45 @@ interface ImplicitReceiverStack {
 
     fun lastDispatchReceiver(): ImplicitDispatchReceiverValue?
     fun receiversAsReversed(): List<ImplicitReceiverValue<*>>
+
+    fun snapshot(): ImplicitReceiverStack
 }
 
-class ImplicitReceiverStackImpl : ImplicitReceiverStack, Iterable<ImplicitReceiverValue<*>> {
+class ImplicitReceiverStackImpl private constructor(
-    private val stack: MutableList<ImplicitReceiverValue<*>> = mutableListOf()
+    private var stack: PersistentList<ImplicitReceiverValue<*>>,
-    private val originalTypes: MutableList<ConeKotlinType> = mutableListOf()
     // This multi-map holds indexes of the stack ^
-    private val indexesPerLabel: SetMultimap<Name, Int> = LinkedHashMultimap.create()
+    private var originalTypes: PersistentList<ConeKotlinType>,
-    private val indexesPerSymbol: MutableMap<FirBasedSymbol<*>, Int> = mutableMapOf()
+    private var indexesPerLabel: PersistentSetMultimap<Name, Int>,
+    private var indexesPerSymbol: PersistentMap<FirBasedSymbol<*>, Int>
+) : ImplicitReceiverStack, Iterable<ImplicitReceiverValue<*>> {
     val size: Int get() = stack.size
 
+    constructor() : this(
+        persistentListOf(),
+        persistentListOf(),
+        PersistentSetMultimap(),
+        persistentMapOf()
+    )
+
     override fun add(name: Name?, value: ImplicitReceiverValue<*>) {
-        stack += value
+        stack = stack.add(value)
-        originalTypes += value.type
+        originalTypes = originalTypes.add(value.type)
         val index = stack.size - 1
         if (name != null) {
-            indexesPerLabel.put(name, index)
+            indexesPerLabel = indexesPerLabel.put(name, index)
         }
-        indexesPerSymbol.put(value.boundSymbol, index)
+        indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)
     }
 
     override fun pop(name: Name?) {
         val index = stack.size - 1
         if (name != null) {
-            indexesPerLabel.remove(name, index)
+            indexesPerLabel = indexesPerLabel.remove(name, index)
         }
-        originalTypes.removeAt(index)
+        originalTypes = originalTypes.removeAt(index)
-        val value = stack.removeAt(index)
+        val value = stack.get(index)
-        indexesPerSymbol.remove(value.boundSymbol)
+        stack = stack.removeAt(index)
+        indexesPerSymbol = indexesPerSymbol.remove(value.boundSymbol)
     }
 
     override operator fun get(name: String?): ImplicitReceiverValue<*>? {
@@ -76,4 +86,37 @@ class ImplicitReceiverStackImpl : ImplicitReceiverStack, Iterable<ImplicitReceiv
     override operator fun iterator(): Iterator<ImplicitReceiverValue<*>> {
         return stack.iterator()
     }
+
+    override fun snapshot(): ImplicitReceiverStack {
+        return ImplicitReceiverStackImpl(stack, originalTypes, indexesPerLabel, indexesPerSymbol)
+    }
+}
+
+class PersistentSetMultimap<K, V> private constructor(private val map: PersistentMap<K, PersistentSet<V>>) {
+
+    constructor() : this(persistentMapOf())
+
+    fun put(key: K, value: V): PersistentSetMultimap<K, V> {
+        val set = map[key] ?: persistentSetOf()
+        val newSet = set.add(value)
+        if (newSet === set) return this
+        val newMap = map.put(key, newSet)
+        return PersistentSetMultimap(newMap)
+    }
+
+    fun remove(key: K, value: V): PersistentSetMultimap<K, V> {
+        val set = map.get(key) ?: return this
+        val newSet = set.remove(value)
+        if (set === newSet) return this
+        val newMap = if (newSet.isEmpty()) {
+            map.remove(key)
+        } else {
+            map.put(key, newSet)
+        }
+        return PersistentSetMultimap(newMap)
+    }
+
+    operator fun get(key: K): Set<V> {
+        return map[key] ?: emptySet()
+    }
 }