From 251a02f71d09f550e9cec670f898a0157a55c4c1 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 25 Apr 2019 16:06:56 +0300
Subject: [PATCH] [NI] Fix recording type of lambda in trace

---
 .../kotlin/types/expressions/FunctionsTypingVisitor.kt        | 4 +++-
 .../functionLiterals/return/LocalReturnNoCoercionToUnit.kt    | 2 +-
 .../functionLiterals/return/LocalReturnUnitAndDontCareType.kt | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
index 681cec62fdb..a42933e3aa3 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.resolve.BindingContextUtils
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.FunctionDescriptorUtil
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
+import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableTypeConstructor
 import org.jetbrains.kotlin.resolve.checkers.UnderscoreChecker
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.scopes.LexicalWritableScope
@@ -38,6 +39,7 @@ import org.jetbrains.kotlin.types.TypeUtils.*
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.expressions.CoercionStrategy.COERCION_TO_UNIT
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.createTypeInfo
+import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
@@ -250,7 +252,6 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
 
         newInferenceLambdaInfo?.let {
             it.lastExpressionInfo.dataFlowInfoAfter = blockReturnedType.dataFlowInfo
-            return null
         }
 
         return computeReturnTypeBasedOnReturnExpressions(functionLiteral, context, typeOfBodyExpression)
@@ -290,6 +291,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         returnedExpressionTypes.addIfNotNull(typeOfBodyExpression)
 
         if (returnedExpressionTypes.isEmpty()) return null
+        if (returnedExpressionTypes.any { it.contains { it.constructor is TypeVariableTypeConstructor }}) return null
         return CommonSupertypes.commonSupertype(returnedExpressionTypes)
     }
 
diff --git a/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt b/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt
index bf99e0f31b5..248d8870097 100644
--- a/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt
+++ b/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt
@@ -2,6 +2,6 @@
 fun test(a: Int) {
     run f@{
       if (a > 0) return@f
-      else return@f <!OI;RETURN_TYPE_MISMATCH!>1<!>
+      else return@f <!RETURN_TYPE_MISMATCH!>1<!>
     }
 }
diff --git a/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnitAndDontCareType.kt b/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnitAndDontCareType.kt
index fc388387eef..66e060c468e 100644
--- a/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnitAndDontCareType.kt
+++ b/compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnitAndDontCareType.kt
@@ -17,5 +17,5 @@ val b = run {
 val c = run {
     if (flag) return@run
 
-    return@run <!OI;RETURN_TYPE_MISMATCH!>4<!>
+    return@run <!RETURN_TYPE_MISMATCH!>4<!>
 }
\ No newline at end of file