[FIR] Extract mutable state from InferenceComponents

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

@@ -146,9 +146,7 @@ class FirCallResolver(
             transformer.components.containingDeclarations,
         )
         towerResolver.reset()
-        val result = towerResolver.runResolver(
+        val result = towerResolver.runResolver(info, transformer.resolutionContext)
-            info,
-        )
         val bestCandidates = result.bestCandidates()
         val reducedCandidates = if (!result.currentApplicability.isSuccess) {
             bestCandidates.toSet()
@@ -256,6 +254,7 @@ class FirCallResolver(
         val localCollector = CandidateCollector(this, resolutionStageRunner)
         val result = towerResolver.runResolver(
             info,
+            transformer.resolutionContext,
             collector = localCollector,
             manager = TowerResolveManager(localCollector),
         )
@@ -317,7 +316,8 @@ class FirCallResolver(
 
         val result = towerResolver.runResolverForDelegatingConstructor(
             callInfo,
-            constructedType
+            constructedType,
+            transformer.resolutionContext
         )
 
         return callResolver.selectDelegatingConstructorCall(delegatedConstructorCall, name, result, callInfo)
@@ -396,11 +396,11 @@ class FirCallResolver(
             }
         }
         if (constructorSymbol == null) return null
-        val candidate = CandidateFactory(components, callInfo).createCandidate(
+        val candidate = CandidateFactory(transformer.resolutionContext, callInfo).createCandidate(
             constructorSymbol!!,
             ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
         )
-        val applicability = resolutionStageRunner.processCandidate(candidate)
+        val applicability = resolutionStageRunner.processCandidate(candidate, transformer.resolutionContext)
         return ResolutionResult(callInfo, applicability, listOf(candidate))
     }
 
@@ -541,8 +541,8 @@ class FirCallResolver(
         source: FirSourceElement?,
         name: Name
     ): FirErrorReferenceWithCandidate {
-        val candidate = CandidateFactory(components, callInfo).createErrorCandidate(diagnostic)
+        val candidate = CandidateFactory(transformer.resolutionContext, callInfo).createErrorCandidate(diagnostic)
-        resolutionStageRunner.processCandidate(candidate, stopOnFirstError = false)
+        resolutionStageRunner.processCandidate(candidate, transformer.resolutionContext, stopOnFirstError = false)
         return FirErrorReferenceWithCandidate(source, name, candidate, diagnostic)
     }
 
@@ -552,7 +552,7 @@ class FirCallResolver(
         diagnostic: ConeDiagnostic
     ): FirErrorReferenceWithCandidate {
         if (!candidate.fullyAnalyzed) {
-            resolutionStageRunner.processCandidate(candidate, stopOnFirstError = false)
+            resolutionStageRunner.processCandidate(candidate, transformer.resolutionContext, stopOnFirstError = false)
         }
         return FirErrorReferenceWithCandidate(source, candidate.callInfo.name, candidate, diagnostic)
     }

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

@@ -10,7 +10,10 @@ import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
@@ -56,8 +59,6 @@ interface BodyResolveComponents : SessionHolder {
     val integerLiteralTypeApproximator: IntegerLiteralTypeApproximationTransformer
     val integerOperatorsTypeUpdater: IntegerOperatorsTypeUpdater
     val outerClassManager: FirOuterClassManager
-
-    val resolutionContext: ResolutionContext
 }
 
 typealias FirLocalScopes = PersistentList<FirLocalScope>

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

@@ -29,8 +29,8 @@ open class CandidateCollector(
         bestGroup = TowerGroup.Last
     }
 
-    open fun consumeCandidate(group: TowerGroup, candidate: Candidate): CandidateApplicability {
+    open fun consumeCandidate(group: TowerGroup, candidate: Candidate, context: ResolutionContext): CandidateApplicability {
-        val applicability = resolutionStageRunner.processCandidate(candidate)
+        val applicability = resolutionStageRunner.processCandidate(candidate, context)
 
         if (applicability > currentApplicability || (applicability == currentApplicability && group < bestGroup)) {
             candidates.clear()
@@ -45,13 +45,12 @@ open class CandidateCollector(
         return applicability
     }
 
-    fun bestCandidates() = candidates
+    fun bestCandidates(): List<Candidate> = candidates
 
-    fun shouldStopAtTheLevel(group: TowerGroup) =
+    fun shouldStopAtTheLevel(group: TowerGroup): Boolean =
         isSuccess() && bestGroup < group
 
     fun isSuccess(): Boolean {
         return currentApplicability.isSuccess
     }
 }
-

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildErrorFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildErrorProperty
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirErrorFunctionSymbol
@@ -21,30 +20,30 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 
 class CandidateFactory private constructor(
-    val bodyResolveComponents: BodyResolveComponents,
+    val context: ResolutionContext,
     val callInfo: CallInfo,
     private val baseSystem: ConstraintStorage
 ) {
 
     companion object {
-        private fun buildBaseSystem(bodyResolveComponents: BodyResolveComponents, callInfo: CallInfo): ConstraintStorage {
+        private fun buildBaseSystem(context: ResolutionContext, callInfo: CallInfo): ConstraintStorage {
-            val system = bodyResolveComponents.inferenceComponents.createConstraintSystem()
+            val system = context.inferenceComponents.createConstraintSystem()
             callInfo.arguments.forEach {
                 system.addSubsystemFromExpression(it)
             }
-            system.addOtherSystem(bodyResolveComponents.inferenceComponents.inferenceSession.currentConstraintSystem)
+            system.addOtherSystem(context.bodyResolveContext.inferenceSession.currentConstraintSystem)
             return system.asReadOnlyStorage()
         }
     }
 
-    constructor(bodyResolveComponents: BodyResolveComponents, callInfo: CallInfo) :
+    constructor(context: ResolutionContext, callInfo: CallInfo) :
-            this(bodyResolveComponents, callInfo, buildBaseSystem(bodyResolveComponents, callInfo))
+            this(context, callInfo, buildBaseSystem(context, callInfo))
 
     fun replaceCallInfo(callInfo: CallInfo): CandidateFactory {
         if (this.callInfo.arguments.size != callInfo.arguments.size) {
             throw AssertionError("Incorrect replacement of call info in CandidateFactory")
         }
-        return CandidateFactory(bodyResolveComponents, callInfo, baseSystem)
+        return CandidateFactory(context, callInfo, baseSystem)
     }
 
     fun createCandidate(
@@ -56,7 +55,7 @@ class CandidateFactory private constructor(
     ): Candidate {
         return Candidate(
             symbol, dispatchReceiverValue, implicitExtensionReceiverValue,
-            explicitReceiverKind, bodyResolveComponents.inferenceComponents.constraintSystemFactory, baseSystem,
+            explicitReceiverKind, context.inferenceComponents.constraintSystemFactory, baseSystem,
             builtInExtensionFunctionReceiverValue?.receiverExpression?.let {
                 callInfo.withReceiverAsArgument(it)
             } ?: callInfo
@@ -78,7 +77,7 @@ class CandidateFactory private constructor(
             dispatchReceiverValue = null,
             implicitExtensionReceiverValue = null,
             explicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
-            bodyResolveComponents.inferenceComponents.constraintSystemFactory,
+            context.inferenceComponents.constraintSystemFactory,
             baseSystem,
             callInfo
         )
@@ -87,7 +86,7 @@ class CandidateFactory private constructor(
     private fun createErrorFunctionSymbol(diagnostic: ConeDiagnostic): FirErrorFunctionSymbol {
         return FirErrorFunctionSymbol().also {
             buildErrorFunction {
-                session = this@CandidateFactory.bodyResolveComponents.session
+                session = context.session
                 resolvePhase = FirResolvePhase.BODY_RESOLVE
                 origin = FirDeclarationOrigin.Synthetic
                 this.diagnostic = diagnostic
@@ -99,7 +98,7 @@ class CandidateFactory private constructor(
     private fun createErrorPropertySymbol(diagnostic: ConeDiagnostic): FirErrorPropertySymbol {
         return FirErrorPropertySymbol(diagnostic).also {
             buildErrorProperty {
-                session = this@CandidateFactory.bodyResolveComponents.session
+                session = context.session
                 resolvePhase = FirResolvePhase.BODY_RESOLVE
                 origin = FirDeclarationOrigin.Synthetic
                 name = FirErrorPropertySymbol.NAME

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

@@ -13,8 +13,8 @@ import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.intrinsics.createCoroutineUnintercepted
 import kotlin.coroutines.resume
 
-class ResolutionStageRunner(private val context: ResolutionContext) {
+class ResolutionStageRunner {
-    fun processCandidate(candidate: Candidate, stopOnFirstError: Boolean = true): CandidateApplicability {
+    fun processCandidate(candidate: Candidate, context: ResolutionContext, stopOnFirstError: Boolean = true): CandidateApplicability {
         val sink = CheckerSinkImpl(stopOnFirstError = stopOnFirstError)
         var finished = false
         sink.continuation = suspend {

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.inference.*
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SyntheticSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -32,7 +33,8 @@ import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
 
 data class ResolutionContext(
     val session: FirSession,
-    val bodyResolveComponents: BodyResolveComponents
+    val bodyResolveComponents: BodyResolveComponents,
+    val bodyResolveContext: BodyResolveContext
 ) {
     val inferenceComponents: InferenceComponents
         get() = bodyResolveComponents.inferenceComponents

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirInvokeResolveTowerExtension.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.builder.FirQualifiedAccessExpressionBuilder
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
-import org.jetbrains.kotlin.fir.resolve.calls.ExpressionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
@@ -24,10 +23,12 @@ import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 internal class FirInvokeResolveTowerExtension(
-    private val components: BodyResolveComponents,
+    private val context: ResolutionContext,
     private val manager: TowerResolveManager,
     private val candidateFactoriesAndCollectors: CandidateFactoriesAndCollectors
 ) {
+    private val components: BodyResolveComponents
+        get() = context.bodyResolveComponents
 
     fun enqueueResolveTasksForQualifier(info: CallInfo, receiver: FirResolvedQualifier) {
         if (info.callKind != CallKind.Function) return
@@ -108,7 +109,7 @@ internal class FirInvokeResolveTowerExtension(
             manager,
             towerDataElementsForName,
             collector,
-            CandidateFactory(components, invokeReceiverInfo),
+            CandidateFactory(context, invokeReceiverInfo),
             onSuccessfulLevel = { towerGroup ->
                 enqueueResolverTasksForInvokeReceiverCandidates(
                     invokeBuiltinExtensionMode, info,
@@ -182,7 +183,7 @@ internal class FirInvokeResolveTowerExtension(
         useImplicitReceiverAsBuiltinInvokeArgument: Boolean,
         receiverGroup: TowerGroup
     ) {
-        val invokeOnGivenReceiverCandidateFactory = CandidateFactory(components, invokeFunctionInfo)
+        val invokeOnGivenReceiverCandidateFactory = CandidateFactory(context, invokeFunctionInfo)
         val task = InvokeFunctionResolveTask(
             components,
             manager,
@@ -385,4 +386,4 @@ private class InvokeFunctionResolveTask(
         group.InvokeResolvePriority(InvokeResolvePriority.COMMON_INVOKE),
         explicitReceiverKind
     )
-}
+}

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

@@ -25,24 +25,25 @@ class FirTowerResolver(
 
     fun runResolver(
         info: CallInfo,
+        context: ResolutionContext,
         collector: CandidateCollector = this.collector,
         manager: TowerResolveManager = this.manager
     ): CandidateCollector {
-        val candidateFactoriesAndCollectors = buildCandidateFactoriesAndCollectors(info, collector)
+        val candidateFactoriesAndCollectors = buildCandidateFactoriesAndCollectors(info, collector, context)
 
-        enqueueResolutionTasks(components, manager, candidateFactoriesAndCollectors, info)
+        enqueueResolutionTasks(context, manager, candidateFactoriesAndCollectors, info)
 
         manager.runTasks()
         return collector
     }
 
     private fun enqueueResolutionTasks(
-        components: BodyResolveComponents,
+        context: ResolutionContext,
         manager: TowerResolveManager,
         candidateFactoriesAndCollectors: CandidateFactoriesAndCollectors,
         info: CallInfo
     ) {
-        val invokeResolveTowerExtension = FirInvokeResolveTowerExtension(components, manager, candidateFactoriesAndCollectors)
+        val invokeResolveTowerExtension = FirInvokeResolveTowerExtension(context, manager, candidateFactoriesAndCollectors)
 
         val mainTask = FirTowerResolveTask(
             components,
@@ -77,7 +78,8 @@ class FirTowerResolver(
 
     fun runResolverForDelegatingConstructor(
         info: CallInfo,
-        constructedType: ConeClassLikeType
+        constructedType: ConeClassLikeType,
+        context: ResolutionContext
     ): CandidateCollector {
         val outerType = components.outerClassManager.outerType(constructedType)
         val scope = constructedType.scope(components.session, components.scopeSession) ?: return collector
@@ -90,7 +92,7 @@ class FirTowerResolver(
             else
                 null
 
-        val candidateFactory = CandidateFactory(components, info)
+        val candidateFactory = CandidateFactory(context, info)
         val resultCollector = collector
 
         scope.processDeclaredConstructors {
@@ -102,7 +104,8 @@ class FirTowerResolver(
                     dispatchReceiver,
                     implicitExtensionReceiverValue = null,
                     builtInExtensionFunctionReceiverValue = null
-                )
+                ),
+                context
             )
         }
 
@@ -111,9 +114,10 @@ class FirTowerResolver(
 
     private fun buildCandidateFactoriesAndCollectors(
         info: CallInfo,
-        collector: CandidateCollector
+        collector: CandidateCollector,
+        context: ResolutionContext
     ): CandidateFactoriesAndCollectors {
-        val candidateFactory = CandidateFactory(components, info)
+        val candidateFactory = CandidateFactory(context, info)
         val stubReceiverCandidateFactory =
             if (info.callKind == CallKind.CallableReference && info.stubReceiver != null)
                 candidateFactory.replaceCallInfo(info.replaceExplicitReceiver(info.stubReceiver))

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

@@ -160,7 +160,7 @@ private class TowerScopeLevelProcessor(
                 val declarationReceiverType = (symbol as? FirCallableSymbol<*>)?.fir?.receiverTypeRef?.coneType
                 if (declarationReceiverType is ConeClassLikeType) {
                     if (!AbstractTypeChecker.isSubtypeOf(
-                            candidateFactory.bodyResolveComponents.inferenceComponents.ctx,
+                            candidateFactory.context.inferenceComponents.ctx,
                             extensionReceiverType,
                             declarationReceiverType.lookupTag.constructClassType(
                                 declarationReceiverType.typeArguments.map { ConeStarProjection }.toTypedArray(),
@@ -181,7 +181,7 @@ private class TowerScopeLevelProcessor(
                 dispatchReceiverValue,
                 implicitExtensionReceiverValue,
                 builtInExtensionFunctionReceiverValue
-            )
+            ), candidateFactory.context
         )
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/AbstractManyCandidatesInferenceSession.kt

@@ -9,16 +9,20 @@ import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.calls.candidate
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 
 abstract class AbstractManyCandidatesInferenceSession(
-    protected val components: BodyResolveComponents
+    protected val resolutionContext: ResolutionContext
 ) : FirInferenceSession() {
     private val errorCalls: MutableList<FirResolvable> = mutableListOf()
     protected val partiallyResolvedCalls: MutableList<Pair<FirResolvable, Candidate>> = mutableListOf()
     private val completedCalls: MutableSet<FirResolvable> = mutableSetOf()
 
+    protected val components: BodyResolveComponents
+        get() = resolutionContext.bodyResolveComponents
+
     override val currentConstraintSystem: ConstraintStorage
         get() = partiallyResolvedCalls.lastOrNull()
             ?.second
@@ -44,4 +48,4 @@ abstract class AbstractManyCandidatesInferenceSession(
 
     protected val FirResolvable.candidate: Candidate
         get() = candidate()!!
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConstraintSystemCompleter.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeFixVariableConstraintPosition
 import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.types.*
@@ -36,6 +37,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<FirStatement>,
         candidateReturnType: ConeKotlinType,
+        context: ResolutionContext,
         collectVariablesFromContext: Boolean = false,
         analyze: (PostponedResolvedAtom) -> Unit
     ) {
@@ -52,7 +54,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
 
             if (
                 completionMode == ConstraintSystemCompletionMode.FULL &&
-                resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(c, variableForFixation, postponedAtoms, analyze)
+                resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(c, variableForFixation, postponedAtoms, context, analyze)
             ) {
                 continue
             }
@@ -81,6 +83,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
         c: ConstraintSystemCompletionContext,
         variableForFixation: VariableFixationFinder.VariableForFixation,
         postponedAtoms: List<PostponedResolvedAtom>,
+        context: ResolutionContext,
         analyze: (PostponedResolvedAtom) -> Unit
     ): Boolean {
         val variable = variableForFixation.variable as ConeTypeVariableTypeConstructor
@@ -110,7 +113,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
                     isSuitable = { isBuiltinFunctionalType(components.session) },
                     typeVariableCreator = { ConeTypeVariableForLambdaReturnType(postponedAtom.atom, "_R") },
                     newAtomCreator = { returnTypeVariable, expectedType ->
-                        postponedAtom.transformToResolvedLambda(csBuilder, components.resolutionContext, expectedType, returnTypeVariable)
+                        postponedAtom.transformToResolvedLambda(csBuilder, context, expectedType, returnTypeVariable)
                     }
                 )
             }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt

@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeStubType
@@ -24,9 +24,9 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImp
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 
 class FirBuilderInferenceSession(
-    components: BodyResolveComponents,
+    resolutionContext: ResolutionContext,
     private val stubsForPostponedVariables: Map<ConeTypeVariable, ConeStubType>,
-) : AbstractManyCandidatesInferenceSession(components) {
+) : AbstractManyCandidatesInferenceSession(resolutionContext) {
     private val commonCalls: MutableList<Pair<FirStatement, Candidate>> = mutableListOf()
 
     override fun <T> shouldRunCompletion(call: T): Boolean where T : FirResolvable, T : FirStatement {
@@ -96,7 +96,7 @@ class FirBuilderInferenceSession(
             context,
             ConstraintSystemCompletionMode.FULL,
             partiallyResolvedCalls.map { it.first as FirStatement },
-            components.session.builtinTypes.unitType.type,
+            components.session.builtinTypes.unitType.type, resolutionContext,
             collectVariablesFromContext = true
         ) {
             error("Shouldn't be called in complete constraint system mode")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
@@ -41,7 +42,7 @@ class FirCallCompleter(
 ) : BodyResolveComponents by components {
     val completer = ConstraintSystemCompleter(components)
     private val inferenceSession
-        get() = inferenceComponents.inferenceSession
+        get() = transformer.context.inferenceSession
 
     data class CompletionResult<T>(val result: T, val callCompleted: Boolean)
 
@@ -65,13 +66,13 @@ class FirCallCompleter(
 
         val completionMode = candidate.computeCompletionMode(inferenceComponents, expectedTypeRef, initialType)
 
-        val analyzer = createPostponedArgumentsAnalyzer()
+        val analyzer = createPostponedArgumentsAnalyzer(transformer.resolutionContext)
         call.transformSingle(InvocationKindTransformer, null)
 
         return when (completionMode) {
             ConstraintSystemCompletionMode.FULL -> {
                 if (inferenceSession.shouldRunCompletion(call)) {
-                    completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(call), initialType) {
+                    completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(call), initialType, transformer.resolutionContext) {
                         analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
                     }
                     val finalSubstitutor =
@@ -94,7 +95,13 @@ class FirCallCompleter(
             }
 
             ConstraintSystemCompletionMode.PARTIAL -> {
-                completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(call), initialType) {
+                completer.complete(
+                    candidate.system.asConstraintSystemCompleterContext(),
+                    completionMode,
+                    listOf(call),
+                    initialType,
+                    transformer.resolutionContext
+                ) {
                     analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
                 }
                 val approximatedCall = call.transformSingle(integerOperatorsTypeUpdater, null)
@@ -119,10 +126,10 @@ class FirCallCompleter(
         )
     }
 
-    fun createPostponedArgumentsAnalyzer(): PostponedArgumentsAnalyzer {
+    fun createPostponedArgumentsAnalyzer(context: ResolutionContext): PostponedArgumentsAnalyzer {
         val lambdaAnalyzer = LambdaAnalyzerImpl()
         return PostponedArgumentsAnalyzer(
-            resolutionContext,
+            context,
             lambdaAnalyzer,
             inferenceComponents,
             transformer.components.callResolver
@@ -179,7 +186,7 @@ class FirCallCompleter(
 
             val builderInferenceSession = runIf(stubsForPostponedVariables.isNotEmpty()) {
                 @Suppress("UNCHECKED_CAST")
-                FirBuilderInferenceSession(components, stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>)
+                FirBuilderInferenceSession(transformer.resolutionContext, stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>)
             }
 
             val localContext = towerDataContextForAnonymousFunctions.get(lambdaArgument.symbol) ?: error(
@@ -187,7 +194,7 @@ class FirCallCompleter(
             )
             transformer.context.withTowerDataContext(localContext) {
                 if (builderInferenceSession != null) {
-                    components.inferenceComponents.withInferenceSession(builderInferenceSession) {
+                    transformer.context.withInferenceSession(builderInferenceSession) {
                         lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
                     }
                 } else {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirDelegatedPropertyInferenceSession.kt

@@ -13,13 +13,16 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.FirNamedReference
-import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
@@ -32,9 +35,9 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 class FirDelegatedPropertyInferenceSession(
     val property: FirProperty,
     initialCall: FirExpression,
-    components: BodyResolveComponents,
+    resolutionContext: ResolutionContext,
     private val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
-) : AbstractManyCandidatesInferenceSession(components) {
+) : AbstractManyCandidatesInferenceSession(resolutionContext) {
     init {
         val initialCandidate = (initialCall as? FirResolvable)
             ?.calleeReference
@@ -65,13 +68,13 @@ class FirDelegatedPropertyInferenceSession(
             addOtherSystem(currentConstraintSystem)
         }
         prepareForCompletion(commonSystem, resolvedCalls)
-        components.inferenceComponents.withInferenceSession(DEFAULT) {
+        resolutionContext.bodyResolveContext.withInferenceSession(DEFAULT) {
             @Suppress("UNCHECKED_CAST")
             components.callCompleter.completer.complete(
                 commonSystem.asConstraintSystemCompleterContext(),
                 ConstraintSystemCompletionMode.FULL,
                 resolvedCalls as List<FirStatement>,
-                unitType
+                unitType, resolutionContext
             ) {
                 postponedArgumentsAnalyzer.analyze(
                     commonSystem.asPostponedArgumentsAnalyzerContext(),

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceComponents.kt

@@ -6,10 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.PrivateForInline
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.types.ConeInferenceContext
+import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintIncorporator
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintInjector
 import org.jetbrains.kotlin.resolve.calls.inference.components.ResultTypeResolver
@@ -17,32 +15,15 @@ import org.jetbrains.kotlin.resolve.calls.inference.components.TrivialConstraint
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 
-class InferenceComponents(
+class InferenceComponents(val session: FirSession) {
-    val ctx: ConeInferenceContext,
+    val ctx: ConeInferenceContext = ConeTypeCheckerContext(isErrorTypeEqualsToAnything = false, isStubTypeEqualsToAnything = false, session)
-    val session: FirSession,
+
-    val returnTypeCalculator: ReturnTypeCalculator,
-    val scopeSession: ScopeSession
-) {
     val approximator: AbstractTypeApproximator = object : AbstractTypeApproximator(ctx) {}
     val trivialConstraintTypeInferenceOracle = TrivialConstraintTypeInferenceOracle.create(ctx)
     private val incorporator = ConstraintIncorporator(approximator, trivialConstraintTypeInferenceOracle, ConeConstraintSystemUtilContext)
     private val injector = ConstraintInjector(incorporator, approximator)
     val resultTypeResolver = ResultTypeResolver(approximator, trivialConstraintTypeInferenceOracle)
 
-    @set:PrivateForInline
-    var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
-
-    @OptIn(PrivateForInline::class)
-    inline fun <R> withInferenceSession(inferenceSession: FirInferenceSession, block: () -> R): R {
-        val oldSession = this.inferenceSession
-        this.inferenceSession = inferenceSession
-        return try {
-            block()
-        } finally {
-            this.inferenceSession = oldSession
-        }
-    }
-
     val constraintSystemFactory = ConstraintSystemFactory()
 
     fun createConstraintSystem(): NewConstraintSystemImpl {

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

@@ -51,7 +51,7 @@ class FirSyntheticCallGenerator(
     private val checkNotNullFunction: FirSimpleFunction = generateSyntheticCheckNotNullFunction()
     private val elvisFunction: FirSimpleFunction = generateSyntheticElvisFunction()
 
-    fun generateCalleeForWhenExpression(whenExpression: FirWhenExpression): FirWhenExpression? {
+    fun generateCalleeForWhenExpression(whenExpression: FirWhenExpression, context: ResolutionContext): FirWhenExpression? {
         val stubReference = whenExpression.calleeReference
         // TODO: Investigate: assertion failed in ModularizedTest
         // assert(stubReference is FirStubReference)
@@ -63,13 +63,14 @@ class FirSyntheticCallGenerator(
         val reference = generateCalleeReferenceWithCandidate(
             whenSelectFunction,
             argumentList,
-            SyntheticCallableId.WHEN.callableName
+            SyntheticCallableId.WHEN.callableName,
+            context = context
         ) ?: return null // TODO
 
         return whenExpression.transformCalleeReference(UpdateReference, reference)
     }
 
-    fun generateCalleeForTryExpression(tryExpression: FirTryExpression): FirTryExpression? {
+    fun generateCalleeForTryExpression(tryExpression: FirTryExpression, context: ResolutionContext): FirTryExpression? {
         val stubReference = tryExpression.calleeReference
         assert(stubReference is FirStubReference)
 
@@ -85,26 +86,28 @@ class FirSyntheticCallGenerator(
         val reference = generateCalleeReferenceWithCandidate(
             trySelectFunction,
             argumentList,
-            SyntheticCallableId.TRY.callableName
+            SyntheticCallableId.TRY.callableName,
+            context = context
         ) ?: return null // TODO
 
         return tryExpression.transformCalleeReference(UpdateReference, reference)
     }
 
-    fun generateCalleeForCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall): FirCheckNotNullCall? {
+    fun generateCalleeForCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall, context: ResolutionContext): FirCheckNotNullCall? {
         val stubReference = checkNotNullCall.calleeReference
         if (stubReference !is FirStubReference) return null
 
         val reference = generateCalleeReferenceWithCandidate(
             checkNotNullFunction,
             checkNotNullCall.argumentList,
-            SyntheticCallableId.CHECK_NOT_NULL.callableName
+            SyntheticCallableId.CHECK_NOT_NULL.callableName,
+            context = context
         ) ?: return null // TODO
 
         return checkNotNullCall.transformCalleeReference(UpdateReference, reference)
     }
 
-    fun generateCalleeForElvisExpression(elvisExpression: FirElvisExpression): FirElvisExpression? {
+    fun generateCalleeForElvisExpression(elvisExpression: FirElvisExpression, context: ResolutionContext): FirElvisExpression? {
         if (elvisExpression.calleeReference !is FirStubReference) return null
 
         val argumentList = buildArgumentList {
@@ -114,7 +117,8 @@ class FirSyntheticCallGenerator(
         val reference = generateCalleeReferenceWithCandidate(
             elvisFunction,
             argumentList,
-            SyntheticCallableId.ELVIS_NOT_NULL.callableName
+            SyntheticCallableId.ELVIS_NOT_NULL.callableName,
+            context = context
         ) ?: return null
 
         return elvisExpression.transformCalleeReference(UpdateReference, reference)
@@ -122,13 +126,18 @@ class FirSyntheticCallGenerator(
 
     fun resolveCallableReferenceWithSyntheticOuterCall(
         callableReferenceAccess: FirCallableReferenceAccess,
-        expectedTypeRef: FirTypeRef?
+        expectedTypeRef: FirTypeRef?,
+        context: ResolutionContext
     ): FirCallableReferenceAccess? {
         val argumentList = buildUnaryArgumentList(callableReferenceAccess)
 
         val reference =
             generateCalleeReferenceWithCandidate(
-                idFunction, argumentList, SyntheticCallableId.ID.callableName, CallKind.SyntheticIdForCallableReferencesResolution
+                idFunction,
+                argumentList,
+                SyntheticCallableId.ID.callableName,
+                CallKind.SyntheticIdForCallableReferencesResolution,
+                context
             ) ?: return callableReferenceAccess.transformCalleeReference(
                 StoreCalleeReference,
                 buildErrorNamedReference {
@@ -149,11 +158,12 @@ class FirSyntheticCallGenerator(
         function: FirSimpleFunction,
         argumentList: FirArgumentList,
         name: Name,
-        callKind: CallKind = CallKind.SyntheticSelect
+        callKind: CallKind = CallKind.SyntheticSelect,
+        context: ResolutionContext
     ): FirNamedReferenceWithCandidate? {
         val callInfo = generateCallInfo(name, argumentList, callKind)
-        val candidate = generateCandidate(callInfo, function)
+        val candidate = generateCandidate(callInfo, function, context)
-        val applicability = resolutionStageRunner.processCandidate(candidate)
+        val applicability = resolutionStageRunner.processCandidate(candidate, context)
         if (applicability <= CandidateApplicability.INAPPLICABLE) {
             return null
         }
@@ -161,8 +171,8 @@ class FirSyntheticCallGenerator(
         return FirNamedReferenceWithCandidate(null, name, candidate)
     }
 
-    private fun generateCandidate(callInfo: CallInfo, function: FirSimpleFunction): Candidate =
+    private fun generateCandidate(callInfo: CallInfo, function: FirSimpleFunction, context: ResolutionContext): Candidate =
-        CandidateFactory(components, callInfo).createCandidate(
+        CandidateFactory(context, callInfo).createCandidate(
             symbol = function.symbol,
             explicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
         )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
+import org.jetbrains.kotlin.fir.resolve.inference.FirInferenceSession
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
@@ -52,6 +53,20 @@ class BodyResolveContext(
 
     val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
 
+    @set:PrivateForInline
+    var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
+
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withInferenceSession(inferenceSession: FirInferenceSession, block: () -> R): R {
+        val oldSession = this.inferenceSession
+        this.inferenceSession = inferenceSession
+        return try {
+            block()
+        } finally {
+            this.inferenceSession = oldSession
+        }
+    }
+
     @OptIn(PrivateForInline::class)
     inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
         val old = towerDataContextsForClassParts

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveUtils.kt

@@ -11,13 +11,13 @@ import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildVarargArgumentsExpression
-import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.arrayElementType
@@ -43,7 +43,7 @@ internal fun inferenceComponents(
     scopeSession: ScopeSession
 ): InferenceComponents {
     val inferenceContext = session.typeContext
-    return InferenceComponents(inferenceContext, session, returnTypeCalculator, scopeSession)
+    return InferenceComponents(session)
 }
 
 internal fun remapArgumentsWithVararg(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt

@@ -5,8 +5,13 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirCallResolver
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.FirQualifiedNameResolver
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.PrivateForInline
+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.*
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
@@ -23,6 +28,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAbstractPhaseTransformer<ResolutionMode>(phase) {
     abstract val context: BodyResolveContext
     abstract val components: BodyResolveTransformerComponents
+    abstract val resolutionContext: ResolutionContext
 
     @set:PrivateForInline
     abstract var implicitTypeOnly: Boolean
@@ -93,8 +99,6 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         val transformer: FirBodyResolveTransformer,
         val context: BodyResolveContext
     ) : BodyResolveComponents {
-        override val resolutionContext: ResolutionContext = ResolutionContext(session, this@BodyResolveTransformerComponents)
-
         override val fileImportsScope: List<FirScope> get() = context.fileImportsScope
         override val towerDataElements: List<FirTowerDataElement> get() = context.towerDataContext.towerDataElements
         override val localScopes: FirLocalScopes get() = context.towerDataContext.localScopes
@@ -112,7 +116,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         override val symbolProvider: FirSymbolProvider = session.firSymbolProvider
 
         override val inferenceComponents: InferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
-        override val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner(resolutionContext)
+        override val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner()
         override val samResolver: FirSamResolver = FirSamResolverImpl(session, scopeSession)
         private val qualifiedResolver: FirQualifiedNameResolver = FirQualifiedNameResolver(this)
         override val callResolver: FirCallResolver = FirCallResolver(

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
@@ -41,6 +42,8 @@ open class FirBodyResolveTransformer(
     final override val components: BodyResolveTransformerComponents =
         BodyResolveTransformerComponents(session, scopeSession, this, context)
 
+    final override val resolutionContext: ResolutionContext = ResolutionContext(session, components, context)
+
     internal open val expressionsTransformer = FirExpressionsResolveTransformer(this)
     protected open val declarationsTransformer = FirDeclarationsResolveTransformer(this)
     private val controlFlowStatementsTransformer = FirControlFlowStatementsResolveTransformer(this)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt

@@ -79,7 +79,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
                 else -> {
                     whenExpression = whenExpression.transformBranches(transformer, ResolutionMode.ContextDependent)
 
-                    whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression) ?: run {
+                    whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression, resolutionContext) ?: run {
                         whenExpression = whenExpression.transformSingle(whenExhaustivenessTransformer, null)
                         dataFlowAnalyzer.exitWhenExpression(whenExpression)
                         whenExpression.resultType = buildErrorTypeRef {
@@ -149,7 +149,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         var callCompleted = false
 
         @Suppress("NAME_SHADOWING")
-        var result = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression)?.let {
+        var result = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression, resolutionContext)?.let {
             val expectedTypeRef = data.expectedType
             val completionResult = callCompleter.completeCall(it, expectedTypeRef)
             callCompleted = completionResult.callCompleted
@@ -218,7 +218,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
         elvisExpression.transformRhs(transformer, ResolutionMode.ContextDependent)
 
-        val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression)?.let {
+        val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext)?.let {
             callCompleter.completeCall(it, data.expectedType).result
         } ?: elvisExpression.also {
             it.resultType = buildErrorTypeRef {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -214,11 +214,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         val inferenceSession = FirDelegatedPropertyInferenceSession(
             property,
             delegateExpression,
-            components,
+            resolutionContext,
-            callCompleter.createPostponedArgumentsAnalyzer()
+            callCompleter.createPostponedArgumentsAnalyzer(resolutionContext)
         )
 
-        components.inferenceComponents.withInferenceSession(inferenceSession) {
+        context.withInferenceSession(inferenceSession) {
             property.transformAccessors()
             val completedCalls = inferenceSession.completeCandidates()
             val finalSubstitutor = inferenceSession.createFinalSubstitutor()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -456,7 +456,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     private inline fun <T> resolveCandidateForAssignmentOperatorCall(block: () -> T): T {
         return dataFlowAnalyzer.withIgnoreFunctionCalls {
             callResolver.withNoArgumentsTransform {
-                inferenceComponents.withInferenceSession(InferenceSessionForAssignmentOperatorCall) {
+                context.withInferenceSession(InferenceSessionForAssignmentOperatorCall) {
                     block()
                 }
             }
@@ -546,7 +546,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         checkNotNullCall.argumentList.transformArguments(transformer, ResolutionMode.ContextDependent)
 
         var callCompleted = false
-        val result = components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall)?.let {
+        val result = components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall, resolutionContext)?.let {
             val completionResult = callCompleter.completeCall(it, data.expectedType)
             callCompleted = completionResult.callCompleted
             completionResult.result
@@ -629,7 +629,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         if (data !is ResolutionMode.ContextDependent) {
             val resolvedReference =
                 components.syntheticCallGenerator.resolveCallableReferenceWithSyntheticOuterCall(
-                    callableReferenceAccess, data.expectedType,
+                    callableReferenceAccess, data.expectedType, resolutionContext,
                 ) ?: callableReferenceAccess
 
             return resolvedReference.compose()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirPartialBodyResolveTransformer.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 
 abstract class FirPartialBodyResolveTransformer(
@@ -21,6 +22,10 @@ abstract class FirPartialBodyResolveTransformer(
     final override inline val components: BodyResolveTransformerComponents
         get() = transformer.components
 
+    @Suppress("OVERRIDE_BY_INLINE")
+    final override inline val resolutionContext: ResolutionContext
+        get() = transformer.resolutionContext
+
     @set:PrivateForInline
     override var implicitTypeOnly: Boolean
         get() = transformer.implicitTypeOnly

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/resolver/SingleCandidateResolver.kt

@@ -47,7 +47,7 @@ class SingleCandidateResolver(
         stubBodyResolveTransformer,
         bodyResolveComponents,
     )
-    private val resolutionStageRunner = ResolutionStageRunner(ResolutionContext(firSession, bodyResolveComponents))
+    private val resolutionStageRunner = ResolutionStageRunner()
 
     fun resolveSingleCandidate(
         resolutionParameters: ResolutionParameters
@@ -62,14 +62,16 @@ class SingleCandidateResolver(
         val dispatchReceiverValue = infoProvider.dispatchReceiverValue()
         val implicitExtensionReceiverValue = infoProvider.implicitExtensionReceiverValue()
 
-        val candidate = CandidateFactory(bodyResolveComponents, callInfo).createCandidate(
+        val resolutionContext = stubBodyResolveTransformer.resolutionContext
+
+        val candidate = CandidateFactory(resolutionContext, callInfo).createCandidate(
             resolutionParameters.callableSymbol,
             explicitReceiverKind = explicitReceiverKind,
             dispatchReceiverValue = dispatchReceiverValue,
             implicitExtensionReceiverValue = implicitExtensionReceiverValue,
         )
 
-        val applicability = resolutionStageRunner.processCandidate(candidate, stopOnFirstError = true)
+        val applicability = resolutionStageRunner.processCandidate(candidate, resolutionContext, stopOnFirstError = true)
         if (applicability.isSuccess) {
             return completeResolvedCandidate(candidate, resolutionParameters)
         }