KT-10717 Type inference for lambda with local return

#KT-10717 Fixed

compiler/frontend/src/org/jetbrains/kotlin/resolve/DescriptorResolver.java

@@ -56,6 +56,7 @@ import org.jetbrains.kotlin.storage.StorageManager;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker;
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices;
+import org.jetbrains.kotlin.types.expressions.FunctionsTypingVisitor;
 import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor;
 
 import java.util.*;
@@ -80,6 +81,7 @@ public class DescriptorResolver {
     @NotNull private final ExpressionTypingServices expressionTypingServices;
     @NotNull private final OverloadChecker overloadChecker;
     @NotNull private final LanguageFeatureSettings languageFeatureSettings;
+    @NotNull private final FunctionsTypingVisitor functionsTypingVisitor;
 
     public DescriptorResolver(
             @NotNull AnnotationResolver annotationResolver,
@@ -90,7 +92,8 @@ public class DescriptorResolver {
             @NotNull VariableTypeResolver variableTypeResolver,
             @NotNull ExpressionTypingServices expressionTypingServices,
             @NotNull OverloadChecker overloadChecker,
-            @NotNull LanguageFeatureSettings languageFeatureSettings
+            @NotNull LanguageFeatureSettings languageFeatureSettings,
+            @NotNull FunctionsTypingVisitor functionsTypingVisitor
     ) {
         this.annotationResolver = annotationResolver;
         this.builtIns = builtIns;
@@ -101,6 +104,7 @@ public class DescriptorResolver {
         this.expressionTypingServices = expressionTypingServices;
         this.overloadChecker = overloadChecker;
         this.languageFeatureSettings = languageFeatureSettings;
+        this.functionsTypingVisitor = functionsTypingVisitor;
     }
 
     public List<KotlinType> resolveSupertypes(
@@ -1049,7 +1053,7 @@ public class DescriptorResolver {
     }
 
     @NotNull
-    /*package*/ static DeferredType inferReturnTypeFromExpressionBody(
+    /*package*/ DeferredType inferReturnTypeFromExpressionBody(
             @NotNull StorageManager storageManager,
             @NotNull final ExpressionTypingServices expressionTypingServices,
             @NotNull final BindingTrace trace,
@@ -1064,7 +1068,9 @@ public class DescriptorResolver {
                 PreliminaryDeclarationVisitor.Companion.createForDeclaration(function, trace);
                 KotlinType type = expressionTypingServices.getBodyExpressionType(
                         trace, scope, dataFlowInfo, function, functionDescriptor);
-                return transformAnonymousTypeIfNeeded(functionDescriptor, function, type, trace);
+                KotlinType result = transformAnonymousTypeIfNeeded(functionDescriptor, function, type, trace);
+                functionsTypingVisitor.checkTypesForReturnStatements(function, trace, result);
+                return result;
             }
         });
     }

compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt

@@ -125,7 +125,7 @@ class FunctionDescriptorResolver(
             builtIns.unitType
         }
         else if (function.hasBody()) {
-            inferReturnTypeFromExpressionBody(storageManager, expressionTypingServices, trace, scope,
+            descriptorResolver.inferReturnTypeFromExpressionBody(storageManager, expressionTypingServices, trace, scope,
                                               dataFlowInfo, function, functionDescriptor)
         }
         else {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.diagnostics.Errors.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getAnnotationEntries
@@ -39,9 +40,12 @@ import org.jetbrains.kotlin.types.CommonSupertypes
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 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.isUnit
 import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.*
 
 internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : ExpressionTypingVisitor(facade) {
 
@@ -264,4 +268,45 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
             it.getTargetLabel()?.let { trace.get(BindingContext.LABEL_TARGET, it) } == functionLiteral
         }
     }
+
+    fun checkTypesForReturnStatements(function: KtDeclarationWithBody, trace: BindingTrace, actualReturnType: KotlinType) {
+        if (function.hasBlockBody()) return
+        if ((function !is KtNamedFunction || function.typeReference != null)
+            && (function !is KtPropertyAccessor || function.returnTypeReference == null)) return
+
+        val bodyExpression = function.bodyExpression ?: return
+        val returns = ArrayList<KtReturnExpression>()
+
+        // data == false means, that we inside other function, so ours return should be with label
+        bodyExpression.accept(object : KtTreeVisitor<Boolean>() {
+            override fun visitReturnExpression(expression: KtReturnExpression, data: Boolean): Void? {
+                val label = expression.getTargetLabel()
+                if ((label != null && trace[BindingContext.LABEL_TARGET, label] == function)
+                    || (label == null && data)
+                ) {
+                    returns.add(expression)
+                }
+                return super.visitReturnExpression(expression, data)
+            }
+
+            override fun visitNamedFunction(function: KtNamedFunction, data: Boolean): Void? {
+                return super.visitNamedFunction(function, false)
+            }
+        }, true)
+
+        for (returnForCheck in returns) {
+            val expression = returnForCheck.returnedExpression
+            if (expression == null) {
+                if (!actualReturnType.isUnit()) {
+                    trace.report(Errors.RETURN_TYPE_MISMATCH.on(returnForCheck, actualReturnType))
+                }
+                continue
+            }
+
+            val expressionType = trace.getType(expression) ?: continue
+            if (!KotlinTypeChecker.DEFAULT.isSubtypeOf(expressionType, actualReturnType)) {
+                trace.report(Errors.TYPE_MISMATCH.on(expression, expressionType, actualReturnType))
+            }
+        }
+    }
 }

compiler/testData/diagnostics/tests/controlStructures/ifToAnyDiscriminatingUsages.kt

@@ -32,15 +32,15 @@ fun testReturnFromAnonFun() =
 
 fun <!IMPLICIT_NOTHING_RETURN_TYPE!>testReturn1<!>() =
         run {
-            return if (true) <!IMPLICIT_CAST_TO_ANY!>42<!>
+            return <!TYPE_MISMATCH!>if (true) <!IMPLICIT_CAST_TO_ANY!>42<!>
-                   else <!IMPLICIT_CAST_TO_ANY!>println()<!>
+                   else <!IMPLICIT_CAST_TO_ANY!>println()<!><!>
         }
 
 fun <!IMPLICIT_NOTHING_RETURN_TYPE!>testReturn2<!>() =
         run {
-            return if (true) <!IMPLICIT_CAST_TO_ANY!>42<!>
+            return <!TYPE_MISMATCH!>if (true) <!IMPLICIT_CAST_TO_ANY!>42<!>
                    else if (true) <!IMPLICIT_CAST_TO_ANY!>42<!>
-                   else <!IMPLICIT_CAST_TO_ANY!>println()<!>
+                   else <!IMPLICIT_CAST_TO_ANY!>println()<!><!>
         }
 
 fun testUsage1() =

compiler/testData/diagnostics/tests/controlStructures/kt10717.kt

@@ -0,0 +1,47 @@
+// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNREACHABLE_CODE -UNUSED_PARAMETER -RETURN_NOT_ALLOWED
+
+fun <!IMPLICIT_NOTHING_RETURN_TYPE!>test1<!>() = run {
+    return <!TYPE_MISMATCH(String; Nothing)!>"OK"<!>
+}
+
+fun <!IMPLICIT_NOTHING_RETURN_TYPE!>test2<!>() = run {
+    fun local(): String {
+        return ""
+    }
+    return <!TYPE_MISMATCH(String; Nothing)!>""<!>
+}
+
+inline fun <T, R> Iterable<T>.map(transform: (T) -> R): List<R> = null!!
+fun test3(a: List<String>, b: List<Int>) = a.map {
+    if (it.length == 3) return <!TYPE_MISMATCH(Nothing?; List<Int>)!>null<!>
+    if (it.length == 4) return <!TYPE_MISMATCH(String; List<Int>)!>""<!>
+    if (it.length == 4) return <!TYPE_MISMATCH(Int; List<Int>)!>5<!>
+    if (it.length == 4) return b
+    1
+}
+
+fun test4() = run {
+    fun test5() {
+        return
+
+        <!RETURN_TYPE_MISMATCH!>return@test4<!>
+
+        return <!RETURN_TYPE_MISMATCH!>return@test4<!>
+
+        return <!TYPE_MISMATCH!>fun() { return; return@test4 <!TYPE_MISMATCH!>""<!> }<!>
+    }
+
+    <!RETURN_TYPE_MISMATCH!>return<!>
+    3
+}
+
+val foo: Int
+    get() = run {
+        if (true) return <!TYPE_MISMATCH!>""<!>
+
+        <!RETURN_TYPE_MISMATCH!>return<!>
+    }
+
+fun test(): Int = run {
+    return <!TYPE_MISMATCH!>""<!>
+}

compiler/testData/diagnostics/tests/controlStructures/kt10717.txt

@@ -0,0 +1,9 @@
+package
+
+public val foo: kotlin.Int
+public fun test(): kotlin.Int
+public fun test1(): kotlin.Nothing
+public fun test2(): kotlin.Nothing
+public fun test3(/*0*/ a: kotlin.collections.List<kotlin.String>, /*1*/ b: kotlin.collections.List<kotlin.Int>): kotlin.collections.List<kotlin.Int>
+public fun test4(): kotlin.Int
+public inline fun </*0*/ T, /*1*/ R> kotlin.collections.Iterable<T>.map(/*0*/ transform: (T) -> R): kotlin.collections.List<R>

compiler/testData/diagnostics/tests/controlStructures/whenToAnyDiscriminatingUsages.kt

@@ -26,22 +26,22 @@ fun testResultOfLambda2() =
 
 fun <!IMPLICIT_NOTHING_RETURN_TYPE!>testReturn1<!>() =
         run {
-            return when {
+            return <!TYPE_MISMATCH!>when {
                 true -> <!IMPLICIT_CAST_TO_ANY!>42<!>
                 else -> <!IMPLICIT_CAST_TO_ANY!>println()<!>
-            }
+            }<!>
         }
 
 fun <!IMPLICIT_NOTHING_RETURN_TYPE!>testReturn2<!>() =
         run {
-            return when {
+            return <!TYPE_MISMATCH!>when {
                 true -> <!IMPLICIT_CAST_TO_ANY!>42<!>
                 else ->
                     when {
                         true -> <!IMPLICIT_CAST_TO_ANY!>42<!>
                         else -> <!IMPLICIT_CAST_TO_ANY!>println()<!>
                     }
-            }
+            }<!>
         }
 
 fun testUsage1() =

compiler/testData/diagnostics/tests/implicitNothing.kt

@@ -17,7 +17,7 @@ fun check() {
     <!UNREACHABLE_CODE!>val <!UNUSED_VARIABLE!>x<!> =<!> null!!
 }
 
-fun <!IMPLICIT_NOTHING_RETURN_TYPE!>nonLocalReturn<!>() = run { return }
+fun <!IMPLICIT_NOTHING_RETURN_TYPE!>nonLocalReturn<!>() = run { <!RETURN_TYPE_MISMATCH!>return<!> }
 
 class Klass {
     fun <!IMPLICIT_NOTHING_RETURN_TYPE!>bar<!>() = null!!

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -3762,6 +3762,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt10717.kt")
+            public void testKt10717() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/controlStructures/kt10717.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("kt1075.kt")
             public void testKt1075() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/controlStructures/kt1075.kt");