[NI] Take into account return without expression in lambdas

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

@@ -63,7 +63,7 @@ class KotlinResolutionCallbacksImpl(
 
     class LambdaInfo(val expectedType: UnwrappedType, val contextDependency: ContextDependency) {
         var dataFlowInfoAfter: DataFlowInfo? = null
-        val returnStatements = ArrayList<Pair<KtReturnExpression, KotlinTypeInfo>>()
+        val returnStatements = ArrayList<Pair<KtReturnExpression, KotlinTypeInfo?>>()
 
         companion object {
             val STUB_EMPTY = LambdaInfo(TypeUtils.NO_EXPECTED_TYPE, ContextDependency.INDEPENDENT)
@@ -86,7 +86,6 @@ class KotlinResolutionCallbacksImpl(
                 createSimplePSICallArgument(trace.bindingContext, outerCallContext.statementFilter, outerCallContext.scope.ownerDescriptor,
                                             CallMaker.makeExternalValueArgument(ktExpression), DataFlowInfo.EMPTY, typeInfo)
 
-
         val expression: KtExpression = (psiCallArgument as? LambdaKotlinCallArgumentImpl)?.ktLambdaExpression ?:
                                (psiCallArgument as FunctionExpressionImpl).ktFunction
 
@@ -113,19 +112,31 @@ class KotlinResolutionCallbacksImpl(
         val functionTypeInfo = expressionTypingServices.getTypeInfo(expression, actualContext)
         trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, LambdaInfo.STUB_EMPTY)
 
-        val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
-            if (expectedReturnType?.isUnit() == true) return@let null // coercion to Unit
-
-            val lastExpressionType = trace.getType(lastExpression)
-            val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
-            createCallArgument(lastExpression, lastExpressionTypeInfo)
-        }
-
+        var hasReturnWithoutExpression = false
         val returnArguments = lambdaInfo.returnStatements.mapNotNullTo(ArrayList()) { (expression, typeInfo) ->
-            expression.returnedExpression?.let { createCallArgument(it, typeInfo) }
+            val returnedExpression = expression.returnedExpression
+            if (returnedExpression != null) {
+                createCallArgument(
+                        returnedExpression,
+                        typeInfo ?: throw AssertionError("typeInfo should be non-null for return with expression")
+                )
+            }
+            else {
+                hasReturnWithoutExpression = true
+                EmptyLabeledReturn(expression, builtIns)
+            }
         }
 
-        returnArguments.addIfNotNull(lastExpressionArgument)
+        if (!hasReturnWithoutExpression) {
+            val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
+                if (expectedReturnType?.isUnit() == true) return@let null // coercion to Unit
+
+                val lastExpressionType = trace.getType(lastExpression)
+                val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
+                createCallArgument(lastExpression, lastExpressionTypeInfo)
+            }
+            returnArguments.addIfNotNull(lastExpressionArgument)
+        }
 
         return returnArguments
     }

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

@@ -166,7 +166,7 @@ class KotlinToResolvedCallTransformer(
             updateTraceForLambdaReturnType(lambdaArgument, trace, returnType)
 
             for (lambdaResult in lambdaArgument.resultArguments) {
-                val resultValueArgument = lambdaResult.psiCallArgument
+                val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
                 val newContext =
                         context.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)
                                 .replaceExpectedType(returnType)

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

@@ -159,6 +159,16 @@ class FakeValueArgumentForLeftCallableReference(val ktExpression: KtCallableRefe
     override fun isExternal(): Boolean = false
 }
 
+class EmptyLabeledReturn(
+        val returnExpression: KtReturnExpression,
+        builtIns: KotlinBuiltIns
+) : ExpressionKotlinCallArgument {
+    override val isSpread: Boolean get() = false
+    override val argumentName: Name? get() = null
+    override val receiver = ReceiverValueWithSmartCastInfo(TransientReceiver(builtIns.unitType), emptySet(), true)
+    override val isSafeCall: Boolean get() = false
+}
+
 // context here is context for value argument analysis
 internal fun createSimplePSICallArgument(
         contextForArgument: BasicCallResolutionContext,

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ControlStructureTypingVisitor.java

@@ -681,6 +681,9 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             if (!noExpectedType(expectedType) && !KotlinBuiltIns.isUnit(expectedType) && !isDontCarePlaceholder(expectedType)) {
                 context.trace.report(RETURN_TYPE_MISMATCH.on(expression, expectedType));
             }
+            if (newInferenceLambdaInfo != null) {
+                newInferenceLambdaInfo.getReturnStatements().add(new kotlin.Pair<>(expression, null));
+            }
         }
         return components.dataFlowAnalyzer.createCheckedTypeInfo(resultType, context, expression);
     }

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

@@ -48,10 +48,3 @@ class ReceiverExpressionKotlinCallArgument private constructor(
         ) = ReceiverExpressionKotlinCallArgument(receiver.prepareReceiverRegardingCaptureTypes(), isSafeCall, isVariableReceiverForInvoke)
     }
 }
