From 725cb88f41c3e7eb822947c77cc0e8d725d72615 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Fri, 5 Oct 2018 13:21:48 +0900
Subject: [PATCH] Add quick-fix to "Unlabeled return inside lambda" inspection
 #KT-27007 Fixed

---
 .../UnlabeledReturnInsideLambdaInspection.kt    | 17 ++++++++++++++---
 .../labeledReturn.kt                            | 10 ++++++----
 .../notInlineFunction.kt                        | 10 ++++++++++
 .../unlabeledReturnInsideLambda/return.kt       | 11 ++++++-----
 .../unlabeledReturnInsideLambda/return.kt.after |  8 ++++++++
 .../returnInFunction.kt                         | 14 ++++++++------
 .../LocalInspectionTestGenerated.java           |  5 +++++
 7 files changed, 57 insertions(+), 18 deletions(-)
 create mode 100644 idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/notInlineFunction.kt
 create mode 100644 idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/UnlabeledReturnInsideLambdaInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/UnlabeledReturnInsideLambdaInspection.kt
index ae43551ce73..304d30dbd8b 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/UnlabeledReturnInsideLambdaInspection.kt
+++ b/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
+                )
             )
         })
 }
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt
index 6f980912e08..847b18c9a3a 100644
--- a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt
+++ b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/labeledReturn.kt
@@ -1,7 +1,9 @@
 // PROBLEM: none
-// WITH_RUNTIME
-fun test() {
-    listOf(1).forEach {
-        if (it == 10) <caret>return@forEach
+inline fun foo(f: () -> Unit) {}
+
+fun test(): Int {
+    foo {
+        <caret>return@test 0
     }
+    return 1
 }
diff --git a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/notInlineFunction.kt b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/notInlineFunction.kt
new file mode 100644
index 00000000000..5c29680a374
--- /dev/null
+++ b/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
+}
diff --git a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt
index e905b43109b..1ea928edd5a 100644
--- a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt
+++ b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt
@@ -1,7 +1,8 @@
-// FIX: none
-// WITH_RUNTIME
-fun test() {
-    listOf(1).forEach {
-        if (it == 10) <caret>return
+inline fun foo(f: () -> Unit) {}
+
+fun test(): Int {
+    foo {
+        return<caret> 0
     }
+    return 1
 }
diff --git a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt.after b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/return.kt.after
new file mode 100644
index 00000000000..bc002c67c99
--- /dev/null
+++ b/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
+}
diff --git a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/returnInFunction.kt b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/returnInFunction.kt
index a5e3692a59b..226464480a6 100644
--- a/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/returnInFunction.kt
+++ b/idea/testData/inspectionsLocal/unlabeledReturnInsideLambda/returnInFunction.kt
@@ -1,9 +1,11 @@
 // PROBLEM: none
-// WITH_RUNTIME
-fun test() {
-    listOf(1).forEach {
-        fun foo() {
-            if (it == 10) <caret>return
+inline fun foo(f: () -> Unit) {}
+
+fun test(): Int {
+    foo {
+        fun bar() {
+            return<caret>
         }
     }
-}
+    return 1
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
index 300f4949a9d..2ebcd1dbfd5 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
+++ b/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");