[NI] Update lambda result arguments after resolution.

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

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
-import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtPsiUtil
@@ -142,31 +141,6 @@ class KotlinResolutionCallbacksImpl(
         kotlinToResolvedCallTransformer.createStubResolvedCallAndWriteItToTrace<CallableDescriptor>(candidate, trace)
     }
 
-    override fun completeLambdaReturnType(lambdaArgument: ResolvedLambdaArgument, returnType: KotlinType) {
-        val psiCallArgument = lambdaArgument.argument.psiCallArgument
-        val ktFunction = when (psiCallArgument) {
-            is LambdaKotlinCallArgumentImpl -> psiCallArgument.ktLambdaExpression.functionLiteral
-            is FunctionExpressionImpl -> psiCallArgument.ktFunction
-            else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument")
-        }
-
-        val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl ?:
-                                 throw AssertionError("No function descriptor for resolved lambda argument")
-        functionDescriptor.setReturnType(returnType)
-
-        for (lambdaResult in lambdaArgument.resultArguments) {
-            val resultValueArgument = lambdaResult.psiCallArgument.valueArgument
-            val deparenthesized = resultValueArgument.getArgumentExpression()?.let {
-                KtPsiUtil.getLastElementDeparenthesized(it, expressionTypingServices.statementFilter)
-            } ?: continue
-
-            val recordedType = trace.getType(deparenthesized)
-            if (recordedType != null && !recordedType.constructor.isDenotable) {
-                argumentTypeResolver.updateResultArgumentTypeIfNotDenotable(trace, expressionTypingServices.statementFilter, returnType, deparenthesized)
-            }
-        }
-    }
-
     override fun completeCallableReference(
             callableReferenceArgument: ResolvedCallableReferenceArgument,
             resultTypeParameters: List<UnwrappedType>

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

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.resolve.calls.tower
 
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.psi.*
@@ -69,6 +70,7 @@ class KotlinToResolvedCallTransformer(
                 for (resolvedCall in allResolvedCalls) {
                     runCallCheckers(resolvedCall, callCheckerContext)
                 }
+                runLambdaArgumentsChecks(context, trace, baseResolvedCall.lambdaArguments)
             }
 
             return result
@@ -147,6 +149,38 @@ class KotlinToResolvedCallTransformer(
     }
 
 
+    private fun runLambdaArgumentsChecks(
+            context: BasicCallResolutionContext,
+            trace: BindingTrace,
+            lambdaArguments: List<ResolvedLambdaArgument>
+    ) {
+        for (lambdaArgument in lambdaArguments) {
+            val returnType = lambdaArgument.finalReturnType
+
+            val psiCallArgument = lambdaArgument.argument.psiCallArgument
+            val ktFunction = when (psiCallArgument) {
+                is LambdaKotlinCallArgumentImpl -> psiCallArgument.ktLambdaExpression.functionLiteral
+                is FunctionExpressionImpl -> psiCallArgument.ktFunction
+                else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument")
+            }
+
+            val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl ?:
+                                     throw AssertionError("No function descriptor for resolved lambda argument")
+            functionDescriptor.setReturnType(returnType)
+
+            for (lambdaResult in lambdaArgument.resultArguments) {
+                val resultValueArgument = lambdaResult.psiCallArgument
+                val newContext =
+                        context.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)
+                                .replaceExpectedType(returnType)
+                                .replaceBindingTrace(trace)
+
+                val argumentExpression = resultValueArgument.valueArgument.getArgumentExpression() ?: continue
+                updateRecordedType(argumentExpression, newContext)
+            }
+        }
+    }
+
     // todo very beginning code
     private fun runArgumentsChecks(
             context: BasicCallResolutionContext,
@@ -168,35 +202,40 @@ class KotlinToResolvedCallTransformer(
                             .replaceCallPosition(callPosition)
                             .replaceBindingTrace(trace)
 
-            // todo
-//            if (valueArgument.isExternal()) continue
+            // todo external argument
 
             val argumentExpression = valueArgument.getArgumentExpression() ?: continue
-            val deparenthesized = argumentExpression.let {
-                KtPsiUtil.getLastElementDeparenthesized(it, context.statementFilter)
-            } ?: continue
-
-            val recordedType = context.trace.getType(deparenthesized)
-            var updatedType = recordedType
-
-            val resolvedCall = deparenthesized.getResolvedCall(trace.bindingContext)
-            if (resolvedCall != null) {
-                updatedType = resolvedCall.resultingDescriptor.returnType ?: updatedType
-            }
-
-            // For the cases like 'foo(1)' the type of '1' depends on expected type (it can be Int, Byte, etc.),
-            // so while the expected type is not known, it's IntegerValueType(1), and should be updated when the expected type is known.
-            if (recordedType != null && !recordedType.constructor.isDenotable) {
-                updatedType = argumentTypeResolver.updateResultArgumentTypeIfNotDenotable(newContext, deparenthesized) ?: updatedType
-            }
-
-            updatedType = updateRecordedTypeForArgument(updatedType, recordedType, argumentExpression, context)
-
+            updateRecordedType(argumentExpression, newContext)
 //            dataFlowAnalyzer.checkType(updatedType, deparenthesized, newContext)
         }
 
     }
 
