Do not add a new label to labeled loop

Relevant situation: break / continue in when #KT-16128 Fixed

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

@@ -25,8 +25,15 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.parents
 
-class AddLoopLabelFix(loop: KtLoopExpression, val jumpExpression: KtElement): KotlinQuickFixAction<KtLoopExpression>(loop) {
-    override fun getText() = "Add label to loop"
+class AddLoopLabelFix(loop: KtLoopExpression, val jumpExpression: KtExpressionWithLabel): KotlinQuickFixAction<KtLoopExpression>(loop) {
+
+    private val existingLabelName = (loop.parent as? KtLabeledExpression)?.getLabelName()
+
+    private val description =
+            if (existingLabelName != null) "Add '@$existingLabelName' to ${jumpExpression.text}"
+            else "Add label to loop"
+
+    override fun getText() = description
     override fun getFamilyName() = text
 
     override fun isAvailable(project: Project, editor: Editor?, file: PsiFile): Boolean {
@@ -35,16 +42,17 @@ class AddLoopLabelFix(loop: KtLoopExpression, val jumpExpression: KtElement): Ko
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
-        val usedLabels = collectUsedLabels(element)
-        val labelName = getUniqueLabelName(usedLabels)
+        val labelName = existingLabelName ?: getUniqueLabelName(collectUsedLabels(element))
 
         val jumpWithLabel = KtPsiFactory(project).createExpression(jumpExpression.text + "@" + labelName)
         jumpExpression.replace(jumpWithLabel)
 
         // TODO(yole) use createExpressionByPattern() once it's available
-        val labeledLoopExpression = KtPsiFactory(project).createLabeledExpression(labelName)
-        labeledLoopExpression.baseExpression!!.replace(element)
-        element.replace(labeledLoopExpression)
+        if (existingLabelName == null) {
+            val labeledLoopExpression = KtPsiFactory(project).createLabeledExpression(labelName)
+            labeledLoopExpression.baseExpression!!.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
     }
@@ -76,7 +84,7 @@ class AddLoopLabelFix(loop: KtLoopExpression, val jumpExpression: KtElement): Ko
 
     companion object: KotlinSingleIntentionActionFactory() {
         override fun createAction(diagnostic: Diagnostic): IntentionAction? {
-            val element = diagnostic.psiElement as? KtElement
+            val element = diagnostic.psiElement as? KtExpressionWithLabel
             assert(element is KtBreakExpression || element is KtContinueExpression)
             assert((element as? KtLabeledExpression)?.getLabelName() == null)
             val loop = element?.getStrictParentOfType<KtLoopExpression>() ?: return null

idea/testData/quickfix/when/breakInWhenInLabeled.kt

@@ -0,0 +1,14 @@
+// "Add '@loop' to break" "true"
+// WITH_RUNTIME
+
+fun foo(chars: CharArray) {
+    val length = chars.size
+    var pos = 0
+    loop@ while (pos < length) {
+        val c = chars[pos]
+        when (c) {
+            '\n' -> br<caret>eak
+        }
+        pos++
+    }
+}

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

@@ -0,0 +1,14 @@
+// "Add '@loop' to break" "true"
+// WITH_RUNTIME
+
+fun foo(chars: CharArray) {
+    val length = chars.size
+    var pos = 0
+    loop@ while (pos < length) {
+        val c = chars[pos]
+        when (c) {
+            '\n' -> break@loop
+        }
+        pos++
+    }
+}

idea/testData/quickfix/when/continueInWhenInLabeled.kt

@@ -0,0 +1,14 @@
+// "Add '@loop' to continue" "true"
+// WITH_RUNTIME
+
+fun foo(chars: CharArray) {
+    val length = chars.size
+    var pos = 0
+    loop@ while (pos < length) {
+        val c = chars[pos]
+        when (c) {
+            '\n' -> continue<caret>
+        }
+        pos++
+    }
+}

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

@@ -0,0 +1,14 @@
+// "Add '@loop' to continue" "true"
+// WITH_RUNTIME
+
+fun foo(chars: CharArray) {
+    val length = chars.size
+    var pos = 0
+    loop@ while (pos < length) {
+        val c = chars[pos]
+        when (c) {
+            '\n' -> continue@loop
+        }
+        pos++
+    }
+}

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

@@ -10640,6 +10640,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("breakInWhenInLabeled.kt")
+        public void testBreakInWhenInLabeled() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/breakInWhenInLabeled.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("commasInConditionWithNoArguments.kt")
         public void testCommasInConditionWithNoArguments() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/commasInConditionWithNoArguments.kt");
@@ -10658,6 +10664,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("continueInWhenInLabeled.kt")
+        public void testContinueInWhenInLabeled() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/continueInWhenInLabeled.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("continueInWhenWithLabel.kt")
         public void testContinueInWhenWithLabel() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/continueInWhenWithLabel.kt");