[FIR] Implement overload resolution by lambda return type

#KT-43129 Fixed

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir
 
+import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
@@ -20,11 +21,11 @@ import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.calls.tower.FirTowerResolver
 import org.jetbrains.kotlin.fir.resolve.calls.tower.TowerResolveManager
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
-import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
+import org.jetbrains.kotlin.fir.resolve.inference.*
-import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreNameReference
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirExpressionsResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
@@ -35,14 +36,17 @@ import org.jetbrains.kotlin.fir.types.builder.buildStarProjection
 import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
+import org.jetbrains.kotlin.resolve.descriptorUtil.OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_CLASS_ID
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.same
 
 class FirCallResolver(
-    private val components: BodyResolveComponents,
+    private val components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     private val qualifiedResolver: FirQualifiedNameResolver,
 ) {
     private val session = components.session
@@ -150,7 +154,7 @@ class FirCallResolver(
         towerResolver.reset()
         val result = towerResolver.runResolver(info, transformer.resolutionContext)
         val bestCandidates = result.bestCandidates()
-        val reducedCandidates = if (!result.currentApplicability.isSuccess) {
+        var reducedCandidates = if (!result.currentApplicability.isSuccess) {
             bestCandidates.toSet()
         } else {
             val onSuperReference = (explicitReceiver as? FirQualifiedAccessExpression)?.calleeReference is FirSuperReference
@@ -159,9 +163,127 @@ class FirCallResolver(
             )
         }
 
+        if (
+            reducedCandidates.size > 1 &&
+            session.languageVersionSettings.supportsFeature(LanguageFeature.OverloadResolutionByLambdaReturnType) &&
+            bestCandidates.all { components.context.inferenceSession.shouldRunCompletion(qualifiedAccess) }
+        ) {
+            val newCandidates = chooseCandidateRegardingOverloadResolutionByLambdaReturnType(
+                qualifiedAccess,
+                reducedCandidates,
+                bestCandidates
+            )
+            if (newCandidates != null) {
+                reducedCandidates = newCandidates
+            }
+        }
+
         return ResolutionResult(info, result.currentApplicability, reducedCandidates)
     }
 
+    private fun <T> chooseCandidateRegardingOverloadResolutionByLambdaReturnType(
+        call: T,
+        reducedCandidates: Set<Candidate>,
+        allCandidates: Collection<Candidate>
+    ): Set<Candidate>? where T : FirResolvable, T : FirStatement {
+        val candidatesWithAnnotation = allCandidates.filter { candidate ->
+            (candidate.symbol.fir as FirAnnotationContainer).annotations.any {
+                it.annotationTypeRef.coneType.classId == OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_CLASS_ID
+            }
+        }
+        if (candidatesWithAnnotation.isEmpty()) return null
+        val candidatesWithoutAnnotation = reducedCandidates - candidatesWithAnnotation
+        val newCandidates = analyzeLambdaAndReduceNumberOfCandidatesRegardingOverloadResolutionByLambdaReturnType(call, reducedCandidates) ?: return null
+        var maximallySpecificCandidates = conflictResolver.chooseMaximallySpecificCandidates(newCandidates, discriminateGenerics = true, discriminateAbstracts = false)
+        if (maximallySpecificCandidates.size > 1 && candidatesWithoutAnnotation.any { it in maximallySpecificCandidates }) {
+            maximallySpecificCandidates = maximallySpecificCandidates.toMutableSet().apply { removeAll(candidatesWithAnnotation) }
+            maximallySpecificCandidates.singleOrNull()?.addDiagnostic(CandidateChosenUsingOverloadResolutionByLambdaAnnotation)
+        }
+        return maximallySpecificCandidates
+    }
+
+    private fun <T> analyzeLambdaAndReduceNumberOfCandidatesRegardingOverloadResolutionByLambdaReturnType(
+        call: T,
+        candidates: Set<Candidate>,
+    ): Set<Candidate>? where T : FirResolvable, T : FirStatement {
+        val lambdas = candidates.flatMap { candidate ->
+            candidate.postponedAtoms
+                .filter { it is ResolvedLambdaAtom && !it.analyzed }
+                .map { candidate to it as ResolvedLambdaAtom }
+        }.groupBy { (_, atom) -> atom.atom }
+            .values.singleOrNull()?.toMap() ?: return null
+
+        if (!lambdas.values.same { it.parameters.size }) return null
+        if (!lambdas.values.all { it.expectedType?.isBuiltinFunctionalType(session) == true }) return null
+
+        val originalCalleeReference = call.calleeReference
+
+        for (candidate in lambdas.keys) {
+            call.replaceCalleeReference(FirNamedReferenceWithCandidate(null, candidate.callInfo.name, candidate))
+            components.callCompleter.runCompletionForCall(
+                candidate,
+                ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA,
+                call,
+                components.initialTypeOfCandidate(candidate)
+            )
+        }
+
+        try {
+            val inputTypesAreSame = lambdas.entries.same { (candidate, lambda) ->
+                val substitutor = candidate.system.buildCurrentSubstitutor() as ConeSubstitutor
+                lambda.inputTypes.map { substitutor.substituteOrSelf(it) }
+            }
+            if (!inputTypesAreSame) return null
+            lambdas.entries.forEach { (candidate, atom) ->
+                components.callCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate)
+            }
+            val iterator = lambdas.entries.iterator()
+            val (firstCandidate, firstAtom) = iterator.next()
+
+            val postponedArgumentsAnalyzer = components.callCompleter.createPostponedArgumentsAnalyzer(transformer.resolutionContext)
+
+            call.replaceCalleeReference(FirNamedReferenceWithCandidate(null, firstCandidate.callInfo.name, firstCandidate))
+            val results = postponedArgumentsAnalyzer.analyzeLambda(
+                firstCandidate.system.asPostponedArgumentsAnalyzerContext(),
+                firstAtom,
+                firstCandidate
+            )
+            postponedArgumentsAnalyzer.applyResultsOfAnalyzedLambdaToCandidateSystem(
+                firstCandidate.system.asPostponedArgumentsAnalyzerContext(),
+                firstAtom,
+                firstCandidate,
+                results
+            )
+            while (iterator.hasNext()) {
+                val (candidate, atom) = iterator.next()
+                call.replaceCalleeReference(FirNamedReferenceWithCandidate(null, candidate.callInfo.name, candidate))
+                postponedArgumentsAnalyzer.applyResultsOfAnalyzedLambdaToCandidateSystem(
+                    candidate.system.asPostponedArgumentsAnalyzerContext(),
+                    atom,
+                    candidate,
+                    results
+                )
+            }
+
+            val errorCandidates = mutableSetOf<Candidate>()
+            val successfulCandidates = mutableSetOf<Candidate>()
+
+            for (candidate in candidates) {
+                if (candidate.isSuccessful()) {
+                    successfulCandidates += candidate
+                } else {
+                    errorCandidates += candidate
+                }
+            }
+            return when {
+                successfulCandidates.isNotEmpty() -> successfulCandidates
+                else -> errorCandidates
+            }
+        } finally {
+            call.replaceCalleeReference(originalCalleeReference)
+        }
+    }
+
     fun <T : FirQualifiedAccess> resolveVariableAccessAndSelectCandidate(qualifiedAccess: T): FirStatement {
         val callee = qualifiedAccess.calleeReference as? FirSimpleNamedReference ?: return qualifiedAccess
 

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
@@ -324,3 +325,19 @@ fun FirAnnotationCall.getCorrespondingClassSymbolOrNull(session: FirSession): Fi
         }
     }
 }
+
+fun <T> BodyResolveComponents.initialTypeOfCandidate(
+    candidate: Candidate,
+    call: T
+): ConeKotlinType where T : FirResolvable, T : FirStatement {
+    return initialTypeOfCandidate(candidate, typeFromCallee(call))
+}
+
+fun BodyResolveComponents.initialTypeOfCandidate(candidate: Candidate): ConeKotlinType {
+    val typeRef = typeFromSymbol(candidate.symbol, makeNullable = false)
+    return initialTypeOfCandidate(candidate, typeRef)
+}
+
+private fun initialTypeOfCandidate(candidate: Candidate, typeRef: FirResolvedTypeRef): ConeKotlinType {
+    return candidate.substitutor.substituteOrSelf(typeRef.type)
+}

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

@@ -31,6 +31,8 @@ import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
+import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 
 data class CallInfo(
     val callKind: CallKind,
@@ -109,6 +111,16 @@ class Candidate(
 
     val diagnostics: MutableList<ResolutionDiagnostic> = mutableListOf()
 
+    fun addDiagnostic(diagnostic: ResolutionDiagnostic) {
+        diagnostics += diagnostic
+    }
+
+    fun isSuccessful(): Boolean {
+        if (system.hasContradiction) return false
+        val currentApplicability = diagnostics.map { it.applicability }.minOrNull() ?: CandidateApplicability.RESOLVED
+        return currentApplicability.isSuccess
+    }
+
     var passedStages: Int = 0
 
     fun dispatchReceiverExpression(): FirExpression = when (explicitReceiverKind) {

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

@@ -60,3 +60,5 @@ object ResolvedWithLowPriority : ResolutionDiagnostic(CandidateApplicability.RES
 object InapplicableWrongReceiver : ResolutionDiagnostic(CandidateApplicability.INAPPLICABLE_WRONG_RECEIVER)
 
 object LowerPriorityToPreserveCompatibilityDiagnostic : ResolutionDiagnostic(CandidateApplicability.RESOLVED_NEED_PRESERVE_COMPATIBILITY)
+
+object CandidateChosenUsingOverloadResolutionByLambdaAnnotation : ResolutionDiagnostic(CandidateApplicability.RESOLVED)

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

@@ -12,8 +12,11 @@ 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.resolve.ResolutionMode
+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.ConeArgumentConstraintPosition
+import org.jetbrains.kotlin.fir.resolve.initialTypeOfCandidate
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
@@ -25,14 +28,17 @@ import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
+import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
+import org.jetbrains.kotlin.types.model.safeSubstitute
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirCallCompleter(
@@ -52,9 +58,7 @@ class FirCallCompleter(
 
         val reference = call.calleeReference as? FirNamedReferenceWithCandidate ?: return CompletionResult(call, true)
         val candidate = reference.candidate
-        val initialSubstitutor = candidate.substitutor
+        val initialType = components.initialTypeOfCandidate(candidate, call)
-
-        val initialType = initialSubstitutor.substituteOrSelf(typeRef.type)
 
         if (call is FirExpression) {
             call.resultType = typeRef.resolvedTypeFromPrototype(initialType)
@@ -72,15 +76,7 @@ class FirCallCompleter(
         return when (completionMode) {
             ConstraintSystemCompletionMode.FULL -> {
                 if (inferenceSession.shouldRunCompletion(call)) {
-                    completer.complete(
+                    runCompletionForCall(candidate, completionMode, call, initialType, analyzer)
-                        candidate.system.asConstraintSystemCompleterContext(),
-                        completionMode,
-                        listOf(call),
-                        initialType,
-                        transformer.resolutionContext
-                    ) {
-                        analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
-                    }
                     val finalSubstitutor = candidate.system.asReadOnlyStorage()
                         .buildAbstractResultingSubstitutor(session.inferenceComponents.ctx) as ConeSubstitutor
                     val completedCall = call.transformSingle(
@@ -99,15 +95,7 @@ class FirCallCompleter(
             }
 
             ConstraintSystemCompletionMode.PARTIAL -> {
-                completer.complete(
+                runCompletionForCall(candidate, completionMode, call, initialType, analyzer)
-                    candidate.system.asConstraintSystemCompleterContext(),
-                    completionMode,
-                    listOf(call),
-                    initialType,
-                    transformer.resolutionContext
-                ) {
-                    analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
-                }
                 inferenceSession.addPartiallyResolvedCall(call)
                 CompletionResult(call, false)
             }
@@ -116,6 +104,47 @@ class FirCallCompleter(
         }
     }
 
+    fun <T> runCompletionForCall(
+        candidate: Candidate,
+        completionMode: ConstraintSystemCompletionMode,
+        call: T,
+        initialType: ConeKotlinType,
+        analyzer: PostponedArgumentsAnalyzer? = null
+    ) where T : FirResolvable, T : FirStatement {
+        @Suppress("NAME_SHADOWING")
+        val analyzer = analyzer ?: createPostponedArgumentsAnalyzer(transformer.resolutionContext)
+        completer.complete(
+            candidate.system.asConstraintSystemCompleterContext(),
+            completionMode,
+            listOf(call),
+            initialType,
+            transformer.resolutionContext
+        ) {
+            analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
+        }
+    }
+
+    fun prepareLambdaAtomForFactoryPattern(
+        atom: ResolvedLambdaAtom,
+        candidate: Candidate
+    ) {
+        val returnVariable = ConeTypeVariableForLambdaReturnType(atom.atom, "_R")
+        val csBuilder = candidate.system.getBuilder()
+        csBuilder.registerVariable(returnVariable)
+        val functionalType = csBuilder.buildCurrentSubstitutor()
+            .safeSubstitute(csBuilder.asConstraintSystemCompleterContext(), atom.expectedType!!) as ConeClassLikeType
+        val size = functionalType.typeArguments.size
+        val expectedType = ConeClassLikeTypeImpl(
+            functionalType.lookupTag,
+            Array(size) { index -> if (index != size - 1) functionalType.typeArguments[index] else returnVariable.defaultType },
+            isNullable = functionalType.isNullable,
+            functionalType.attributes
+        )
+        csBuilder.addSubtypeConstraint(expectedType, functionalType, ConeArgumentConstraintPosition())
+        atom.replaceExpectedType(expectedType)
+        atom.replaceTypeVariableForLambdaReturnType(returnVariable)
+    }
+
     fun createCompletionResultsWriter(
         substitutor: ConeSubstitutor,
         mode: FirCallCompletionResultsWriterTransformer.Mode = FirCallCompletionResultsWriterTransformer.Mode.Normal

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

@@ -55,7 +55,7 @@ class PostponedArgumentsAnalyzer(
         argument: PostponedResolvedAtom,
         candidate: Candidate
     ) {
-        return when (argument) {
+        when (argument) {
             is ResolvedLambdaAtom ->
                 analyzeLambda(c, argument, candidate)
 
@@ -97,12 +97,12 @@ class PostponedArgumentsAnalyzer(
         }
     }
 
-    private fun analyzeLambda(
+    fun analyzeLambda(
         c: PostponedArgumentsAnalyzerContext,
         lambda: ResolvedLambdaAtom,
         candidate: Candidate
         //diagnosticHolder: KotlinDiagnosticsHolder
-    ) {
+    ): ReturnArgumentsAnalysisResult {
         val unitType = components.session.builtinTypes.unitType.type
         val stubsForPostponedVariables = c.bindingStubsForPostponedVariables()
         val currentSubstitutor = c.buildCurrentSubstitutor(stubsForPostponedVariables.mapKeys { it.key.freshTypeConstructor(c) })
@@ -122,7 +122,7 @@ class PostponedArgumentsAnalyzer(
             else -> null
         }
 
-        val (returnArguments, inferenceSession) = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
+        val results = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
             lambda,
             receiver,
             parameters,
@@ -130,7 +130,18 @@ class PostponedArgumentsAnalyzer(
             rawReturnType,
             stubsForPostponedVariables
         )
+        applyResultsOfAnalyzedLambdaToCandidateSystem(c, lambda, candidate, results, ::substitute)
+        return results
+    }
 
+    fun applyResultsOfAnalyzedLambdaToCandidateSystem(
+        c: PostponedArgumentsAnalyzerContext,
+        lambda: ResolvedLambdaAtom,
+        candidate: Candidate,
+        results: ReturnArgumentsAnalysisResult,
+        substitute: (ConeKotlinType) -> ConeKotlinType = c.createSubstituteFunctorForLambdaAnalysis()
+    ) {
+        val (returnArguments, inferenceSession) = results
         if (inferenceSession != null) {
             val storageSnapshot = c.getBuilder().currentStorage()
 
@@ -154,7 +165,7 @@ class PostponedArgumentsAnalyzer(
         val checkerSink: CheckerSink = CheckerSinkImpl()
 
         var hasExpressionInReturnArguments = false
-        val lambdaReturnType = lambda.returnType.let(::substitute).takeUnless { it.isUnit }
+        val lambdaReturnType = lambda.returnType.let(substitute).takeUnless { it.isUnit }
         returnArguments.forEach {
             if (it !is FirExpression) return@forEach
             hasExpressionInReturnArguments = true
@@ -172,12 +183,22 @@ class PostponedArgumentsAnalyzer(
 
         if (!hasExpressionInReturnArguments && lambdaReturnType != null) {
             /*LambdaArgumentConstraintPosition(lambda)*/
-            c.getBuilder().addEqualityConstraint(lambdaReturnType, unitType, SimpleConstraintSystemConstraintPosition)
+            c.getBuilder().addEqualityConstraint(
+                lambdaReturnType,
+                components.session.builtinTypes.unitType.type,
+                SimpleConstraintSystemConstraintPosition
+            )
         }
 
         lambda.analyzed = true
         lambda.returnStatements = returnArguments
     }
+
+    fun PostponedArgumentsAnalyzerContext.createSubstituteFunctorForLambdaAnalysis(): (ConeKotlinType) -> ConeKotlinType {
+        val stubsForPostponedVariables = bindingStubsForPostponedVariables()
+        val currentSubstitutor = buildCurrentSubstitutor(stubsForPostponedVariables.mapKeys { it.key.freshTypeConstructor(this) })
+        return { currentSubstitutor.safeSubstitute(this, it) as ConeKotlinType }
+    }
 }
 
 fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(

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

@@ -41,12 +41,12 @@ sealed class PostponedResolvedAtom : PostponedResolvedAtomMarker {
 
 class ResolvedLambdaAtom(
     val atom: FirAnonymousFunction,
-    override val expectedType: ConeKotlinType?,
+    expectedType: ConeKotlinType?,
     val isSuspend: Boolean,
     val receiver: ConeKotlinType?,
     val parameters: List<ConeKotlinType>,
     val returnType: ConeKotlinType,
-    val typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType?,
+    typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType?,
     candidateOfOuterCall: Candidate?
 ) : PostponedResolvedAtom() {
     init {
@@ -55,10 +55,24 @@ class ResolvedLambdaAtom(
         }
     }
 
+    var typeVariableForLambdaReturnType = typeVariableForLambdaReturnType
+        private set
+
+    override var expectedType = expectedType
+        private set
+
     lateinit var returnStatements: Collection<FirStatement>
 
     override val inputTypes: Collection<ConeKotlinType> get() = receiver?.let { parameters + it } ?: parameters
     override val outputType: ConeKotlinType get() = returnType
+
+    fun replaceExpectedType(expectedType: ConeKotlinType) {
+        this.expectedType = expectedType
+    }
+
+    fun replaceTypeVariableForLambdaReturnType(typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType) {
+        this.typeVariableForLambdaReturnType = typeVariableForLambdaReturnType
+    }
 }
 
 class LambdaWithTypeVariableAsExpectedTypeAtom(

compiler/testData/diagnostics/testsWithStdLib/factoryPattern/overloadByLambdaReturnType_disabled.fir.kt

@@ -26,7 +26,7 @@ fun test_1() {
 
 fun test_2() {
     val x = create { 1 }
-    <!INAPPLICABLE_CANDIDATE!>takeInt<!>(x)
+    takeInt(x)
 }
 
 fun test_3() {
@@ -44,5 +44,5 @@ fun test_4() {
 
 fun test_5() {
     val x = create("") { 1 }
-    <!INAPPLICABLE_CANDIDATE!>takeInt<!>(x)
+    takeInt(x)
 }

compiler/testData/diagnostics/testsWithStdLib/factoryPattern/overloadByLambdaReturnType_enabled.fir.kt

@@ -26,7 +26,7 @@ fun test_1() {
 
 fun test_2() {
     val x = create { 1 }
-    <!INAPPLICABLE_CANDIDATE!>takeInt<!>(x)
+    takeInt(x)
 }
 
 fun test_3() {
@@ -44,7 +44,7 @@ fun test_4() {
 
 fun test_5() {
     val x = create("") { 1 }
-    <!INAPPLICABLE_CANDIDATE!>takeInt<!>(x)
+    takeInt(x)
 }
 
 interface A

core/compiler.common/src/org/jetbrains/kotlin/resolve/descriptorUtil/annotationsForResolve.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.descriptorUtil
 
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
@@ -15,7 +16,9 @@ val HIDES_MEMBERS_ANNOTATION_FQ_NAME = FqName("kotlin.internal.HidesMembers")
 val ONLY_INPUT_TYPES_FQ_NAME = FqName("kotlin.internal.OnlyInputTypes")
 val DYNAMIC_EXTENSION_FQ_NAME = FqName("kotlin.internal.DynamicExtension")
 val BUILDER_INFERENCE_ANNOTATION_FQ_NAME = FqName("kotlin.BuilderInference")
-val OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_FQ_NAME = FqName.fromSegments(listOf("kotlin", "OverloadResolutionByLambdaReturnType"))
+
+val OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_CLASS_ID = ClassId(FqName("kotlin"), Name.identifier("OverloadResolutionByLambdaReturnType"))
+val OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_FQ_NAME = OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_CLASS_ID.asSingleFqName()
 
 // @HidesMembers annotation only has effect for members with these names
 val HIDES_MEMBERS_NAME_LIST = setOf(Name.identifier("forEach"), Name.identifier("addSuppressed"))