-
-class EmptyLabeledReturn(builtIns: KotlinBuiltIns) : ExpressionKotlinCallArgument {
-    override val isSpread: Boolean get() = false
-    override val argumentName: Name? get() = null
-    override val receiver = ReceiverValueWithSmartCastInfo(TransientReceiver(builtIns.unitType), emptySet(), true)
-    override val isSafeCall: Boolean get() = false
-}

compiler/testData/ir/irText/lambdas/nonLocalReturn.kt

@@ -18,6 +18,16 @@ fun test2() {
     }
 }
 
+// TODO we don't see 'return@lambda' inside internal lambda when we analyze an external lambda,
+// so type information from NI is actually incorrect, see KT-18392
+fun test3() {
+    run lambda@{
+        run {
+            return@lambda
+        }
+    }
+}
+
 fun testLrmFoo1(ints: List<Int>) {
     ints.forEach lit@ {
         if (it == 0) return@lit

compiler/testData/ir/irText/lambdas/nonLocalReturn.txt

@@ -29,6 +29,24 @@ FILE /nonLocalReturn.kt
               RETURN type=kotlin.Nothing from='<anonymous>(): Unit'
                 GET_OBJECT 'Unit' type=kotlin.Unit
           FUNCTION_REFERENCE '<anonymous>(): Unit' type=() -> kotlin.Unit origin=LAMBDA
+  FUN public fun test3(): kotlin.Unit
+    BLOCK_BODY
+      CALL 'run(() -> Nothing): Nothing' type=kotlin.Nothing origin=null
+        <R>: Nothing
+        block: BLOCK type=() -> kotlin.Nothing origin=LAMBDA
+          FUN LOCAL_FUNCTION_FOR_LAMBDA local final fun <anonymous>(): kotlin.Nothing
+            BLOCK_BODY
+              RETURN type=kotlin.Nothing from='<anonymous>(): Nothing'
+                CALL 'run(() -> Nothing): Nothing' type=kotlin.Nothing origin=null
+                  <R>: Nothing
+                  block: BLOCK type=() -> kotlin.Nothing origin=LAMBDA
+                    FUN LOCAL_FUNCTION_FOR_LAMBDA local final fun <anonymous>(): kotlin.Nothing
+                      BLOCK_BODY
+                        RETURN type=kotlin.Nothing from='<anonymous>(): Nothing'
+                          TYPE_OP type=kotlin.Nothing origin=IMPLICIT_CAST typeOperand=kotlin.Nothing
+                            GET_OBJECT 'Unit' type=kotlin.Unit
+                    FUNCTION_REFERENCE '<anonymous>(): Nothing' type=() -> kotlin.Nothing origin=LAMBDA
+          FUNCTION_REFERENCE '<anonymous>(): Nothing' type=() -> kotlin.Nothing origin=LAMBDA
   FUN public fun testLrmFoo1(ints: kotlin.collections.List<kotlin.Int>): kotlin.Unit
     VALUE_PARAMETER value-parameter ints: kotlin.collections.List<kotlin.Int>
     BLOCK_BODY