FIR: Move some responsibilities of TowerResolveManager to FirTowerResolverSession

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

@@ -261,11 +261,12 @@ class FirCallResolver(
             expectedType, constraintSystemBuilder,
         )
         // No reset here!
+        val localCollector = CandidateCollector(this, resolutionStageRunner)
         val result = towerResolver.runResolver(
             implicitReceiverStack.receiversAsReversed(),
             info,
-            collector = CandidateCollector(this, resolutionStageRunner),
+            collector = localCollector,
-            manager = TowerResolveManager(),
+            manager = TowerResolveManager(localCollector),
         )
         val bestCandidates = result.bestCandidates()
         val noSuccessfulCandidates = result.currentApplicability < CandidateApplicability.SYNTHETIC_RESOLVED

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirTowerResolver.kt

@@ -14,7 +14,7 @@ class FirTowerResolver(
     resolutionStageRunner: ResolutionStageRunner,
 ) {
     private val collector = CandidateCollector(components, resolutionStageRunner)
-    private val manager = TowerResolveManager()
+    private val manager = TowerResolveManager(collector)
 
     fun runResolver(
         implicitReceiverValues: List<ImplicitReceiverValue<*>>,
@@ -22,9 +22,9 @@ class FirTowerResolver(
         collector: CandidateCollector = this.collector,
         manager: TowerResolveManager = this.manager
     ): CandidateCollector {
-        val towerResolverSession = FirTowerResolverSession(components, implicitReceiverValues, manager)
+        val callResolutionContext = buildCallResolutionContext(info, collector)
-        manager.callResolutionContext = buildCallResolutionContext(info, collector, towerResolverSession)
 
+        val towerResolverSession = FirTowerResolverSession(components, implicitReceiverValues, manager, callResolutionContext)
         towerResolverSession.runResolution(info)
 
         manager.runTasks()
@@ -33,8 +33,7 @@ class FirTowerResolver(
 
     private fun buildCallResolutionContext(
         info: CallInfo,
-        collector: CandidateCollector,
+        collector: CandidateCollector
-        towerResolverSession: FirTowerResolverSession
     ): CallResolutionContext {
         val candidateFactory = CandidateFactory(components, info)
         val stubReceiverCandidateFactory =
@@ -63,8 +62,7 @@ class FirTowerResolver(
             invokeBuiltinExtensionReceiverCandidateFactory,
             stubReceiverCandidateFactory,
             invokeReceiverCollector,
-            this,
+            this
-            towerResolverSession
         )
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirTowerResolverSession.kt

@@ -27,10 +27,11 @@ import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.descriptorUtil.HIDES_MEMBERS_NAME_LIST
 
-class FirTowerResolverSession(
+class FirTowerResolverSession internal constructor(
     val components: BodyResolveComponents,
     implicitReceiverValues: List<ImplicitReceiverValue<*>>,
-    private val manager: TowerResolveManager
+    private val manager: TowerResolveManager,
+    private val callResolutionContext: CallResolutionContext,
 ) {
     private data class ImplicitReceiver(
         val receiver: ImplicitReceiverValue<*>,
@@ -73,6 +74,35 @@ class FirTowerResolverSession(
         }
     }
 
+    private suspend fun processLevel(
+        towerLevel: SessionBasedTowerLevel,
+        callInfo: CallInfo,
+        group: TowerGroup,
+        explicitReceiverKind: ExplicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
+        invokeResolveMode: InvokeResolveMode? = null
+    ): ProcessorAction {
+        manager.requestGroup(group)
+
+        val result = with(LevelHandler(callInfo, explicitReceiverKind, group, callResolutionContext, manager, this)) {
+            towerLevel.handleLevel(invokeResolveMode)
+        }
+
+        manager.stopIfSuccess()
+
+        return result
+    }
+
+    private suspend fun processLevelForPropertyWithInvoke(
+        towerLevel: SessionBasedTowerLevel,
+        callInfo: CallInfo,
+        group: TowerGroup,
+        explicitReceiverKind: ExplicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
+    ) {
+        if (callInfo.callKind == CallKind.Function) {
+            processLevel(towerLevel, callInfo, group, explicitReceiverKind, InvokeResolveMode.RECEIVER_FOR_INVOKE_BUILTIN_EXTENSION)
+        }
+    }
+
     private fun FirScope.toScopeTowerLevel(
         extensionReceiver: ReceiverValue? = null,
         extensionsOnly: Boolean = false,
@@ -96,7 +126,7 @@ class FirTowerResolverSession(
     ) {
         if (qualifierReceiver == null) return
         for ((depth, qualifierScope) in qualifierReceiver.callableScopes().withIndex()) {
-            manager.processLevel(
+            processLevel(
                 qualifierScope.toScopeTowerLevel(noInnerConstructors = true),
                 info.noStubReceiver(), TowerGroup.Qualifier(depth)
             )
@@ -112,7 +142,7 @@ class FirTowerResolverSession(
             qualifierReceiver.classSymbol != qualifierReceiver.originalSymbol
         ) return
         val scope = qualifierReceiver.classifierScope() ?: return
-        manager.processLevel(
+        processLevel(
             scope.toScopeTowerLevel(noInnerConstructors = true), info.noStubReceiver(),
             if (prioritized) TowerGroup.ClassifierPrioritized else TowerGroup.Classifier
         )
@@ -189,7 +219,7 @@ class FirTowerResolverSession(
         info: CallInfo,
         index: Int,
         depth: Int?
-    ) = manager.processLevel(
+    ) = processLevel(
         topLevelScope.toScopeTowerLevel(
             extensionReceiver = receiverValue, extensionsOnly = true
         ),
@@ -201,7 +231,7 @@ class FirTowerResolverSession(
         nonEmptyLocalScopes: MutableList<FirLocalScope>
     ) {
         for ((index, localScope) in localScopes.withIndex()) {
-            val result = manager.processLevel(
+            val result = processLevel(
                 localScope.toScopeTowerLevel(), info, TowerGroup.Local(index)
             )
             if (result != ProcessorAction.NONE) {
@@ -234,7 +264,7 @@ class FirTowerResolverSession(
             if (implicitReceiverValue is ImplicitDispatchReceiverValue) {
                 val scope = implicitReceiverValue.scope(session, components.scopeSession)
                 if (scope != null) {
-                    manager.processLevel(
+                    processLevel(
                         FirStaticScope(scope).toScopeTowerLevel(),
                         info, parentGroup.Static(depth)
                     )
@@ -251,7 +281,7 @@ class FirTowerResolverSession(
         nonEmptyLocalScopes: List<FirLocalScope>,
         emptyTopLevelScopes: MutableSet<FirScope>
     ) {
-        val implicitResult = manager.processLevel(
+        val implicitResult = processLevel(
             receiver.toMemberScopeTowerLevel(), info, parentGroup.Member
         )
 
@@ -260,7 +290,7 @@ class FirTowerResolverSession(
         }
 
         for ((localIndex, localScope) in nonEmptyLocalScopes.withIndex()) {
-            manager.processLevel(
+            processLevel(
                 localScope.toScopeTowerLevel(extensionReceiver = receiver),
                 info, parentGroup.Local(localIndex)
             )
@@ -268,7 +298,7 @@ class FirTowerResolverSession(
 
         for ((implicitDispatchReceiverValue, dispatchDepth) in implicitReceiversUsableAsValues) {
             if (implicitDispatchReceiverValue in implicitReceiverValuesWithEmptyScopes) continue
-            manager.processLevel(
+            processLevel(
                 implicitDispatchReceiverValue.toMemberScopeTowerLevel(extensionReceiver = receiver),
                 info, parentGroup.Implicit(dispatchDepth)
             )
@@ -276,7 +306,7 @@ class FirTowerResolverSession(
 
         for ((topIndex, topLevelScope) in topLevelScopes.withIndex()) {
             if (topLevelScope in emptyTopLevelScopes) continue
-            val result = manager.processLevel(
+            val result = processLevel(
                 topLevelScope.toScopeTowerLevel(extensionReceiver = receiver),
                 info, parentGroup.Top(topIndex)
             )
@@ -294,7 +324,7 @@ class FirTowerResolverSession(
             // NB: this check does not work for variables
             // because we do not search for objects if we have extension receiver
             if (info.callKind != CallKind.VariableAccess && topLevelScope in emptyTopLevelScopes) continue
-            manager.processLevel(
+            processLevel(
                 topLevelScope.toScopeTowerLevel(), info, TowerGroup.Top(index)
             )
         }
@@ -349,12 +379,12 @@ class FirTowerResolverSession(
         info: CallInfo,
         towerGroup: TowerGroup,
     ): ProcessorAction {
-        val result = manager.processLevel(
+        val result = processLevel(
             scope.toScopeTowerLevel(extensionReceiver = explicitReceiverValue),
             info, towerGroup, ExplicitReceiverKind.EXTENSION_RECEIVER
         )
 
-        manager.processLevelForPropertyWithInvoke(
+        processLevelForPropertyWithInvoke(
             scope.toScopeTowerLevel(), info, towerGroup
         )
 
@@ -365,7 +395,7 @@ class FirTowerResolverSession(
         explicitReceiverValue: ExpressionReceiverValue,
         info: CallInfo
     ) {
-        manager.processLevel(
+        processLevel(
             explicitReceiverValue.toMemberScopeTowerLevel(), info, TowerGroup.Member, ExplicitReceiverKind.DISPATCH_RECEIVER
         )
     }
@@ -381,31 +411,31 @@ class FirTowerResolverSession(
         val parentGroup = TowerGroup.Implicit(depth)
 
         // Member extensions
-        manager.processLevel(
+        processLevel(
             implicitReceiverValue.toMemberScopeTowerLevel(extensionReceiver = explicitReceiverValue),
             info, parentGroup.Member, ExplicitReceiverKind.EXTENSION_RECEIVER
         )
         // properties for invoke on implicit receiver
-        manager.processLevelForPropertyWithInvoke(
+        processLevelForPropertyWithInvoke(
             implicitReceiverValue.toMemberScopeTowerLevel(), info, parentGroup.Member
         )
 
         for ((localIndex, localScope) in nonEmptyLocalScopes.withIndex()) {
-            manager.processLevelForPropertyWithInvoke(
+            processLevelForPropertyWithInvoke(
                 localScope.toScopeTowerLevel(extensionReceiver = implicitReceiverValue),
                 info, parentGroup.Local(localIndex)
             )
         }
 
         for ((implicitDispatchReceiverValue, dispatchDepth) in implicitReceiversUsableAsValues) {
-            manager.processLevelForPropertyWithInvoke(
+            processLevelForPropertyWithInvoke(
                 implicitDispatchReceiverValue.toMemberScopeTowerLevel(extensionReceiver = implicitReceiverValue),
                 info, parentGroup.Implicit(dispatchDepth)
             )
         }
 
         for ((topIndex, topLevelScope) in topLevelScopes.withIndex()) {
-            manager.processLevelForPropertyWithInvoke(
+            processLevelForPropertyWithInvoke(
                 topLevelScope.toScopeTowerLevel(extensionReceiver = implicitReceiverValue),
                 info, parentGroup.Top(topIndex)
             )
@@ -423,7 +453,7 @@ class FirTowerResolverSession(
             )
             else -> null
         } ?: return
-        manager.processLevel(
+        processLevel(
             scope.toScopeTowerLevel(),
             info, TowerGroup.Member, explicitReceiverKind = ExplicitReceiverKind.DISPATCH_RECEIVER
         )
@@ -433,7 +463,7 @@ class FirTowerResolverSession(
         info: CallInfo,
         invokeReceiverValue: ExpressionReceiverValue
     ) {
-        manager.processLevel(
+        processLevel(
             invokeReceiverValue.toMemberScopeTowerLevel(),
             info, TowerGroup.Member,
             ExplicitReceiverKind.DISPATCH_RECEIVER,
@@ -441,7 +471,7 @@ class FirTowerResolverSession(
         )
 
         for ((index, localScope) in localScopes.withIndex()) {
-            manager.processLevel(
+            processLevel(
                 localScope.toScopeTowerLevel(extensionReceiver = invokeReceiverValue),
                 info, TowerGroup.Local(index),
                 ExplicitReceiverKind.EXTENSION_RECEIVER,
@@ -452,7 +482,7 @@ class FirTowerResolverSession(
         for ((implicitReceiverValue, depth) in implicitReceiversUsableAsValues) {
             // NB: companions are processed via implicitReceiverValues!
             val parentGroup = TowerGroup.Implicit(depth)
-            manager.processLevel(
+            processLevel(
                 implicitReceiverValue.toMemberScopeTowerLevel(extensionReceiver = invokeReceiverValue),
                 info, parentGroup.Member,
                 ExplicitReceiverKind.EXTENSION_RECEIVER,
@@ -461,7 +491,7 @@ class FirTowerResolverSession(
         }
 
         for ((index, topLevelScope) in topLevelScopes.withIndex()) {
-            manager.processLevel(
+            processLevel(
                 topLevelScope.toScopeTowerLevel(extensionReceiver = invokeReceiverValue),
                 info, TowerGroup.Top(index),
                 ExplicitReceiverKind.EXTENSION_RECEIVER,
@@ -475,7 +505,7 @@ class FirTowerResolverSession(
         info: CallInfo,
         invokeReceiverValue: ExpressionReceiverValue
     ) {
-        manager.processLevel(
+        processLevel(
             invokeReceiverValue.toMemberScopeTowerLevel(),
             info, TowerGroup.Member,
             ExplicitReceiverKind.DISPATCH_RECEIVER,
@@ -490,7 +520,7 @@ class FirTowerResolverSession(
     ) {
         for ((implicitReceiverValue, depth) in implicitReceiversUsableAsValues) {
             val parentGroup = TowerGroup.Implicit(depth)
-            manager.processLevel(
+            processLevel(
                 invokeReceiverValue.toMemberScopeTowerLevel(
                     extensionReceiver = implicitReceiverValue,
                     implicitExtensionInvokeMode = true

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/TowerLevelHandler.kt

@@ -31,8 +31,7 @@ internal class CallResolutionContext(
     val invokeBuiltinExtensionReceiverCandidateFactory: CandidateFactory?,
     val stubReceiverCandidateFactory: CandidateFactory?,
     val invokeReceiverCollector: CandidateCollector?,
-    val towerResolver: FirTowerResolver,
+    val towerResolver: FirTowerResolver
-    val towerResolverSession: FirTowerResolverSession
 ) {
     // TODO: Get rid of the property, storing state here looks like a hack
     internal lateinit var invokeOnGivenReceiverCandidateFactory: CandidateFactory
@@ -43,12 +42,12 @@ internal class LevelHandler(
     private val explicitReceiverKind: ExplicitReceiverKind,
     private val group: TowerGroup,
     private val callResolutionContext: CallResolutionContext,
-    private val manager: TowerResolveManager
+    private val manager: TowerResolveManager,
+    private val towerResolverSession: FirTowerResolverSession
 ) {
 
     private val resultCollector: CandidateCollector get() = callResolutionContext.resultCollector
     private val towerResolver: FirTowerResolver get() = callResolutionContext.towerResolver
-    private val towerResolverSession: FirTowerResolverSession get() = callResolutionContext.towerResolverSession
 
     private fun enqueueResolverTask(group: TowerGroup = TowerGroup.Start, task: suspend () -> Unit) =
         manager.enqueueResolverTask(group, task)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/TowerResolveManager.kt

@@ -5,27 +5,22 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls
 
-import org.jetbrains.kotlin.fir.scopes.ProcessorAction
-import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import java.util.*
 import kotlin.coroutines.*
 
-class TowerResolveManager internal constructor() {
+class TowerResolveManager private constructor(private val isSuccess: () -> Boolean) {
+
+    constructor(collector: CandidateCollector) : this({ collector.isSuccess() })
 
     private val queue = PriorityQueue<SuspendedResolverTask>()
 
-    // TODO: Get rid of the property (it should be passed in a different way)
-    internal lateinit var callResolutionContext: CallResolutionContext
-
-    private fun isSuccess() = callResolutionContext.resultCollector.isSuccess()
-
     fun reset() {
         queue.clear()
     }
 
     private suspend fun suspendResolverTask(group: TowerGroup) = suspendCoroutine<Unit> { queue += SuspendedResolverTask(it, group) }
 
-    private suspend fun requestGroup(requested: TowerGroup) {
+    suspend fun requestGroup(requested: TowerGroup) {
         val peeked = queue.peek()
 
         // Task ordering should be FIFO
@@ -37,33 +32,8 @@ class TowerResolveManager internal constructor() {
 
     private suspend fun stopResolverTask(): Nothing = suspendCoroutine { }
 
-    suspend fun processLevel(
+    suspend fun stopIfSuccess() {
-        towerLevel: SessionBasedTowerLevel,
-        callInfo: CallInfo,
-        group: TowerGroup,
-        explicitReceiverKind: ExplicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
-        invokeResolveMode: InvokeResolveMode? = null
-    ): ProcessorAction {
-        requestGroup(group)
-
-        val result = with(LevelHandler(callInfo, explicitReceiverKind, group, callResolutionContext, this)) {
-            towerLevel.handleLevel(invokeResolveMode)
-        }
-
         if (isSuccess()) stopResolverTask()
-
-        return result
-    }
-
-    suspend fun processLevelForPropertyWithInvoke(
-        towerLevel: SessionBasedTowerLevel,
-        callInfo: CallInfo,
-        group: TowerGroup,
-        explicitReceiverKind: ExplicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
-    ) {
-        if (callInfo.callKind == CallKind.Function) {
-            processLevel(towerLevel, callInfo, group, explicitReceiverKind, InvokeResolveMode.RECEIVER_FOR_INVOKE_BUILTIN_EXTENSION)
-        }
     }
 
     private data class SuspendedResolverTask(