[NI] Fix updated type for lambda, use original type for substitution

#KT-24367 Fixed

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/CoroutineInferenceSession.kt

@@ -198,8 +198,9 @@ class CoroutineInferenceSession(
         context: BasicCallResolutionContext
     ): ResolvedAtomCompleter {
         return ResolvedAtomCompleter(
-            resultSubstitutor, context.trace, context, kotlinToResolvedCallTransformer, expressionTypingServices,
+            resultSubstitutor, context.trace, context, kotlinToResolvedCallTransformer,
-            argumentTypeResolver, doubleColonExpressionResolver, deprecationResolver, moduleDescriptor, context.dataFlowValueFactory
+            expressionTypingServices, argumentTypeResolver, doubleColonExpressionResolver, builtIns,
+            deprecationResolver, moduleDescriptor, context.dataFlowValueFactory
         )
     }
 }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
@@ -59,7 +60,8 @@ class KotlinToResolvedCallTransformer(
     private val doubleColonExpressionResolver: DoubleColonExpressionResolver,
     private val additionalDiagnosticReporter: AdditionalDiagnosticReporter,
     private val moduleDescriptor: ModuleDescriptor,
-    private val dataFlowValueFactory: DataFlowValueFactory
+    private val dataFlowValueFactory: DataFlowValueFactory,
+    private val builtIns: KotlinBuiltIns
 ) {
     companion object {
         private val REPORT_MISSING_NEW_INFERENCE_DIAGNOSTIC
@@ -113,7 +115,7 @@ class KotlinToResolvedCallTransformer(
                 val resultSubstitutor = baseResolvedCall.constraintSystem.buildResultingSubstitutor()
                 val ktPrimitiveCompleter = ResolvedAtomCompleter(
                     resultSubstitutor, context.trace, context, this, expressionTypingServices, argumentTypeResolver,
-                    doubleColonExpressionResolver, deprecationResolver, moduleDescriptor, dataFlowValueFactory
+                    doubleColonExpressionResolver, builtIns, deprecationResolver, moduleDescriptor, dataFlowValueFactory
                 )
 
                 for (subKtPrimitive in candidate.subResolvedAtoms) {

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.createFunctionType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
@@ -44,6 +46,7 @@ class ResolvedAtomCompleter(
     private val expressionTypingServices: ExpressionTypingServices,
     private val argumentTypeResolver: ArgumentTypeResolver,
     private val doubleColonExpressionResolver: DoubleColonExpressionResolver,
+    private val builtIns: KotlinBuiltIns,
     deprecationResolver: DeprecationResolver,
     moduleDescriptor: ModuleDescriptor,
     private val dataFlowValueFactory: DataFlowValueFactory
@@ -124,7 +127,17 @@ class ResolvedAtomCompleter(
         functionDescriptor.setReturnType(returnType)
 
         val existingLambdaType = trace.getType(ktArgumentExpression) ?: throw AssertionError("No type for resolved lambda argument")
-        trace.recordType(ktArgumentExpression, resultSubstitutor.substituteKeepAnnotations(existingLambdaType.unwrap()))
+        val substitutedFunctionalType = createFunctionType(
+            builtIns,
+            existingLambdaType.annotations,
+            lambda.receiver?.let { resultSubstitutor.substituteKeepAnnotations(it) },
+            lambda.parameters.map { resultSubstitutor.substituteKeepAnnotations(it) },
+            null, // parameter names transforms to special annotations, so they are already taken from parameter types
+            returnType,
+            lambda.isSuspend
+        )
+
+        trace.recordType(ktArgumentExpression, substitutedFunctionalType)
 
         // Mainly this is needed for builder-like inference, when we have type `SomeType<K, V>.() -> Unit` and now we want to update those K, V
         val extensionReceiverParameter = functionDescriptor.extensionReceiverParameter

idea/testData/codeInsight/expressionType/typeOfLambda.kt

@@ -0,0 +1,7 @@
+fun foo() {
+    myMap {<caret> it: Int -> it }
+}
+
+fun <T> myMap(transform: (T) -> T): T = TODO()
+
+// TYPE: { it: Int -> it } -> <html>(Int) &rarr; Int</html>

idea/tests/org/jetbrains/kotlin/idea/codeInsight/ExpressionTypeTestGenerated.java

@@ -109,6 +109,11 @@ public class ExpressionTypeTestGenerated extends AbstractExpressionTypeTest {
         runTest("idea/testData/codeInsight/expressionType/ThisInLambda.kt");
     }
 
+    @TestMetadata("typeOfLambda.kt")
+    public void testTypeOfLambda() throws Exception {
+        runTest("idea/testData/codeInsight/expressionType/typeOfLambda.kt");
+    }
+
     @TestMetadata("VariableDeclaration.kt")
     public void testVariableDeclaration() throws Exception {
         runTest("idea/testData/codeInsight/expressionType/VariableDeclaration.kt");