Add quick-fix to "Unlabeled return inside lambda" inspection #KT-27007 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/UnlabeledReturnInsideLambdaInspection.kt

@@ -5,24 +5,35 @@
 
 package org.jetbrains.kotlin.idea.inspections
 
+import com.intellij.codeInspection.IntentionWrapper
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.psi.PsiElementVisitor
+import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.quickfix.ChangeToLabeledReturnFix
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.returnExpressionVisitor
 
 class UnlabeledReturnInsideLambdaInspection : AbstractKotlinInspection() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor =
         returnExpressionVisitor(fun(returnExpression: KtReturnExpression) {
             if (returnExpression.labelQualifier != null) return
-            if (returnExpression.getParentOfType<KtLambdaExpression>(true, KtNamedFunction::class.java) == null) return
+            val lambda = returnExpression.getParentOfType<KtLambdaExpression>(true, KtNamedFunction::class.java) ?: return
+            val parentFunction = lambda.getStrictParentOfType<KtNamedFunction>() ?: return
+            if (returnExpression.analyze().diagnostics.forElement(returnExpression).any { it.severity == Severity.ERROR }) return
             holder.registerProblem(
-                returnExpression,
+                returnExpression.returnKeyword,
                 "Unlabeled return inside lambda",
-                ProblemHighlightType.GENERIC_ERROR_OR_WARNING
+                ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
+                IntentionWrapper(
+                    ChangeToLabeledReturnFix(returnExpression, labeledReturn = "return@${parentFunction.name}"),
+                    returnExpression.containingFile
+                )
             )
         })
 }

idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt

@@ -1,7 +1,9 @@
 // PROBLEM: none
-// WITH_RUNTIME
+inline fun foo(f: () -> Unit) {}
-fun test() {
+
-    listOf(1).forEach {
+fun test(): Int {
-        if (it == 10) <caret>return@forEach
+    foo {
+        <caret>return@test 0
     }
+    return 1
 }

idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/notInlineFunction.kt

@@ -0,0 +1,10 @@
+// PROBLEM: none
+// DISABLE-ERRORS
+fun foo(f: () -> Unit) {}
+
+fun test(): Int {
+    foo {
+        return<caret> 0
+    }
+    return 1
+}

idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt

@@ -1,7 +1,8 @@
-// FIX: none
+inline fun foo(f: () -> Unit) {}
-// WITH_RUNTIME
+
-fun test() {
+fun test(): Int {
-    listOf(1).forEach {
+    foo {
-        if (it == 10) <caret>return
+        return<caret> 0
     }
+    return 1
 }

idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt.after

@@ -0,0 +1,8 @@
+inline fun foo(f: () -> Unit) {}
+
+fun test(): Int {
+    foo {
+        return@test 0
+    }
+    return 1
+}

idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/returnInFunction.kt

@@ -1,9 +1,11 @@
 // PROBLEM: none
-// WITH_RUNTIME
+inline fun foo(f: () -> Unit) {}
-fun test() {
+
-    listOf(1).forEach {
+fun test(): Int {
-        fun foo() {
+    foo {
-            if (it == 10) <caret>return
+        fun bar() {
+            return<caret>
         }
     }
-}
+    return 1
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -6041,6 +6041,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             runTest("idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt");
         }
 
+        @TestMetadata("notInlineFunction.kt")
+        public void testNotInlineFunction() throws Exception {
+            runTest("idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/notInlineFunction.kt");
+        }
+
         @TestMetadata("return.kt")
         public void testReturn() throws Exception {
             runTest("idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt");