[FIR] Extract overload by lambda return type into separate component

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

@@ -5,7 +5,6 @@
 
 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
@@ -21,7 +20,8 @@ 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.*
+import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
+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
@@ -36,20 +36,18 @@ 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: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     private val qualifiedResolver: FirQualifiedNameResolver,
 ) {
     private val session = components.session
+    private val overloadByLambdaReturnTypeResolver = FirOverloadByLambdaReturnTypeResolver(components)
 
     private lateinit var transformer: FirExpressionsResolveTransformer
 
@@ -61,7 +59,7 @@ class FirCallResolver(
         components, components.resolutionStageRunner,
     )
 
-    private val conflictResolver: ConeCallConflictResolver =
+    val conflictResolver: ConeCallConflictResolver =
         session.callConflictResolverFactory.create(TypeSpecificityComparator.NONE, session.inferenceComponents)
 
     @PrivateForInline
@@ -163,146 +161,11 @@ class FirCallResolver(
             )
         }
 
-        if (
+        reducedCandidates = overloadByLambdaReturnTypeResolver.reduceCandidates(qualifiedAccess, bestCandidates, reducedCandidates)
-            reducedCandidates.size > 1 &&
-            session.languageVersionSettings.supportsFeature(LanguageFeature.OverloadResolutionByLambdaReturnType)
-        ) {
-            /*
-             * Inference session may look into candidate of call, and for that it uses callee reference.
-             * So we need replace reference with proper candidate before calling inference session
-             */
-            val shouldRunCompletion = if (components.context.inferenceSession != FirInferenceSession.DEFAULT) {
-                var shouldRunCompletion = true
-                val originalReference = qualifiedAccess.calleeReference
-                val inferenceSession = components.context.inferenceSession
-                for (candidate in bestCandidates) {
-                    qualifiedAccess.replaceCalleeReference(FirNamedReferenceWithCandidate(null, candidate.callInfo.name, candidate))
-                    shouldRunCompletion = shouldRunCompletion && inferenceSession.shouldRunCompletion(qualifiedAccess)
-                    if (!shouldRunCompletion) break
-                }
-                qualifiedAccess.replaceCalleeReference(originalReference)
-                shouldRunCompletion
-            } else {
-                true
-            }
-            if (shouldRunCompletion) {
-                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/FirOverloadByLambdaReturnTypeResolver.kt

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir
+
+import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.fir.expressions.FirResolvable
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.CandidateChosenUsingOverloadResolutionByLambdaAnnotation
+import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
+import org.jetbrains.kotlin.fir.resolve.inference.FirInferenceSession
+import org.jetbrains.kotlin.fir.resolve.inference.ResolvedLambdaAtom
+import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
+import org.jetbrains.kotlin.fir.resolve.initialTypeOfCandidate
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
+import org.jetbrains.kotlin.resolve.descriptorUtil.OVERLOAD_RESOLUTION_BY_LAMBDA_ANNOTATION_CLASS_ID
+import org.jetbrains.kotlin.utils.addToStdlib.same
+
+class FirOverloadByLambdaReturnTypeResolver(
+    val components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents
+) {
+    private val session = components.session
+    private val callCompleter: FirCallCompleter
+        get() = components.callCompleter
+    private val inferenceSession: FirInferenceSession
+        get() = components.transformer.context.inferenceSession
+
+    fun <T> reduceCandidates(
+        qualifiedAccess: T,
+        allCandidates: Collection<Candidate>,
+        bestCandidates: Set<Candidate>
+    ): Set<Candidate> where T : FirStatement, T : FirResolvable {
+        if (
+            bestCandidates.size <= 1 ||
+            !session.languageVersionSettings.supportsFeature(LanguageFeature.OverloadResolutionByLambdaReturnType)
+        ) return bestCandidates
+
+        /*
+         * Inference session may look into candidate of call, and for that it uses callee reference.
+         * So we need replace reference with proper candidate before calling inference session
+         */
+        val shouldRunCompletion = if (inferenceSession != FirInferenceSession.DEFAULT) {
+            var shouldRunCompletion = true
+            val originalReference = qualifiedAccess.calleeReference
+            for (candidate in bestCandidates) {
+                qualifiedAccess.replaceCalleeReference(FirNamedReferenceWithCandidate(null, candidate.callInfo.name, candidate))
+                shouldRunCompletion = shouldRunCompletion && inferenceSession.shouldRunCompletion(qualifiedAccess)
+                if (!shouldRunCompletion) break
+            }
+            qualifiedAccess.replaceCalleeReference(originalReference)
+            shouldRunCompletion
+        } else {
+            true
+        }
+
+        if (!shouldRunCompletion) return bestCandidates
+
+        return reduceCandidatesImpl(
+            qualifiedAccess,
+            bestCandidates,
+            allCandidates
+        ) ?: bestCandidates
+    }
+
+    private fun <T> reduceCandidatesImpl(
+        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 = components.callResolver.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))
+            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) ->
+                callCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate)
+            }
+            val iterator = lambdas.entries.iterator()
+            val (firstCandidate, firstAtom) = iterator.next()
+
+            val postponedArgumentsAnalyzer = callCompleter.createPostponedArgumentsAnalyzer(
+                components.transformer.resolutionContext
+            )
+
+            call.replaceCalleeReference(FirNamedReferenceWithCandidate(null, firstCandidate.callInfo.name, firstCandidate))
+            val results = postponedArgumentsAnalyzer.analyzeLambda(
+                firstCandidate.system.asPostponedArgumentsAnalyzerContext(),
+                firstAtom,
+                firstCandidate
+            )
+            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)
+        }
+    }
+}