+    fun updateRecordedType(
+            expression: KtExpression,
+            context: BasicCallResolutionContext
+    ): KotlinType? {
+        val deparenthesized = expression.let {
+            KtPsiUtil.getLastElementDeparenthesized(it, context.statementFilter)
+        } ?: return null
+
+        val recordedType = context.trace.getType(deparenthesized)
+        var updatedType = recordedType
+
+        val resolvedCall = deparenthesized.getResolvedCall(context.trace.bindingContext)
+        if (resolvedCall != null) {
+            updatedType = resolvedCall.resultingDescriptor.returnType ?: updatedType
+        }
+
+        // For the cases like 'foo(1)' the type of '1' depends on expected type (it can be Int, Byte, etc.),
+        // so while the expected type is not known, it's IntegerValueType(1), and should be updated when the expected type is known.
+        if (recordedType != null && !recordedType.constructor.isDenotable) {
+            updatedType = argumentTypeResolver.updateResultArgumentTypeIfNotDenotable(context, deparenthesized) ?: updatedType
+        }
+
+        return updateRecordedTypeForArgument(updatedType, recordedType, expression, context)
+    }
+
     // See CallCompleter#updateRecordedTypeForArgument
     private fun updateRecordedTypeForArgument(
             updatedType: KotlinType?,

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

@@ -37,8 +37,6 @@ interface KotlinResolutionCallbacks {
     // todo this is hack for some client which try to read ResolvedCall from trace before all calls completed
     fun bindStubResolvedCallForCandidate(candidate: KotlinResolutionCandidate)
 
-    fun completeLambdaReturnType(lambdaArgument: ResolvedLambdaArgument, returnType: KotlinType)
-
     fun completeCallableReference(callableReferenceArgument: ResolvedCallableReferenceArgument,
                                   resultTypeParameters: List<UnwrappedType>)
 

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

@@ -114,12 +114,12 @@ class KotlinCallCompleter(
             it.candidate.toCompletedCall(currentSubstitutor)
         }
         c.lambdaArguments.forEach {
-            resolutionCallbacks.completeLambdaReturnType(it, currentSubstitutor.safeSubstitute(it.returnType))
+            it.finalReturnType = currentSubstitutor.safeSubstitute(it.returnType)
         }
         c.callableReferenceArguments.forEach {
             resolutionCallbacks.completeCallableReference(it, it.myTypeVariables.map { currentSubstitutor.safeSubstitute(it.defaultType) })
         }
-        return ResolvedKotlinCall.CompletedResolvedKotlinCall(completedCall, competedCalls)
+        return ResolvedKotlinCall.CompletedResolvedKotlinCall(completedCall, competedCalls, c.lambdaArguments)
     }
 
     private fun KotlinResolutionCandidate.toCompletedCall(substitutor: NewTypeSubstitutor): CompletedKotlinCall {

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

@@ -48,6 +48,7 @@ class ResolvedLambdaArgument(
     override val outputType: UnwrappedType get() = returnType
 
     lateinit var resultArguments: List<KotlinCallArgument>
+    lateinit var finalReturnType: UnwrappedType
 }
 
 class ResolvedCallableReferenceArgument(

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

@@ -27,7 +27,8 @@ import org.jetbrains.kotlin.types.UnwrappedType
 sealed class ResolvedKotlinCall {
     class CompletedResolvedKotlinCall(
             val completedCall: CompletedKotlinCall,
-            val allInnerCalls: Collection<CompletedKotlinCall>
+            val allInnerCalls: Collection<CompletedKotlinCall>,
+            val lambdaArguments: List<ResolvedLambdaArgument>
     ): ResolvedKotlinCall()
 
     class OnlyResolvedKotlinCall(