Quickfix to add a loop label when 'break' or 'continue' is used in a loop inside 'when'

#KT-7202 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt

@@ -321,6 +321,9 @@ public class JetPsiFactory(private val project: Project) {
         return (createExpression("this@" + labelName) as JetThisExpression).getTargetLabel()!!
     }
 
+    public fun createLabeledExpression(labelName: String): JetLabeledExpression
+        = createExpression("$labelName@ 1") as JetLabeledExpression
+
     public fun createFieldIdentifier(fieldName: String): JetExpression {
         return createExpression("$" + fieldName)
     }

compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java

@@ -562,6 +562,8 @@ public class JetPsiUtil {
             return true;
         }
 
+        if (parentExpression instanceof JetLabeledExpression) return false;
+
         // 'x ?: ...' case
         if (parentExpression instanceof JetBinaryExpression && parentOperation == JetTokens.ELVIS && currentInner == ((JetBinaryExpression) parentExpression).getRight()) {
             return false;

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2015 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 com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.parents
+import java.util.*
+
+public class AddLoopLabelFix(loop: JetLoopExpression, val jumpExpression: JetElement): JetIntentionAction<JetLoopExpression>(loop) {
+    override fun getText() = "Add label to loop"
+    override fun getFamilyName() = getText()
+
+    override fun isAvailable(project: Project, editor: Editor?, file: PsiFile?): Boolean {
+        return super.isAvailable(project, editor, file)
+    }
+
+    override fun invoke(project: Project, editor: Editor, file: JetFile) {
+        val usedLabels = collectUsedLabels(element)
+        val labelName = getUniqueLabelName(usedLabels)
+
+        val jumpWithLabel = JetPsiFactory(project).createExpression(jumpExpression.getText() + "@" + labelName)
+        jumpExpression.replace(jumpWithLabel)
+
+        // TODO(yole) use createExpressionByPattern() once it's available
+        val labeledLoopExpression = JetPsiFactory(project).createLabeledExpression(labelName)
+        labeledLoopExpression.getBaseExpression().replace(element)
+        element.replace(labeledLoopExpression)
+
+        // TODO(yole) We should initiate in-place rename for the label here, but in-place rename for labels is not yet implemented
+    }
+
+    private fun collectUsedLabels(element: JetElement): Set<String> {
+        val usedLabels = hashSetOf<String>()
+        element.acceptChildren(object : JetTreeVisitorVoid() {
+            override fun visitLabeledExpression(expression: JetLabeledExpression) {
+                super.visitLabeledExpression(expression)
+                usedLabels.add(expression.getLabelName())
+            }
+        })
+        element.parents(false).forEach {
+            if (it is JetLabeledExpression) {
+                usedLabels.add(it.getLabelName())
+            }
+        }
+        return usedLabels
+    }
+
+    private fun getUniqueLabelName(existingNames: Collection<String>): String {
+        var index = 0
+        var result = "loop"
+        while (result in existingNames) {
+            result = "loop${++index}"
+        }
+        return result
+    }
+
+    companion object: JetSingleIntentionActionFactory() {
+        override fun createAction(diagnostic: Diagnostic): IntentionAction? {
+            val element = diagnostic.getPsiElement() as? JetElement
+            assert(element is JetBreakExpression || element is JetContinueExpression)
+            assert((element as? JetLabeledExpression)?.getLabelName() == null)
+            val loop = element?.getStrictParentOfType<JetLoopExpression>() ?: return null
+            return AddLoopLabelFix(loop, element!!)
+        }
+    }
+}

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

@@ -185,6 +185,7 @@ public class QuickFixRegistrar {
 
         QuickFixes.factories.put(ELSE_MISPLACED_IN_WHEN, MoveWhenElseBranchFix.createFactory());
         QuickFixes.factories.put(NO_ELSE_IN_WHEN, AddWhenElseBranchFix.createFactory());
+        QuickFixes.factories.put(BREAK_OR_CONTINUE_IN_WHEN, AddLoopLabelFix.Companion);
 
         QuickFixes.factories.put(NO_TYPE_ARGUMENTS_ON_RHS, AddStarProjectionsFix.createFactoryForIsExpression());
         QuickFixes.factories.put(WRONG_NUMBER_OF_TYPE_ARGUMENTS, AddStarProjectionsFix.createFactoryForJavaClass());

idea/testData/quickfix/when/breakInWhen.kt

@@ -0,0 +1,16 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    for (y in 0..10) {
+        when(i) {
+            0 -> bre<caret>ak
+            2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+        }
+    }
+}

idea/testData/quickfix/when/breakInWhen.kt.after

@@ -0,0 +1,16 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    loop@ for (y in 0..10) {
+        when(i) {
+            0 -> break@loop
+            2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+        }
+    }
+}

idea/testData/quickfix/when/continueInWhen.kt

@@ -0,0 +1,16 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    for (y in 0..10) {
+        when(i) {
+            0 -> co<caret>ntinue
+            2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+        }
+    }
+}

idea/testData/quickfix/when/continueInWhen.kt.after

@@ -0,0 +1,16 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    loop@ for (y in 0..10) {
+        when(i) {
+            0 -> continue@loop
+            2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+        }
+    }
+}

idea/testData/quickfix/when/continueInWhenWithLabel.kt

@@ -0,0 +1,18 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    loop@ for (x in 0..10) {
+        for (y in 0..10) {
+            when(i) {
+                0 -> co<caret>ntinue
+                2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+            }
+        }
+    }
+}

idea/testData/quickfix/when/continueInWhenWithLabel.kt.after

@@ -0,0 +1,18 @@
+// "Add label to loop" "true"
+fun breakContinueInWhen(i: Int) {
+    loop@ for (x in 0..10) {
+        loop1@ for (y in 0..10) {
+            when(i) {
+                0 -> co<caret>ntinue@loop1
+                2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+            }
+        }
+    }
+}

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

@@ -5418,6 +5418,24 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/when"), Pattern.compile("^(\\w+)\\.kt$"), true);
         }
 
+        @TestMetadata("breakInWhen.kt")
+        public void testBreakInWhen() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/when/breakInWhen.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("continueInWhen.kt")
+        public void testContinueInWhen() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/when/continueInWhen.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("continueInWhenWithLabel.kt")
+        public void testContinueInWhenWithLabel() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/when/continueInWhenWithLabel.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("elseNotLastInWhen.kt")
         public void testElseNotLastInWhen() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/when/elseNotLastInWhen.kt");