ConvertToForEachFunctionCallIntention: convert 'continue' to 'return@forEach'

Fixed #KT-11764

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt

@@ -17,13 +17,14 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.util.CommentSaver
-import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtForExpression
-import org.jetbrains.kotlin.psi.KtPsiFactory
-import org.jetbrains.kotlin.psi.createExpressionByPattern
 import org.jetbrains.kotlin.psi.psiUtil.contentRange
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import java.util.*
 
 class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpression>(KtForExpression::class.java, "Replace with a 'forEach' function call") {
     override fun isApplicableTo(element: KtForExpression, caretOffset: Int): Boolean {
@@ -35,15 +36,58 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
     override fun applyTo(element: KtForExpression, editor: Editor?) {
         val commentSaver = CommentSaver(element)
 
+        val labelName = element.getLabelName()
+
         val body = element.body!!
         val loopParameter = element.loopParameter!!
 
         val functionBodyArgument: Any = if (body is KtBlockExpression) body.contentRange() else body
 
-        val foreachExpression = KtPsiFactory(element).createExpressionByPattern(
+        val psiFactory = KtPsiFactory(element)
+        val foreachExpression = psiFactory.createExpressionByPattern(
                 "$0.forEach{$1->$2}", element.loopRange!!, loopParameter, functionBodyArgument)
-        val result = element.replace(foreachExpression)
+        val result = element.replace(foreachExpression) as KtElement
+
+        result.findDescendantOfType<KtFunctionLiteral>()!!.getContinuesWithLabel(labelName).forEach {
+            it.replace(psiFactory.createExpression("return@forEach"))
+        }
 
         commentSaver.restore(result)
     }
+
+    private fun KtElement.getContinuesWithLabel(labelName: String?): List<KtContinueExpression> {
+        val continueElements = ArrayList<KtContinueExpression>()
+
+        forEachDescendantOfType<KtContinueExpression>({ it.shouldEnterForUnqualified(this) }) {
+            if (it.getLabelName() == null) {
+                continueElements += it
+            }
+        }
+
+        if (labelName != null) {
+            forEachDescendantOfType<KtContinueExpression>({ it.shouldEnterForQualified(this, labelName) }) {
+                if (it.getLabelName() == labelName) {
+                    continueElements += it
+                }
+            }
+        }
+
+        return continueElements
+    }
+
+    private fun PsiElement.shouldEnterForUnqualified(allow: PsiElement): Boolean {
+        if (this == allow) return true
+        if (shouldNeverEnter()) return false
+        return this !is KtLoopExpression
+    }
+
+    private fun PsiElement.shouldEnterForQualified(allow: PsiElement, labelName: String): Boolean {
+        if (this == allow) return true
+        if (shouldNeverEnter()) return false
+        return this !is KtLoopExpression || getLabelName() != labelName
+    }
+
+    private fun PsiElement.shouldNeverEnter() = this is KtLambdaExpression || this is KtClassOrObject || this is KtFunction
+
+    private fun KtLoopExpression.getLabelName() = (parent as? KtExpressionWithLabel)?.getLabelName()
 }

idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+fun main() {
+    outer@
+    <caret>for (i in 1..100) {
+        if (i % 2 == 0) continue
+        inner@
+        for (j in 1..100) {
+            continue@inner
+        }
+        for (j in 1..100) {
+            for (k in 1..1) {
+                continue@outer
+            }
+            continue
+        }
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt.after

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+fun main() {
+    outer@
+    <caret>(1..100).forEach { i ->
+        if (i % 2 == 0) return@forEach
+        inner@
+        for (j in 1..100) {
+            continue@inner
+        }
+        for (j in 1..100) {
+            for (k in 1..1) {
+                return@forEach
+            }
+            continue
+        }
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+fun main() {
+    <caret>for (i in 1..100) {
+        if (i % 2 == 0) continue
+        inner@
+        for (j in 1..100) {
+            continue@inner
+        }
+        for (j in 1..100) {
+            continue
+        }
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt.after

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+fun main() {
+    <caret>(1..100).forEach { i ->
+        if (i % 2 == 0) return@forEach
+        inner@
+        for (j in 1..100) {
+            continue@inner
+        }
+        for (j in 1..100) {
+            continue
+        }
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4576,6 +4576,20 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("continueToReturnWithLabel.kt")
+        public void testContinueToReturnWithLabel() throws Exception {
+            String fileName = KotlinTestUtils
+                    .navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("continueToReturnWithoutLabel.kt")
+        public void testContinueToReturnWithoutLabel() throws Exception {
+            String fileName = KotlinTestUtils
+                    .navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("emptyBody.kt")
         public void testEmptyBody() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt");