diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DataFlowVariable.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DataFlowVariable.kt index b8327872bb3..19adc6669aa 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DataFlowVariable.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DataFlowVariable.kt @@ -5,8 +5,6 @@ package org.jetbrains.kotlin.fir.resolve.dfa -import com.google.common.collect.LinkedHashMultimap -import com.google.common.collect.Multimap import org.jetbrains.kotlin.fir.FirElement import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.declarations.FirTypedDeclaration @@ -31,7 +29,7 @@ data class DataFlowVariable( } class DataFlowVariableStorage(private val session: FirSession) { - private val dfi2FirMap: Multimap = LinkedHashMultimap.create() + private val dfi2FirMap: MutableMap = mutableMapOf() private val fir2DfiMap: MutableMap = mutableMapOf() private var counter: Int = 1 @@ -57,12 +55,13 @@ class DataFlowVariableStorage(private val session: FirSession) { removeVariable(variable) } - fun removeVariable(variable: DataFlowVariable) { - val firExpressions = dfi2FirMap.removeAll(variable) - firExpressions.forEach(fir2DfiMap::remove) + fun removeVariable(variable: DataFlowVariable): FirElement? { + return dfi2FirMap.remove(variable)?.also { + fir2DfiMap.remove(it) + } } - operator fun get(variable: DataFlowVariable): Collection { + operator fun get(variable: DataFlowVariable): FirElement? { return dfi2FirMap[variable] } @@ -81,12 +80,12 @@ class DataFlowVariableStorage(private val session: FirSession) { } private fun storeVariable(variable: DataFlowVariable, fir: FirElement) { - dfi2FirMap.put(variable, fir) - fir2DfiMap.put(fir, variable) + dfi2FirMap[variable] = fir + fir2DfiMap[fir] = variable } @Deprecated("only for debug") - fun getByName(name: String): DataFlowVariable? = dfi2FirMap.keySet().firstOrNull { it.name == name } + fun getByName(name: String): DataFlowVariable? = dfi2FirMap.keys.firstOrNull { it.name == name } private val FirElement.type: FirTypeRef get() = when (this) {