"Change to return with label" quick fix: apply for type mismatch

#KT-32280 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/ChangeToLabeledReturnFix.kt

@@ -9,15 +9,19 @@ import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.util.findLabelAndCall
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 
 class ChangeToLabeledReturnFix(
     element: KtReturnExpression, val labeledReturn: String
@@ -56,9 +60,24 @@ class ChangeToLabeledReturnFix(
         }
 
         override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> {
-            val expression = diagnostic.psiElement as? KtReturnExpression ?: return emptyList()
-            return findAccessibleLabels(expression.analyze(), expression).map {
-                ChangeToLabeledReturnFix(expression, labeledReturn = "return@${it.render()}")
+            val element = diagnostic.psiElement as? KtElement ?: return emptyList()
+            val context by lazy { element.analyze() }
+            val returnExpression = when (diagnostic.factory) {
+                Errors.RETURN_NOT_ALLOWED -> {
+                    diagnostic.psiElement as? KtReturnExpression
+                }
+                Errors.TYPE_MISMATCH, Errors.CONSTANT_EXPECTED_TYPE_MISMATCH, Errors.NULL_FOR_NONNULL_TYPE -> {
+                    val returnExpression = diagnostic.psiElement.getStrictParentOfType<KtReturnExpression>() ?: return emptyList()
+                    val lambda = returnExpression.getStrictParentOfType<KtLambdaExpression>() ?: return emptyList()
+                    val lambdaReturnType = context[BindingContext.FUNCTION, lambda.functionLiteral]?.returnType ?: return emptyList()
+                    val returnType = returnExpression.returnedExpression?.getType(context) ?: return emptyList()
+                    if (!returnType.isSubtypeOf(lambdaReturnType)) return emptyList()
+                    returnExpression
+                }
+                else -> null
+            } ?: return emptyList()
+            return findAccessibleLabels(context, returnExpression).map {
+                ChangeToLabeledReturnFix(returnExpression, labeledReturn = "return@${it.render()}")
             }
         }
     }

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt

@@ -570,6 +570,9 @@ class QuickFixRegistrar : QuickFixContributor {
         MUST_BE_INITIALIZED_OR_BE_ABSTRACT.registerFactory(AddModifierFix.AddLateinitFactory)
 
         RETURN_NOT_ALLOWED.registerFactory(ChangeToLabeledReturnFix)
+        TYPE_MISMATCH.registerFactory(ChangeToLabeledReturnFix)
+        CONSTANT_EXPECTED_TYPE_MISMATCH.registerFactory(ChangeToLabeledReturnFix)
+        NULL_FOR_NONNULL_TYPE.registerFactory(ChangeToLabeledReturnFix)
 
         WRONG_ANNOTATION_TARGET.registerFactory(AddAnnotationTargetFix)
         WRONG_ANNOTATION_TARGET_WITH_USE_SITE_TARGET.registerFactory(MoveReceiverAnnotationFix, AddAnnotationTargetFix)

idea/testData/quickfix/changeToLabeledReturn/constantExpectedTypeMismatch.kt

@@ -0,0 +1,9 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int) {}
+
+fun test() {
+    foo { i ->
+        if (i == 1) return 1<caret>
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/constantExpectedTypeMismatch.kt.after

@@ -0,0 +1,9 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int) {}
+
+fun test() {
+    foo { i ->
+        if (i == 1) return@foo 1
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/nullForNonnullType.kt

@@ -0,0 +1,11 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int?) {}
+
+fun baz(): Int = 0
+
+fun test() {
+    foo { i ->
+        if (i == 1) return null<caret>
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/nullForNonnullType.kt.after

@@ -0,0 +1,11 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int?) {}
+
+fun baz(): Int = 0
+
+fun test() {
+    foo { i ->
+        if (i == 1) return@foo null
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/nullForNonnullType2.kt

@@ -0,0 +1,14 @@
+// "Change to 'return@foo'" "false"
+// ACTION: Add braces to 'if' statement
+// ACTION: Change return type of enclosing function 'test' to 'Unit?'
+// DISABLE-ERRORS
+inline fun foo(f: (Int) -> Int) {}
+
+fun baz(): Int = 0
+
+fun test() {
+    foo { i ->
+        if (i == 1) return null<caret>
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/typeMismatch.kt

@@ -0,0 +1,11 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int) {}
+
+fun baz(): Int = 0
+
+fun test() {
+    foo { i ->
+        if (i == 1) return baz()<caret>
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/typeMismatch.kt.after

@@ -0,0 +1,11 @@
+// "Change to 'return@foo'" "true"
+inline fun foo(f: (Int) -> Int) {}
+
+fun baz(): Int = 0
+
+fun test() {
+    foo { i ->
+        if (i == 1) return@foo baz()
+        0
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/typeMismatch2.kt

@@ -0,0 +1,19 @@
+// "Change to 'return@foo'" "false"
+// ACTION: Add braces to 'if' statement
+// ACTION: Change parameter 'f' type of function 'foo' to '(Int) -> Unit'
+// ACTION: Change return type of called function 'baz' to 'Unit'
+// ACTION: Change return type of enclosing function 'test' to 'String'
+// ACTION: Convert to single-line lambda
+// ACTION: Enable a trailing comma by default in the formatter
+// ACTION: Move lambda argument into parentheses
+// ACTION: Specify explicit lambda signature
+// DISABLE-ERRORS
+inline fun foo(f: (Int) -> Int) {}
+
+fun baz(): String = ""
+
+fun test() {
+    foo { i ->
+        if (i == 1) return baz()<caret>
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -2687,6 +2687,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/quickfix/changeToLabeledReturn"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), null, true);
         }
 
+        @TestMetadata("constantExpectedTypeMismatch.kt")
+        public void testConstantExpectedTypeMismatch() throws Exception {
+            runTest("idea/testData/quickfix/changeToLabeledReturn/constantExpectedTypeMismatch.kt");
+        }
+
         @TestMetadata("multipleInner.kt")
         public void testMultipleInner() throws Exception {
             runTest("idea/testData/quickfix/changeToLabeledReturn/multipleInner.kt");
@@ -2706,6 +2711,26 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
         public void testNormal() throws Exception {
             runTest("idea/testData/quickfix/changeToLabeledReturn/normal.kt");
         }
+
+        @TestMetadata("nullForNonnullType.kt")
+        public void testNullForNonnullType() throws Exception {
+            runTest("idea/testData/quickfix/changeToLabeledReturn/nullForNonnullType.kt");
+        }
+
+        @TestMetadata("nullForNonnullType2.kt")
+        public void testNullForNonnullType2() throws Exception {
+            runTest("idea/testData/quickfix/changeToLabeledReturn/nullForNonnullType2.kt");
+        }
+
+        @TestMetadata("typeMismatch.kt")
+        public void testTypeMismatch() throws Exception {
+            runTest("idea/testData/quickfix/changeToLabeledReturn/typeMismatch.kt");
+        }
+
+        @TestMetadata("typeMismatch2.kt")
+        public void testTypeMismatch2() throws Exception {
+            runTest("idea/testData/quickfix/changeToLabeledReturn/typeMismatch2.kt");
+        }
     }
 
     @TestMetadata("idea/testData/quickfix/changeToMutableCollection")