FIR resolve: integrate implicit receivers into ImplicitReceiverStack

It can work both by order & by name

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

@@ -11,11 +11,10 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedQualifierImpl
-import org.jetbrains.kotlin.fir.expressions.impl.FirThisReceiverExpressionImpl
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreNameReference
@@ -32,14 +31,13 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
 
 class FirCallResolver(
     private val transformer: FirBodyResolveTransformer,
     private val topLevelScopes: List<FirScope>,
     private val localScopes: List<FirLocalScope>,
-    private val implicitReceiverStack: List<ImplicitReceiverValue<*>>,
+    override val implicitReceiverStack: ImplicitReceiverStack,
     private val qualifiedResolver: FirQualifiedNameResolver
 ) : BodyResolveComponents by transformer {
 
@@ -72,7 +70,7 @@ class FirCallResolver(
         )
 
         val consumer = createFunctionConsumer(session, name, info, inferenceComponents, resolver.collector, resolver)
-        val result = resolver.runTowerResolver(consumer, implicitReceiverStack.asReversed())
+        val result = resolver.runTowerResolver(consumer, implicitReceiverStack.receiversAsReversed())
         val bestCandidates = result.bestCandidates()
         val reducedCandidates = if (result.currentApplicability < CandidateApplicability.SYNTHETIC_RESOLVED) {
             bestCandidates.toSet()
@@ -171,7 +169,7 @@ class FirCallResolver(
             info, inferenceComponents,
             resolver.collector
         )
-        val result = resolver.runTowerResolver(consumer, implicitReceiverStack.asReversed())
+        val result = resolver.runTowerResolver(consumer, implicitReceiverStack.receiversAsReversed())
 
         val candidates = result.bestCandidates()
         val nameReference = createResolvedNamedReference(

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

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2019 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.resolve
+
+import com.google.common.collect.LinkedHashMultimap
+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.name.Name
+
+class ImplicitReceiverStack {
+    private val stack: MutableList<ImplicitReceiverValue<*>> = mutableListOf()
+    // This multi-map holds indexes of the stack ^
+    private val indexesPerLabel: SetMultimap<Name, Int> = LinkedHashMultimap.create()
+
+    fun add(name: Name, value: ImplicitReceiverValue<*>) {
+        stack += value
+        indexesPerLabel.put(name, stack.size - 1)
+    }
+
+    fun pop(name: Name) {
+        indexesPerLabel.remove(name, stack.size - 1)
+        stack.removeAt(stack.size - 1)
+    }
+
+    operator fun get(name: String?): ImplicitReceiverValue<*>? {
+        if (name == null) return stack.lastOrNull()
+        return indexesPerLabel[Name.identifier(name)].lastOrNull()?.let { stack[it] }
+    }
+
+    fun lastDispatchReceiver(): ImplicitDispatchReceiverValue? {
+        return stack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
+    }
+
+    fun receiversAsReversed(): List<ImplicitReceiverValue<*>> = stack.asReversed()
+}

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

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.*
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -274,8 +273,7 @@ fun <T : FirResolvable> BodyResolveComponents.typeFromCallee(access: T): FirReso
         }
         is FirThisReference -> {
             val labelName = newCallee.labelName
-            val implicitReceivers = if (labelName == null) implicitReceiverPerLabel.values() else implicitReceiverPerLabel[Name.identifier(labelName)]
-            val implicitReceiver = implicitReceivers.lastOrNull()
+            val implicitReceiver = implicitReceiverStack[labelName]
             FirResolvedTypeRefImpl(null, implicitReceiver?.type ?: ConeKotlinErrorType("Unresolved this@$labelName"), emptyList())
         }
         else -> error("Failed to extract type from: $newCallee")

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

@@ -5,21 +5,17 @@
 
 package org.jetbrains.kotlin.fir.resolve
 
-import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.phasedFir
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.FirSymbolOwner
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.name.Name
 
 interface SessionHolder {
     val session: FirSession
@@ -27,7 +23,7 @@ interface SessionHolder {
 
 interface BodyResolveComponents : SessionHolder {
     val returnTypeCalculator: ReturnTypeCalculator
-    val implicitReceiverPerLabel: SetMultimap<Name, ImplicitReceiverValue<*>>
+    val implicitReceiverStack: ImplicitReceiverStack
     val noExpectedType: FirTypeRef
     val symbolProvider: FirSymbolProvider
     val file: FirFile

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
-import com.google.common.collect.LinkedHashMultimap
-import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
@@ -67,8 +65,7 @@ open class FirBodyResolveTransformer(
 
     private val localScopes = mutableListOf<FirLocalScope>()
     private val topLevelScopes = mutableListOf<FirScope>()
-    final override val implicitReceiverPerLabel: SetMultimap<Name, ImplicitReceiverValue<*>> = LinkedHashMultimap.create()
-    private val implicitReceiverStack = mutableListOf<ImplicitReceiverValue<*>>()
+    final override val implicitReceiverStack: ImplicitReceiverStack = ImplicitReceiverStack()
     final override val inferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
 
     private var primaryConstructorParametersScope: FirLocalScope? = null
@@ -200,9 +197,7 @@ open class FirBodyResolveTransformer(
         when (val callee = qualifiedAccessExpression.calleeReference) {
             is FirExplicitThisReference -> {
                 val labelName = callee.labelName
-                val implicitReceivers =
-                    if (labelName == null) implicitReceiverPerLabel.values() else implicitReceiverPerLabel[Name.identifier(labelName)]
-                val implicitReceiver = implicitReceivers.lastOrNull()
+                val implicitReceiver = implicitReceiverStack[labelName]
                 callee.boundSymbol = implicitReceiver?.boundSymbol
                 qualifiedAccessExpression.resultType = FirResolvedTypeRefImpl(
                     null, implicitReceiver?.type ?: ConeKotlinErrorType("Unresolved this@$labelName"),
@@ -214,7 +209,7 @@ open class FirBodyResolveTransformer(
                 if (callee.superTypeRef is FirResolvedTypeRef) {
                     qualifiedAccessExpression.resultType = callee.superTypeRef
                 } else {
-                    val superTypeRef = implicitReceiverStack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
+                    val superTypeRef = implicitReceiverStack.lastDispatchReceiver()
                         ?.boundSymbol?.phasedFir?.superTypeRefs?.firstOrNull()
                         ?: FirErrorTypeRefImpl(qualifiedAccessExpression.psi, "No super type")
                     qualifiedAccessExpression.resultType = superTypeRef
@@ -793,11 +788,9 @@ open class FirBodyResolveTransformer(
                 throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
             }
         }
-        implicitReceiverPerLabel.put(labelName, implicitReceiverValue)
-        implicitReceiverStack += implicitReceiverValue
+        implicitReceiverStack.add(labelName, implicitReceiverValue)
         val result = block()
-        implicitReceiverStack.removeAt(implicitReceiverStack.size - 1)
-        implicitReceiverPerLabel.remove(labelName, implicitReceiverValue)
+        implicitReceiverStack.pop(labelName)
         return result
     }