[FIR] Introduce ResolutionContext and get rid of components in Candidate

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -492,7 +492,7 @@ class FirElementSerializer private constructor(
         }
 
         if (parameter.isVararg) {
-            val varargElementType = parameter.returnTypeRef.coneType.varargElementType(session)
+            val varargElementType = parameter.returnTypeRef.coneType.varargElementType()
             if (useTypeTable()) {
                 builder.varargElementTypeId = typeId(varargElementType)
             } else {

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

@@ -10,10 +10,7 @@ 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.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.*
-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
@@ -59,6 +56,8 @@ 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/Arguments.kt

@@ -36,6 +36,7 @@ fun Candidate.resolveArgumentExpression(
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef?,
     sink: CheckerSink,
+    context: ResolutionContext,
     isReceiver: Boolean,
     isDispatch: Boolean
 ) {
@@ -45,6 +46,7 @@ fun Candidate.resolveArgumentExpression(
             argument as FirResolvable,
             expectedType,
             sink,
+            context,
             isReceiver,
             isDispatch
         )
@@ -62,6 +64,7 @@ fun Candidate.resolveArgumentExpression(
                     nestedQualifier,
                     expectedType,
                     sink,
+                    context,
                     isReceiver,
                     isDispatch,
                     useNullableArgumentType = true
@@ -86,13 +89,14 @@ fun Candidate.resolveArgumentExpression(
                     argument,
                     expectedType,
                     sink,
+                    context,
                     isReceiver,
                     isDispatch
                 )
             else
-                preprocessCallableReference(argument, expectedType)
+                preprocessCallableReference(argument, expectedType, context)
         // TODO:!
-        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef)
+        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef, context)
         // TODO:!
         //TODO: Collection literal
         is FirWrappedArgumentExpression -> resolveArgumentExpression(
@@ -101,6 +105,7 @@ fun Candidate.resolveArgumentExpression(
             expectedType,
             expectedTypeRef,
             sink,
+            context,
             isReceiver,
             isDispatch
         )
@@ -110,10 +115,11 @@ fun Candidate.resolveArgumentExpression(
             expectedType,
             expectedTypeRef,
             sink,
+            context,
             isReceiver,
             isDispatch
         )
-        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isDispatch)
+        else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, context, isReceiver, isDispatch)
     }
 }
 
@@ -123,6 +129,7 @@ private fun Candidate.resolveBlockArgument(
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef?,
     sink: CheckerSink,
+    context: ResolutionContext,
     isReceiver: Boolean,
     isDispatch: Boolean
 ) {
@@ -146,6 +153,7 @@ private fun Candidate.resolveBlockArgument(
             expectedType,
             expectedTypeRef,
             sink,
+            context,
             isReceiver,
             isDispatch
         )
@@ -157,6 +165,7 @@ fun Candidate.resolveSubCallArgument(
     argument: FirResolvable,
     expectedType: ConeKotlinType?,
     sink: CheckerSink,
+    context: ResolutionContext,
     isReceiver: Boolean,
     isDispatch: Boolean,
     useNullableArgumentType: Boolean = false
@@ -166,6 +175,7 @@ fun Candidate.resolveSubCallArgument(
         argument as FirExpression,
         expectedType,
         sink,
+        context,
         isReceiver,
         isDispatch,
         useNullableArgumentType
@@ -180,7 +190,7 @@ fun Candidate.resolveSubCallArgument(
         sink.components.returnTypeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe()).type
     }
     val argumentType = candidate.substitutor.substituteOrSelf(type)
-    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, useNullableArgumentType)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, context, isReceiver, isDispatch, useNullableArgumentType)
 }
 
 fun Candidate.resolvePlainExpressionArgument(
@@ -188,13 +198,14 @@ fun Candidate.resolvePlainExpressionArgument(
     argument: FirExpression,
     expectedType: ConeKotlinType?,
     sink: CheckerSink,
+    context: ResolutionContext,
     isReceiver: Boolean,
     isDispatch: Boolean,
     useNullableArgumentType: Boolean = false
 ) {
     if (expectedType == null) return
     val argumentType = argument.typeRef.coneTypeSafe<ConeKotlinType>() ?: return
-    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, useNullableArgumentType)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, context, isReceiver, isDispatch, useNullableArgumentType)
     checkApplicabilityForIntegerOperatorCall(sink, argument)
 }
 
