diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ImplicitReceiverStack.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ImplicitReceiverStack.kt index b2646f7d45f..ca541f62053 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ImplicitReceiverStack.kt +++ b/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 com.google.common.collect.SetMultimap +import kotlinx.collections.immutable.* 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> + + fun snapshot(): ImplicitReceiverStack } -class ImplicitReceiverStackImpl : ImplicitReceiverStack, Iterable> { - private val stack: MutableList> = mutableListOf() - private val originalTypes: MutableList = mutableListOf() +class ImplicitReceiverStackImpl private constructor( + private var stack: PersistentList>, // This multi-map holds indexes of the stack ^ - private val indexesPerLabel: SetMultimap = LinkedHashMultimap.create() - private val indexesPerSymbol: MutableMap, Int> = mutableMapOf() + private var originalTypes: PersistentList, + private var indexesPerLabel: PersistentSetMultimap, + private var indexesPerSymbol: PersistentMap, Int> +) : ImplicitReceiverStack, Iterable> { val size: Int get() = stack.size + constructor() : this( + persistentListOf(), + persistentListOf(), + PersistentSetMultimap(), + persistentMapOf() + ) + override fun add(name: Name?, value: ImplicitReceiverValue<*>) { - stack += value - originalTypes += value.type + stack = stack.add(value) + 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) - val value = stack.removeAt(index) - indexesPerSymbol.remove(value.boundSymbol) + originalTypes = originalTypes.removeAt(index) + val value = stack.get(index) + stack = stack.removeAt(index) + indexesPerSymbol = indexesPerSymbol.remove(value.boundSymbol) } override operator fun get(name: String?): ImplicitReceiverValue<*>? { @@ -76,4 +86,37 @@ class ImplicitReceiverStackImpl : ImplicitReceiverStack, Iterable> { return stack.iterator() } + + override fun snapshot(): ImplicitReceiverStack { + return ImplicitReceiverStackImpl(stack, originalTypes, indexesPerLabel, indexesPerSymbol) + } +} + +class PersistentSetMultimap private constructor(private val map: PersistentMap>) { + + constructor() : this(persistentMapOf()) + + fun put(key: K, value: V): PersistentSetMultimap { + 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 { + 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 { + return map[key] ?: emptySet() + } } \ No newline at end of file