[NI] Update lambda argument expression type when return type is known

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

@@ -16,6 +16,8 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.builtins.replaceReturnType
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
@@ -40,9 +42,7 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
-import org.jetbrains.kotlin.types.ErrorUtils
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.expressions.DataFlowAnalyzer
 import java.util.*
 
@@ -162,16 +162,7 @@ class KotlinToResolvedCallTransformer(
         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)
+            updateTraceForLambdaReturnType(lambdaArgument, trace, returnType)
 
             for (lambdaResult in lambdaArgument.resultArguments) {
                 val resultValueArgument = lambdaResult.psiCallArgument
@@ -186,6 +177,31 @@ class KotlinToResolvedCallTransformer(
         }
     }
 
+    private fun updateTraceForLambdaReturnType(lambdaArgument: ResolvedLambdaArgument, trace: BindingTrace, returnType: UnwrappedType) {
+        val psiCallArgument = lambdaArgument.argument.psiCallArgument
+
+        val ktArgumentExpression: KtExpression
+        val ktFunction: KtElement
+        when (psiCallArgument) {
+            is LambdaKotlinCallArgumentImpl -> {
+                ktArgumentExpression = psiCallArgument.ktLambdaExpression
+                ktFunction = ktArgumentExpression.functionLiteral
+            }
+            is FunctionExpressionImpl -> {
+                ktArgumentExpression = psiCallArgument.ktFunction
+                ktFunction = ktArgumentExpression
+            }
+            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)
+
+        val existingLambdaType = trace.getType(ktArgumentExpression) ?: throw AssertionError("No type for resolved lambda argument")
+        trace.recordType(ktArgumentExpression, existingLambdaType.replaceReturnType(returnType))
+    }
+
     // todo very beginning code
     private fun runArgumentsChecks(
             context: BasicCallResolutionContext,

core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt

@@ -29,10 +29,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.constants.ConstantValueFactory
 import org.jetbrains.kotlin.resolve.constants.StringValue
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.KotlinTypeFactory
-import org.jetbrains.kotlin.types.SimpleType
-import org.jetbrains.kotlin.types.TypeProjection
+import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.replaceAnnotations
 import org.jetbrains.kotlin.utils.DFS
@@ -122,21 +119,27 @@ private fun FqNameUnsafe.getFunctionalClassKind(): FunctionClassDescriptor.Kind?
 
 
 fun KotlinType.getReceiverTypeFromFunctionType(): KotlinType? {
-    assert(isBuiltinFunctionalType) { "Not a function type: ${this}" }
+    assert(isBuiltinFunctionalType) { "Not a function type: $this" }
     return if (isTypeAnnotatedWithExtensionFunctionType) arguments.first().type else null
 }
 
 fun KotlinType.getReturnTypeFromFunctionType(): KotlinType {
-    assert(isBuiltinFunctionalType) { "Not a function type: ${this}" }
+    assert(isBuiltinFunctionalType) { "Not a function type: $this" }
     return arguments.last().type
 }
 
+fun KotlinType.replaceReturnType(newReturnType: KotlinType): KotlinType {
+    assert(isBuiltinFunctionalType) { "Not a function type: $this"}
+    val argumentsWithNewReturnType = arguments.toMutableList().apply { set(size - 1, TypeProjectionImpl(newReturnType)) }
+    return replace(newArguments = argumentsWithNewReturnType)
+}
+
 fun KotlinType.getValueParameterTypesFromFunctionType(): List<TypeProjection> {
-    assert(isBuiltinFunctionalType) { "Not a function type: ${this}" }
+    assert(isBuiltinFunctionalType) { "Not a function type: $this" }
     val arguments = arguments
     val first = if (isBuiltinExtensionFunctionalType) 1 else 0
     val last = arguments.size - 1
-    assert(first <= last) { "Not an exact function type: ${this}" }
+    assert(first <= last) { "Not an exact function type: $this" }
     return arguments.subList(first, last)
 }