@@ -210,6 +221,7 @@ fun Candidate.resolvePlainArgumentType(
     argumentType: ConeKotlinType,
     expectedType: ConeKotlinType?,
     sink: CheckerSink,
+    context: ResolutionContext,
     isReceiver: Boolean,
     isDispatch: Boolean,
     useNullableArgumentType: Boolean = false
@@ -217,7 +229,7 @@ fun Candidate.resolvePlainArgumentType(
     val position = SimpleConstraintSystemConstraintPosition //TODO
 
     val session = sink.components.session
-    val capturedType = prepareCapturedType(argumentType)
+    val capturedType = prepareCapturedType(argumentType, context)
 
     val argumentTypeForApplicabilityCheck =
         if (useNullableArgumentType)
@@ -230,24 +242,24 @@ fun Candidate.resolvePlainArgumentType(
     )
 }
 
-fun Candidate.prepareCapturedType(argumentType: ConeKotlinType): ConeKotlinType {
+fun Candidate.prepareCapturedType(argumentType: ConeKotlinType, context: ResolutionContext): ConeKotlinType {
-    return captureTypeFromExpressionOrNull(argumentType) ?: argumentType
+    return captureTypeFromExpressionOrNull(argumentType, context) ?: argumentType
 }
 
-private fun Candidate.captureTypeFromExpressionOrNull(argumentType: ConeKotlinType): ConeKotlinType? {
+private fun Candidate.captureTypeFromExpressionOrNull(argumentType: ConeKotlinType, context: ResolutionContext): ConeKotlinType? {
     if (argumentType is ConeFlexibleType) {
-        return captureTypeFromExpressionOrNull(argumentType.lowerBound)
+        return captureTypeFromExpressionOrNull(argumentType.lowerBound, context)
     }
 
     if (argumentType !is ConeClassLikeType) return null
 
-    argumentType.fullyExpandedType(bodyResolveComponents.session).let {
+    argumentType.fullyExpandedType(context.session).let {
-        if (it !== argumentType) return captureTypeFromExpressionOrNull(it)
+        if (it !== argumentType) return captureTypeFromExpressionOrNull(it, context)
     }
 
     if (argumentType.typeArguments.isEmpty()) return null
 
-    return bodyResolveComponents.inferenceComponents.ctx.captureFromArguments(
+    return context.inferenceComponents.ctx.captureFromArguments(
         argumentType, CaptureStatus.FROM_EXPRESSION
     ) as? ConeKotlinType
 }
@@ -289,47 +301,55 @@ internal fun Candidate.resolveArgument(
     argument: FirExpression,
     parameter: FirValueParameter?,
     isReceiver: Boolean,
-    sink: CheckerSink
+    sink: CheckerSink,
+    context: ResolutionContext
 ) {
 
     argument.resultType.ensureResolvedTypeDeclaration(sink.components.session)
 
-    val expectedType = prepareExpectedType(sink.components.session, argument, parameter)
+    val expectedType = prepareExpectedType(sink.components.session, argument, parameter, context)
     resolveArgumentExpression(
         this.system.getBuilder(),
         argument,
         expectedType,
         parameter?.returnTypeRef,
         sink,
+        context,
         isReceiver,
         false
     )
 }
 
-private fun Candidate.prepareExpectedType(session: FirSession, argument: FirExpression, parameter: FirValueParameter?): ConeKotlinType? {
+private fun Candidate.prepareExpectedType(
+    session: FirSession,
+    argument: FirExpression,
+    parameter: FirValueParameter?,
+    context: ResolutionContext
+): ConeKotlinType? {
     if (parameter == null) return null
     if (parameter.returnTypeRef is FirILTTypeRefPlaceHolder && argument.resultType is FirResolvedTypeRef)
         return argument.resultType.coneType
     val basicExpectedType = argument.getExpectedType(session, parameter/*, LanguageVersionSettings*/)
-    val expectedType = getExpectedTypeWithSAMConversion(session, argument, basicExpectedType) ?: basicExpectedType
+    val expectedType = getExpectedTypeWithSAMConversion(session, argument, basicExpectedType, context) ?: basicExpectedType
     return this.substitutor.substituteOrSelf(expectedType)
 }
 
 private fun Candidate.getExpectedTypeWithSAMConversion(
     session: FirSession,
     argument: FirExpression,
-    candidateExpectedType: ConeKotlinType
+    candidateExpectedType: ConeKotlinType,
+    context: ResolutionContext
 ): ConeKotlinType? {
     if (candidateExpectedType.isBuiltinFunctionalType(session)) return null
     // TODO: if (!callComponents.languageVersionSettings.supportsFeature(LanguageFeature.SamConversionPerArgument)) return null
     val firFunction = symbol.fir as? FirFunction<*> ?: return null
-    if (!samResolver.shouldRunSamConversionForFunction(firFunction)) return null
+    if (!context.bodyResolveComponents.samResolver.shouldRunSamConversionForFunction(firFunction)) return null
 
     if (!argument.isFunctional(session)) return null
 
     // TODO: resolvedCall.registerArgumentWithSamConversion(argument, SamConversionDescription(convertedTypeByOriginal, convertedTypeByCandidate!!))
 
-    return samResolver.getFunctionTypeForPossibleSamType(candidateExpectedType).apply {
+    return context.bodyResolveComponents.samResolver.getFunctionTypeForPossibleSamType(candidateExpectedType).apply {
         usesSAM = true
     }
 }
@@ -351,12 +371,12 @@ internal fun FirExpression.getExpectedType(
     }
 
     return if (parameter.isVararg && shouldUnwrapVarargType) {
-        parameter.returnTypeRef.coneType.varargElementType(session)
+        parameter.returnTypeRef.coneType.varargElementType()
     } else {
         parameter.returnTypeRef.coneType
     }
 }
 
-fun ConeKotlinType.varargElementType(session: FirSession): ConeKotlinType {
+fun ConeKotlinType.varargElementType(): ConeKotlinType {
     return this.arrayElementType() ?: error("Failed to extract! ${this.render()}!")
 }

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionStub
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.DoubleColonLHS
+import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.inference.PostponedResolvedAtom
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -78,21 +79,19 @@ class Candidate(
     val dispatchReceiverValue: ReceiverValue?,
     val implicitExtensionReceiverValue: ImplicitReceiverValue<*>?,
     val explicitReceiverKind: ExplicitReceiverKind,
-    val bodyResolveComponents: BodyResolveComponents,
+    val constraintSystemFactory: InferenceComponents.ConstraintSystemFactory,
     private val baseSystem: ConstraintStorage,
     val callInfo: CallInfo
 ) {
 
     var systemInitialized: Boolean = false
     val system: NewConstraintSystemImpl by lazy {
-        val system = bodyResolveComponents.inferenceComponents.createConstraintSystem()
+        val system = constraintSystemFactory.createConstraintSystem()
         system.addOtherSystem(baseSystem)
         systemInitialized = true
         system
     }
 
-    val samResolver get() = bodyResolveComponents.samResolver
-
     lateinit var substitutor: ConeSubstitutor
     lateinit var freshVariables: List<ConeTypeVariable>
     var resultingTypeForCallableReference: ConeKotlinType? = null

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

@@ -56,7 +56,7 @@ class CandidateFactory private constructor(
     ): Candidate {
         return Candidate(
             symbol, dispatchReceiverValue, implicitExtensionReceiverValue,
-            explicitReceiverKind, bodyResolveComponents, baseSystem,
+            explicitReceiverKind, bodyResolveComponents.inferenceComponents.constraintSystemFactory, baseSystem,
             builtInExtensionFunctionReceiverValue?.receiverExpression?.let {
                 callInfo.withReceiverAsArgument(it)
             } ?: callInfo
@@ -78,7 +78,7 @@ class CandidateFactory private constructor(
             dispatchReceiverValue = null,
             implicitExtensionReceiverValue = null,
             explicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
-            bodyResolveComponents,
+            bodyResolveComponents.inferenceComponents.constraintSystemFactory,
             baseSystem,
             callInfo
         )

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

@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 
 internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val declaration = candidate.symbol.fir
         if (declaration !is FirTypeParameterRefsOwner || declaration.typeParameters.isEmpty()) {
             candidate.substitutor = ConeSubstitutor.Empty
@@ -66,7 +66,7 @@ internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
                     getTypePreservingFlexibilityWrtTypeVariable(
                         typeArgument.typeRef.coneType,
                         typeParameter,
-                        candidate.bodyResolveComponents.inferenceComponents.ctx
+                        context.bodyResolveComponents.inferenceComponents.ctx
                     ),
                     SimpleConstraintSystemConstraintPosition // TODO
                 )

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls
 
-import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 import kotlin.coroutines.Continuation
@@ -14,15 +13,15 @@ import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.intrinsics.createCoroutineUnintercepted
 import kotlin.coroutines.resume
 
-class ResolutionStageRunner(val components: InferenceComponents) {
+class ResolutionStageRunner(private val context: ResolutionContext) {
     fun processCandidate(candidate: Candidate, stopOnFirstError: Boolean = true): CandidateApplicability {
-        val sink = CheckerSinkImpl(components, stopOnFirstError = stopOnFirstError)
+        val sink = CheckerSinkImpl(context.bodyResolveComponents.inferenceComponents, stopOnFirstError = stopOnFirstError)
         var finished = false
         sink.continuation = suspend {
             candidate.callInfo.callKind.resolutionSequence.forEachIndexed { index, stage ->
                 if (index < candidate.passedStages) return@forEachIndexed
                 candidate.passedStages++
-                stage.check(candidate, sink, candidate.callInfo)
+                stage.check(candidate, candidate.callInfo, sink, context)
             }
         }.createCoroutineUnintercepted(completion = object : Continuation<Unit> {
             override val context: CoroutineContext

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 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.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SyntheticSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -29,14 +30,25 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
 
+data class ResolutionContext(
+    val session: FirSession,
+    val bodyResolveComponents: BodyResolveComponents
+) {
+    val inferenceComponents: InferenceComponents
+        get() = bodyResolveComponents.inferenceComponents
+
+    val returnTypeCalculator: ReturnTypeCalculator
+        get() = bodyResolveComponents.returnTypeCalculator
+}
+
 abstract class ResolutionStage {
-    abstract suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo)
+    abstract suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext)
 }
 
 abstract class CheckerStage : ResolutionStage()
 
 internal object CheckExplicitReceiverConsistency : ResolutionStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val receiverKind = candidate.explicitReceiverKind
         val explicitReceiver = callInfo.explicitReceiver
         // TODO: add invoke cases
@@ -71,7 +83,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
             return this == DISPATCH_RECEIVER || this == BOTH_RECEIVERS
         }
 
-        override fun Candidate.getReceiverType(): ConeKotlinType? {
+        override fun Candidate.getReceiverType(context: ResolutionContext): ConeKotlinType? {
             return dispatchReceiverValue?.type
         }
     }
@@ -85,25 +97,25 @@ internal sealed class CheckReceivers : ResolutionStage() {
             return this == EXTENSION_RECEIVER || this == BOTH_RECEIVERS
         }
 
-        override fun Candidate.getReceiverType(): ConeKotlinType? {
+        override fun Candidate.getReceiverType(context: ResolutionContext): ConeKotlinType? {
             val callableSymbol = symbol as? FirCallableSymbol<*> ?: return null
             val callable = callableSymbol.fir
             val receiverType = callable.receiverTypeRef?.coneType
             if (receiverType != null) return receiverType
             val returnTypeRef = callable.returnTypeRef as? FirResolvedTypeRef ?: return null
-            if (!returnTypeRef.type.isExtensionFunctionType(bodyResolveComponents.session)) return null
+            if (!returnTypeRef.type.isExtensionFunctionType(context.session)) return null
             return (returnTypeRef.type.typeArguments.firstOrNull() as? ConeKotlinTypeProjection)?.type
         }
     }
 
-    abstract fun Candidate.getReceiverType(): ConeKotlinType?
+    abstract fun Candidate.getReceiverType(context: ResolutionContext): ConeKotlinType?
 
     abstract fun ExplicitReceiverKind.shouldBeCheckedAgainstExplicit(): Boolean
 
     abstract fun ExplicitReceiverKind.shouldBeCheckedAgainstImplicit(): Boolean
 
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
-        val expectedReceiverType = candidate.getReceiverType()
+        val expectedReceiverType = candidate.getReceiverType(context)
         val explicitReceiverExpression = callInfo.explicitReceiver
         val explicitReceiverKind = candidate.explicitReceiverKind
 
@@ -118,6 +130,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
                     expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType),
                     expectedTypeRef = explicitReceiverExpression.typeRef,
                     sink = sink,
+                    context = context,
                     isReceiver = true,
                     isDispatch = this is Dispatch
                 )
@@ -130,6 +143,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
                         argumentType = argumentExtensionReceiverValue.type,
                         expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType.type),
                         sink = sink,
+                        context = context,
                         isReceiver = true,
                         isDispatch = this is Dispatch
                     )
@@ -148,7 +162,7 @@ private fun FirExpression.isSuperReferenceExpression(): Boolean {
 }
 
 internal object MapArguments : ResolutionStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val symbol = candidate.symbol as? FirFunctionSymbol<*> ?: return sink.reportDiagnostic(HiddenCandidate)
         val function = symbol.fir
 
@@ -161,7 +175,7 @@ internal object MapArguments : ResolutionStage() {
 }
 
 internal object CheckArguments : CheckerStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val argumentMapping =
             candidate.argumentMapping ?: error("Argument should be already mapped while checking arguments!")
         for (argument in callInfo.arguments) {
@@ -170,7 +184,8 @@ internal object CheckArguments : CheckerStage() {
                 argument,
                 parameter,
                 isReceiver = false,
-                sink = sink
+                sink = sink,
+                context = context
             )
             if (candidate.system.hasContradiction) {
                 sink.yieldDiagnostic(InapplicableCandidate)
@@ -181,11 +196,11 @@ internal object CheckArguments : CheckerStage() {
 }
 
 internal object EagerResolveOfCallableReferences : CheckerStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         if (candidate.postponedAtoms.isEmpty()) return
         for (atom in candidate.postponedAtoms) {
             if (atom is ResolvedCallableReferenceAtom) {
-                if (!candidate.bodyResolveComponents.callResolver.resolveCallableReference(candidate.csBuilder, atom)) {
+                if (!context.bodyResolveComponents.callResolver.resolveCallableReference(candidate.csBuilder, atom)) {
                     sink.yieldDiagnostic(InapplicableCandidate)
                 }
             }
@@ -194,10 +209,10 @@ internal object EagerResolveOfCallableReferences : CheckerStage() {
 }
 
 internal object CheckCallableReferenceExpectedType : CheckerStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val outerCsBuilder = callInfo.outerCSBuilder ?: return
         val expectedType = callInfo.expectedType
-        val candidateSymbol = candidate.symbol as? FirCallableSymbol<*> ?: return
+        if (candidate.symbol !is FirCallableSymbol<*>) return
 
         val resultingReceiverType = when (callInfo.lhs) {
             is DoubleColonLHS.Type -> callInfo.lhs.type.takeIf { callInfo.explicitReceiver !is FirResolvedQualifier }
@@ -206,7 +221,7 @@ internal object CheckCallableReferenceExpectedType : CheckerStage() {
 
         val fir: FirCallableDeclaration<*> = candidate.symbol.fir
 
-        val returnTypeRef = candidate.bodyResolveComponents.returnTypeCalculator.tryCalculateReturnType(fir)
+        val returnTypeRef = context.bodyResolveComponents.returnTypeCalculator.tryCalculateReturnType(fir)
         // If the expected type is a suspend function type and the current argument of interest is a function reference, we need to do
         // "suspend conversion." Here, during resolution, we bypass constraint system by making resulting type be KSuspendFunction.
         // Then, during conversion, we need to create an adapter function and replace the function reference created here with an adapted
@@ -215,9 +230,8 @@ internal object CheckCallableReferenceExpectedType : CheckerStage() {
         val requireSuspendConversion = expectedType?.isSuspendFunctionType(callInfo.session) == true
         val resultingType: ConeKotlinType = when (fir) {
             is FirFunction -> callInfo.session.createAdaptedKFunctionType(
-                callInfo.session,
+                fir,
-                fir, resultingReceiverType, returnTypeRef,
+                resultingReceiverType, returnTypeRef, expectedParameterTypes = expectedType?.typeArguments,
-                expectedParameterTypes = expectedType?.typeArguments,
                 isSuspend = (fir as? FirSimpleFunction)?.isSuspend == true || requireSuspendConversion,
                 expectedReturnType = extractInputOutputTypesFromCallableReferenceExpectedType(expectedType, callInfo.session)?.outputType
             )
@@ -272,7 +286,6 @@ private fun createKPropertyType(
 }
 
 private fun FirSession.createAdaptedKFunctionType(
-    session: FirSession,
     function: FirFunction<*>,
     receiverType: ConeKotlinType?,
     returnTypeRef: FirResolvedTypeRef,
@@ -326,7 +339,7 @@ private fun FirSession.createAdaptedKFunctionType(
         lastVarargParameter != null
     ) {
         val varargArrayType = lastVarargParameter.returnTypeRef.coneType
-        val varargElementType = varargArrayType.varargElementType(session)
+        val varargElementType = varargArrayType.varargElementType()
         val expectedParameterType = (expectedParameterTypes[parameterTypes.size + shift] as? ConeKotlinTypeProjection)?.type
         // Expect an array or potentially array (i.e., type variable). Pass vararg parameter as-is.
         if (expectedParameterType.isPotentiallyArray()) {
@@ -355,7 +368,7 @@ private fun FirSession.createAdaptedKFunctionType(
 }
 
 internal object DiscriminateSynthetics : CheckerStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         if (candidate.symbol is SyntheticSymbol) {
             sink.reportDiagnostic(ResolvedWithLowPriority)
         }
@@ -363,7 +376,7 @@ internal object DiscriminateSynthetics : CheckerStage() {
 }
 
 internal object CheckVisibility : CheckerStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val visibilityChecker = callInfo.session.visibilityChecker
         val symbol = candidate.symbol
         val declaration = symbol.fir
@@ -406,7 +419,7 @@ internal object CheckLowPriorityInOverloadResolution : CheckerStage() {
     private val LOW_PRIORITY_IN_OVERLOAD_RESOLUTION_CLASS_ID: ClassId =
         ClassId(FqName("kotlin.internal"), Name.identifier("LowPriorityInOverloadResolution"))
 
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val annotations = when (val fir = candidate.symbol.fir) {
             is FirSimpleFunction -> fir.annotations
             is FirProperty -> fir.annotations
@@ -427,7 +440,7 @@ internal object CheckLowPriorityInOverloadResolution : CheckerStage() {
 internal object PostponedVariablesInitializerResolutionStage : ResolutionStage() {
     val BUILDER_INFERENCE_CLASS_ID: ClassId = ClassId.fromString("kotlin/BuilderInference")
 
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val argumentMapping = candidate.argumentMapping ?: return
         // TODO: convert type argument mapping to map [FirTypeParameterSymbol, FirTypedProjection?]
         if (candidate.typeArgumentMapping is TypeArgumentMapping.Mapped) return

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

@@ -24,7 +24,7 @@ sealed class TypeArgumentMapping {
 }
 
 internal object MapTypeArguments : ResolutionStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val typeArguments = callInfo.typeArguments
         if (typeArguments.isEmpty()) {
             candidate.typeArgumentMapping = TypeArgumentMapping.NoExplicitArguments
@@ -43,7 +43,7 @@ internal object MapTypeArguments : ResolutionStage() {
 }
 
 internal object NoTypeArguments : ResolutionStage() {
-    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         if (callInfo.typeArguments.isNotEmpty()) {
             sink.yieldDiagnostic(InapplicableCandidate)
         }

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

@@ -81,7 +81,6 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
         c: ConstraintSystemCompletionContext,
         variableForFixation: VariableFixationFinder.VariableForFixation,
         postponedAtoms: List<PostponedResolvedAtom>,
-        /*diagnosticsHolder: KotlinDiagnosticsHolder,*/
         analyze: (PostponedResolvedAtom) -> Unit
     ): Boolean {
         val variable = variableForFixation.variable as ConeTypeVariableTypeConstructor
@@ -111,7 +110,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
                     isSuitable = { isBuiltinFunctionalType(components.session) },
                     typeVariableCreator = { ConeTypeVariableForLambdaReturnType(postponedAtom.atom, "_R") },
                     newAtomCreator = { returnTypeVariable, expectedType ->
-                        postponedAtom.transformToResolvedLambda(csBuilder, expectedType, returnTypeVariable)
+                        postponedAtom.transformToResolvedLambda(csBuilder, components.resolutionContext, expectedType, returnTypeVariable)
                     }
                 )
             }

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

@@ -122,7 +122,9 @@ class FirCallCompleter(
     fun createPostponedArgumentsAnalyzer(): PostponedArgumentsAnalyzer {
         val lambdaAnalyzer = LambdaAnalyzerImpl()
         return PostponedArgumentsAnalyzer(
-            lambdaAnalyzer, inferenceComponents,
+            resolutionContext,
+            lambdaAnalyzer,
+            inferenceComponents,
             transformer.components.callResolver
         )
     }

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

@@ -43,7 +43,15 @@ class InferenceComponents(
         }
     }
 
+    val constraintSystemFactory = ConstraintSystemFactory()
+
     fun createConstraintSystem(): NewConstraintSystemImpl {
         return NewConstraintSystemImpl(injector, ctx)
     }
+
+    inner class ConstraintSystemFactory {
+        fun createConstraintSystem(): NewConstraintSystemImpl {
+            return this@InferenceComponents.createConstraintSystem()
+        }
+    }
 }

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -20,6 +21,7 @@ fun Candidate.preprocessLambdaArgument(
     argument: FirAnonymousFunction,
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef?,
+    context: ResolutionContext,
     forceResolution: Boolean = false,
     returnTypeVariable: ConeTypeVariableForLambdaReturnType? = null
 ): PostponedResolvedAtom {
@@ -28,8 +30,8 @@ fun Candidate.preprocessLambdaArgument(
     }
 
     val resolvedArgument =
-        extractLambdaInfoFromFunctionalType(expectedType, expectedTypeRef, argument, returnTypeVariable, bodyResolveComponents, this)
+        extractLambdaInfoFromFunctionalType(expectedType, expectedTypeRef, argument, returnTypeVariable, context.bodyResolveComponents, this)
-            ?: extraLambdaInfo(expectedType, argument, csBuilder, bodyResolveComponents.session, this)
+            ?: extraLambdaInfo(expectedType, argument, csBuilder, context.session, this)
 
     if (expectedType != null) {
         // TODO: add SAM conversion processing
@@ -47,10 +49,11 @@ fun Candidate.preprocessLambdaArgument(
 
 fun Candidate.preprocessCallableReference(
     argument: FirCallableReferenceAccess,
-    expectedType: ConeKotlinType?
+    expectedType: ConeKotlinType?,
+    context: ResolutionContext
 ) {
-    val lhs = bodyResolveComponents.doubleColonExpressionResolver.resolveDoubleColonLHS(argument)
+    val lhs = context.bodyResolveComponents.doubleColonExpressionResolver.resolveDoubleColonLHS(argument)
-    postponedAtoms += ResolvedCallableReferenceAtom(argument, expectedType, lhs, bodyResolveComponents.session)
+    postponedAtoms += ResolvedCallableReferenceAtom(argument, expectedType, lhs, context.session)
 }
 
 private fun extraLambdaInfo(

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

@@ -44,6 +44,7 @@ interface LambdaAnalyzer {
 }
 
 class PostponedArgumentsAnalyzer(
+    private val resolutionContext: ResolutionContext,
     private val lambdaAnalyzer: LambdaAnalyzer,
     private val components: InferenceComponents,
     private val callResolver: FirCallResolver
@@ -53,14 +54,13 @@ class PostponedArgumentsAnalyzer(
         c: PostponedArgumentsAnalyzerContext,
         argument: PostponedResolvedAtom,
         candidate: Candidate
-        //diagnosticsHolder: KotlinDiagnosticsHolder
     ) {
         return when (argument) {
             is ResolvedLambdaAtom ->
-                analyzeLambda(c, argument, candidate/*, diagnosticsHolder*/)
+                analyzeLambda(c, argument, candidate)
 
             is LambdaWithTypeVariableAsExpectedTypeAtom ->
-                analyzeLambda(c, argument.transformToResolvedLambda(c.getBuilder()), candidate/*, diagnosticsHolder*/)
+                analyzeLambda(c, argument.transformToResolvedLambda(c.getBuilder(), resolutionContext), candidate)
 
             is ResolvedCallableReferenceAtom -> processCallableReference(argument, candidate)
 
@@ -163,6 +163,7 @@ class PostponedArgumentsAnalyzer(
                 lambda.returnType.let(::substitute),
                 lambda.atom.returnTypeRef, // TODO: proper ref
                 checkerSink,
+                context = resolutionContext,
                 isReceiver = false,
                 isDispatch = false
             )
@@ -181,7 +182,7 @@ class PostponedArgumentsAnalyzer(
 
 fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(
     csBuilder: ConstraintSystemBuilder,
-    /*diagnosticHolder: KotlinDiagnosticsHolder,*/
+    context: ResolutionContext,
     expectedType: ConeKotlinType? = null,
     returnTypeVariable: ConeTypeVariableForLambdaReturnType? = null
 ): ResolvedLambdaAtom {
@@ -192,8 +193,9 @@ fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(
         atom,
         fixedExpectedType,
         expectedTypeRef,
+        context,
         forceResolution = true,
-        returnTypeVariable
+        returnTypeVariable = returnTypeVariable
     ) as ResolvedLambdaAtom
     analyzed = true
     return resolvedAtom

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

@@ -352,7 +352,7 @@ class FirCallCompletionResultsWriterTransformer(
     private fun Candidate.createArgumentsMapping(): ExpectedArgumentType? {
         return argumentMapping?.map { (argument, valueParameter) ->
             val expectedType = if (valueParameter.isVararg) {
-                valueParameter.returnTypeRef.substitute(this).varargElementType(session)
+                valueParameter.returnTypeRef.substitute(this).varargElementType()
             } else {
                 valueParameter.returnTypeRef.substitute(this)
             }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
@@ -92,6 +93,8 @@ 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
@@ -109,7 +112,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(inferenceComponents)
+        override val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner(resolutionContext)
         override val samResolver: FirSamResolver = FirSamResolverImpl(session, scopeSession)
         private val qualifiedResolver: FirQualifiedNameResolver = FirQualifiedNameResolver(this)
         override val callResolver: FirCallResolver = FirCallResolver(

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(bodyResolveComponents.inferenceComponents)
+    private val resolutionStageRunner = ResolutionStageRunner(ResolutionContext(firSession, bodyResolveComponents))
 
     fun resolveSingleCandidate(
         resolutionParameters: ResolutionParameters