Introduce quick-fix adding label to 'return' in closures #KT-16851 Fixed

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.quickfix
+
+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.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.parentsWithSelf
+import org.jetbrains.kotlin.renderer.render
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
+
+class ChangeToLabeledReturnFix(
+        element: KtReturnExpression, val labeledReturn: String
+) : KotlinQuickFixAction<KtReturnExpression>(element) {
+
+    override fun getFamilyName() = "Change to return with label"
+    override fun getText() = "Change to '$labeledReturn'"
+
+    override fun invoke(project: Project, editor: Editor?, file: KtFile) {
+        val returnExpression = element ?: return
+        val factory = KtPsiFactory(project)
+        val returnedExpression = returnExpression.returnedExpression
+        val newExpression = if (returnedExpression == null)
+            factory.createExpression(labeledReturn)
+        else
+            factory.createExpressionByPattern("$0 $1", labeledReturn, returnedExpression)
+        returnExpression.replace(newExpression)
+    }
+
+    companion object : KotlinIntentionActionsFactory() {
+        private fun findAccessibleLabels(bindingContext: BindingContext, position: KtReturnExpression): List<Name> {
+            val result = mutableListOf<Name>()
+            for (parent in position.parentsWithSelf) {
+                if (parent is KtFunctionLiteral) {
+                    val (label, call) = parent.findLabelAndCall()
+                    if (label != null) {
+                        result.add(label)
+                    }
+
+                    // check if the current function literal is inlined and stop processing outer declarations if it's not
+                    val callee = call?.calleeExpression as? KtReferenceExpression ?: break
+                    if (!InlineUtil.isInline(bindingContext[BindingContext.REFERENCE_TARGET, callee])) break
+                }
+            }
+            return result
+        }
+
+        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()}")
+            }
+        }
+    }
+}

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

@@ -491,5 +491,7 @@ class QuickFixRegistrar : QuickFixContributor {
         CONFLICTING_OVERLOADS.registerFactory(ChangeSuspendInHierarchyFix)
 
         MUST_BE_INITIALIZED_OR_BE_ABSTRACT.registerFactory(AddModifierFix.AddLateinitFactory)
+
+        RETURN_NOT_ALLOWED.registerFactory(ChangeToLabeledReturnFix)
     }
 }

idea/testData/quickfix/changeToLabeledReturn/multipleInner.kt

@@ -0,0 +1,17 @@
+// "Change to 'return@forEach'" "true"
+// ACTION: Change to 'return@foo'
+// ACTION: Introduce local variable
+// ERROR: The integer literal does not conform to the expected type Unit
+// WITH_RUNTIME
+
+fun foo(f:()->Int){}
+
+fun bar() {
+
+    foo {
+        listOf(1).forEach {
+            return<caret> 1
+        }
+        return@foo 1
+    }
+}

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

@@ -0,0 +1,17 @@
+// "Change to 'return@forEach'" "true"
+// ACTION: Change to 'return@foo'
+// ACTION: Introduce local variable
+// ERROR: The integer literal does not conform to the expected type Unit
+// WITH_RUNTIME
+
+fun foo(f:()->Int){}
+
+fun bar() {
+
+    foo {
+        listOf(1).forEach {
+            return@forEach 1
+        }
+        return@foo 1
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/multipleOuter.kt

@@ -0,0 +1,16 @@
+// "Change to 'return@foo'" "true"
+// ACTION: Change to 'return@forEach'
+// ACTION: Introduce local variable
+// WITH_RUNTIME
+
+fun foo(f:()->Int){}
+
+fun bar() {
+
+    foo {
+        listOf(1).forEach {
+            return<caret> 1
+        }
+        return@foo 1
+    }
+}

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

@@ -0,0 +1,16 @@
+// "Change to 'return@foo'" "true"
+// ACTION: Change to 'return@forEach'
+// ACTION: Introduce local variable
+// WITH_RUNTIME
+
+fun foo(f:()->Int){}
+
+fun bar() {
+
+    foo {
+        listOf(1).forEach {
+            return@foo 1
+        }
+        return@foo 1
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/noCandidates.kt

@@ -0,0 +1,10 @@
+// "Change to 'return@init'" "false"
+// ACTION: Introduce local variable
+// ERROR: 'return' is not allowed here
+// WITH_RUNTIME
+
+class Foo {
+    init {
+        return<caret> 1
+    }
+}

idea/testData/quickfix/changeToLabeledReturn/normal.kt

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

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

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

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

@@ -1154,6 +1154,39 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
 
     }
 
+    @TestMetadata("idea/testData/quickfix/changeToLabeledReturn")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ChangeToLabeledReturn extends AbstractQuickFixTest {
+        public void testAllFilesPresentInChangeToLabeledReturn() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/changeToLabeledReturn"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
+        }
+
+        @TestMetadata("multipleInner.kt")
+        public void testMultipleInner() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/changeToLabeledReturn/multipleInner.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("multipleOuter.kt")
+        public void testMultipleOuter() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/changeToLabeledReturn/multipleOuter.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("noCandidates.kt")
+        public void testNoCandidates() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/changeToLabeledReturn/noCandidates.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("normal.kt")
+        public void testNormal() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/changeToLabeledReturn/normal.kt");
+            doTest(fileName);
+        }
+    }
+
     @TestMetadata("idea/testData/quickfix/changeToUseSpreadOperator")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)