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

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -96,7 +96,6 @@ class KotlinResolutionCallbacksImpl(
         expectedReturnType: UnwrappedType?,
         annotations: Annotations,
         stubsForPostponedVariables: Map<NewTypeVariable, StubType>,
-        shouldRunInIndependentContext: Boolean
     ): ReturnArgumentsAnalysisResult {
         val psiCallArgument = lambdaArgument.psiCallArgument as PSIFunctionKotlinCallArgument
         val outerCallContext = psiCallArgument.outerCallContext
@@ -138,7 +137,7 @@ class KotlinResolutionCallbacksImpl(
 
         val lambdaInfo = LambdaInfo(
             expectedReturnType ?: TypeUtils.NO_EXPECTED_TYPE,
-            if (expectedReturnType != null || shouldRunInIndependentContext) ContextDependency.INDEPENDENT else ContextDependency.DEPENDENT
+            if (expectedReturnType == null) ContextDependency.DEPENDENT else ContextDependency.INDEPENDENT
         )
 
         val builtIns = outerCallContext.scope.ownerDescriptor.builtIns
@@ -198,10 +197,8 @@ class KotlinResolutionCallbacksImpl(
         val functionTypeInfo = expressionTypingServices.getTypeInfo(psiCallArgument.expression, actualContext)
         (temporaryTrace ?: trace).record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, psiCallArgument.ktFunction, LambdaInfo.STUB_EMPTY)
 
-        val inferedReturnType = functionTypeInfo.type?.arguments?.last()?.type?.takeIf { functionTypeInfo.type?.isFunctionTypeOrSubtype == true }
-
         if (coroutineSession?.hasInapplicableCall() == true) {
-            return ReturnArgumentsAnalysisResult(ReturnArgumentsInfo.empty, coroutineSession, inferedReturnType, hasInapplicableCallForBuilderInference = true)
+            return ReturnArgumentsAnalysisResult(ReturnArgumentsInfo.empty, coroutineSession, hasInapplicableCallForBuilderInference = true)
         } else {
             temporaryTrace?.commit()
         }
@@ -249,7 +246,6 @@ class KotlinResolutionCallbacksImpl(
                 returnArgumentFound
             ),
             coroutineSession,
-            inferedReturnType
         )
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -181,6 +181,8 @@ class ResolvedAtomCompleter(
 
     private val ResolvedLambdaAtom.isCoercedToUnit: Boolean
         get() {
+            val resultArgumentsInfo = this.resultArgumentsInfo
+                ?: return (subResolvedAtoms!!.single() as ResolvedLambdaAtom).isCoercedToUnit
             val returnTypes =
                 resultArgumentsInfo.nonErrorArguments.map {
                     val type = it.safeAs<SimpleKotlinCallArgument>()?.receiver?.receiverValue?.type ?: return@map null
@@ -199,6 +201,8 @@ class ResolvedAtomCompleter(
         }
 
     private fun completeLambda(lambda: ResolvedLambdaAtom) {
+        val lambda = lambda.unwrap()
+        val resultArgumentsInfo = lambda.resultArgumentsInfo!!
         val returnType = if (lambda.isCoercedToUnit) {
             builtIns.unitType
         } else {
@@ -213,7 +217,7 @@ class ResolvedAtomCompleter(
             )
         updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType)
 
-        for (lambdaResult in lambda.resultArgumentsInfo.nonErrorArguments) {
+        for (lambdaResult in resultArgumentsInfo.nonErrorArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
             val newContext =
                 topLevelCallContext.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)