[NI] Use types and systems from return arguments instead of return type of lambda

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ExternalComponents.kt

@@ -59,7 +59,6 @@ data class ReturnArgumentsInfo(
 data class ReturnArgumentsAnalysisResult(
     val returnArgumentsInfo: ReturnArgumentsInfo,
     val inferenceSession: InferenceSession?,
-    val lambdaReturnType: KotlinType? = null,
     val hasInapplicableCallForBuilderInference: Boolean = false
 )
 
@@ -73,7 +72,6 @@ interface KotlinResolutionCallbacks {
         expectedReturnType: UnwrappedType?, // null means, that return type is not proper i.e. it depends on some type variables
         annotations: Annotations,
         stubsForPostponedVariables: Map<NewTypeVariable, StubType>,
-        shouldRunInIndependentContext: Boolean = false
     ): ReturnArgumentsAnalysisResult
 
     fun bindStubResolvedCallForCandidate(candidate: ResolvedCallAtom)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/KotlinCallCompleter.kt

@@ -105,7 +105,9 @@ class KotlinCallCompleter(
 
         val newAtoms = mutableMapOf<KotlinResolutionCandidate, ResolvedLambdaAtom>()
         for ((candidate, atom) in lambdas.entries) {
-            newAtoms[candidate] = kotlinConstraintSystemCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate, candidate)
+            val newAtom = kotlinConstraintSystemCompleter.prepareLambdaAtomForFactoryPattern(atom, candidate, candidate)
+            newAtoms[candidate] = newAtom
+            candidate.addResolvedKtPrimitive(newAtom)
         }
 
         val diagnosticHolderForLambda = KotlinDiagnosticsHolder.SimpleHolder()
@@ -116,18 +118,16 @@ class KotlinCallCompleter(
             resolutionCallbacks,
             firstAtom,
             diagnosticHolderForLambda,
-            shouldRunInIndependentContext = true
         )
 
-        val lambdaReturnType = results.lambdaReturnType ?: return candidates
-        lambdas.getValue(firstCandidate).setAnalyzedResults(results.returnArgumentsInfo, listOf(firstAtom))
-        firstCandidate.csBuilder.addSubtypeConstraint(lambdaReturnType, firstAtom.returnType, LambdaArgumentConstraintPosition(firstAtom))
-
         while (iterator.hasNext()) {
             val (candidate, atom) = iterator.next()
-            atom.setAnalyzedResults(results.returnArgumentsInfo, firstAtom.subResolvedAtoms!!)
-            lambdas.getValue(candidate).setAnalyzedResults(results.returnArgumentsInfo, listOf(atom))
-            candidate.csBuilder.addSubtypeConstraint(lambdaReturnType, atom.returnType, LambdaArgumentConstraintPosition(atom))
+            postponedArgumentsAnalyzer.applyResultsOfAnalyzedLambdaToCandidateSystem(
+                candidate.getSystem().asPostponedArgumentsAnalyzerContext(),
+                atom,
+                results,
+                diagnosticHolderForLambda
+            )
         }
 
         val errorCandidates = mutableSetOf<KotlinResolutionCandidate>()

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponeArgumentsChecks.kt

@@ -226,16 +226,16 @@ fun ResolvedLambdaAtom.transformToResolvedLambda(
     expectedType: UnwrappedType,
     returnTypeVariable: TypeVariableForLambdaReturnType? = null
 ): ResolvedLambdaAtom {
-    val resolvedLambdaAtom = preprocessLambdaArgument(
+    return preprocessLambdaArgument(
         csBuilder,
         atom,
         expectedType,
         diagnosticsHolder,
         forceResolution = true,
         returnTypeVariable = returnTypeVariable
-    ) as ResolvedLambdaAtom
-
-    return resolvedLambdaAtom
+    ).also {
+        this.setAnalyzedResults(null, listOf(it))
+    } as ResolvedLambdaAtom
 }
 
 private fun preprocessCallableReference(

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt

@@ -66,23 +66,33 @@ class PostponedArgumentsAnalyzer(
         }
     }
 
+    data class SubstitutorAndStubsForLambdaAnalysis(
+        val stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>,
+        val substitute: (KotlinType) -> UnwrappedType
+    )
+
+    fun Context.createSubstituteFunctorForLambdaAnalysis(): SubstitutorAndStubsForLambdaAnalysis {
+        val stubsForPostponedVariables = bindingStubsForPostponedVariables()
+        val currentSubstitutor = buildCurrentSubstitutor(stubsForPostponedVariables.mapKeys { it.key.freshTypeConstructor(this) })
+        return SubstitutorAndStubsForLambdaAnalysis(stubsForPostponedVariables) {
+            currentSubstitutor.safeSubstitute(this, it) as UnwrappedType
+        }
+    }
+
     fun analyzeLambda(
         c: Context,
         resolutionCallbacks: KotlinResolutionCallbacks,
         lambda: ResolvedLambdaAtom,
         diagnosticHolder: KotlinDiagnosticsHolder,
-        shouldRunInIndependentContext: Boolean = false
     ): ReturnArgumentsAnalysisResult {
-        val stubsForPostponedVariables = c.bindingStubsForPostponedVariables()
-        val currentSubstitutor = c.buildCurrentSubstitutor(stubsForPostponedVariables.mapKeys { it.key.freshTypeConstructor(c) })
-
-        fun substitute(type: UnwrappedType) = currentSubstitutor.safeSubstitute(c, type) as UnwrappedType
+        val substitutorAndStubsForLambdaAnalysis = c.createSubstituteFunctorForLambdaAnalysis()
+        val substitute = substitutorAndStubsForLambdaAnalysis.substitute
 
         // Expected type has a higher priority against which lambda should be analyzed
         // Mostly, this is needed to report more specific diagnostics on lambda parameters
         fun expectedOrActualType(expected: UnwrappedType?, actual: UnwrappedType?): UnwrappedType? {
-            val expectedSubstituted = expected?.let(::substitute)
-            return if (expectedSubstituted != null && c.canBeProper(expectedSubstituted)) expectedSubstituted else actual?.let(::substitute)
+            val expectedSubstituted = expected?.let(substitute)
+            return if (expectedSubstituted != null && c.canBeProper(expectedSubstituted)) expectedSubstituted else actual?.let(substitute)
         }
 
         val builtIns = c.getBuilder().builtIns
@@ -104,7 +114,7 @@ class PostponedArgumentsAnalyzer(
         val parameters =
             expectedParametersToMatchAgainst?.mapIndexed { index, expected ->
                 expectedOrActualType(expected, lambda.parameters.getOrNull(index)) ?: builtIns.nothingType
-            } ?: lambda.parameters.map(::substitute)
+            } ?: lambda.parameters.map(substitute)
 
         val rawReturnType = lambda.returnType
 
@@ -129,15 +139,25 @@ class PostponedArgumentsAnalyzer(
             parameters,
             expectedTypeForReturnArguments,
             convertedAnnotations ?: Annotations.EMPTY,
-            stubsForPostponedVariables.cast(),
-            shouldRunInIndependentContext
+            substitutorAndStubsForLambdaAnalysis.stubsForPostponedVariables.cast(),
         )
-        val (returnArgumentsInfo, inferenceSession, inferedReturnType, hasInapplicableCallForBuilderInference) =
+        applyResultsOfAnalyzedLambdaToCandidateSystem(c, lambda, returnArgumentsAnalysisResult, diagnosticHolder, substitute)
+        return returnArgumentsAnalysisResult
+    }
+
+    fun applyResultsOfAnalyzedLambdaToCandidateSystem(
+        c: Context,
+        lambda: ResolvedLambdaAtom,
+        returnArgumentsAnalysisResult: ReturnArgumentsAnalysisResult,
+        diagnosticHolder: KotlinDiagnosticsHolder,
+        substitute: (KotlinType) -> UnwrappedType = c.createSubstituteFunctorForLambdaAnalysis().substitute
+    ) {
+        val (returnArgumentsInfo, inferenceSession, hasInapplicableCallForBuilderInference) =
             returnArgumentsAnalysisResult
 
         if (hasInapplicableCallForBuilderInference) {
             c.getBuilder().removePostponedVariables()
-            return returnArgumentsAnalysisResult
+            return
         }
 
         val returnArguments = returnArgumentsInfo.nonErrorArguments
@@ -153,7 +173,7 @@ class PostponedArgumentsAnalyzer(
 
         val subResolvedKtPrimitives = allReturnArguments.map {
             resolveKtPrimitive(
-                c.getBuilder(), it, lambda.returnType.let(::substitute),
+                c.getBuilder(), it, lambda.returnType.let(substitute),
                 diagnosticHolder, ReceiverInfo.notReceiver, convertedType = null,
                 inferenceSession
             )
@@ -161,7 +181,7 @@ class PostponedArgumentsAnalyzer(
 
         if (!returnArgumentsInfo.returnArgumentsExist) {
             val unitType = lambda.returnType.builtIns.unitType
-            val lambdaReturnType = lambda.returnType.let(::substitute)
+            val lambdaReturnType = lambda.returnType.let(substitute)
             c.getBuilder().addSubtypeConstraint(unitType, lambdaReturnType, LambdaArgumentConstraintPosition(lambda))
         }
 
@@ -173,7 +193,7 @@ class PostponedArgumentsAnalyzer(
             val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot, diagnosticHolder)
             if (postponedVariables == null) {
                 c.getBuilder().removePostponedVariables()
-                return returnArgumentsAnalysisResult
+                return
             }
 
             for ((constructor, resultType) in postponedVariables) {
@@ -184,8 +204,6 @@ class PostponedArgumentsAnalyzer(
                 c.getBuilder().addEqualityConstraint(variable.defaultType(c), resultType, CoroutinePosition())
             }
         }
-
-        return returnArgumentsAnalysisResult
     }
 
     private fun UnwrappedType?.receiver(): UnwrappedType? {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt

@@ -135,11 +135,17 @@ class ResolvedLambdaAtom(
     val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?,
     override val expectedType: UnwrappedType?
 ) : PostponedResolvedAtom() {
-    lateinit var resultArgumentsInfo: ReturnArgumentsInfo
+    /**
+     * [resultArgumentsInfo] can be null only if lambda was analyzed in process of resolve
+     *   ambiguity by lambda return type
+     * There is a contract that [resultArgumentsInfo] will be not null for unwrapped lambda atom
+     *   (see [unwrap])
+     */
+    var resultArgumentsInfo: ReturnArgumentsInfo? = null
         private set
 
     fun setAnalyzedResults(
-        resultArguments: ReturnArgumentsInfo,
+        resultArguments: ReturnArgumentsInfo?,
         subResolvedAtoms: List<ResolvedAtom>
     ) {
         this.resultArgumentsInfo = resultArguments
@@ -150,6 +156,10 @@ class ResolvedLambdaAtom(
     override val outputType: UnwrappedType get() = returnType
 }
 
+fun ResolvedLambdaAtom.unwrap(): ResolvedLambdaAtom {
+    return if (resultArgumentsInfo != null) this else subResolvedAtoms!!.single() as ResolvedLambdaAtom
+}
+
 abstract class ResolvedCallableReferenceAtom(
     override val atom: CallableReferenceKotlinCallArgument,
     override val expectedType: UnwrappedType?