diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt index 55a7f5d0b53..40db8fce89c 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt +++ b/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.KtForExpression -import org.jetbrains.kotlin.psi.KtPsiFactory -import org.jetbrains.kotlin.psi.createExpressionByPattern +import org.jetbrains.kotlin.psi.* 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::class.java, "Replace with a 'forEach' function call") { override fun isApplicableTo(element: KtForExpression, caretOffset: Int): Boolean { @@ -35,15 +36,58 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention$2}", element.loopRange!!, loopParameter, functionBodyArgument) - val result = element.replace(foreachExpression) + val result = element.replace(foreachExpression) as KtElement + + result.findDescendantOfType()!!.getContinuesWithLabel(labelName).forEach { + it.replace(psiFactory.createExpression("return@forEach")) + } commentSaver.restore(result) } + + private fun KtElement.getContinuesWithLabel(labelName: String?): List { + val continueElements = ArrayList() + + forEachDescendantOfType({ it.shouldEnterForUnqualified(this) }) { + if (it.getLabelName() == null) { + continueElements += it + } + } + + if (labelName != null) { + forEachDescendantOfType({ 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() } diff --git a/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt new file mode 100644 index 00000000000..8322645393f --- /dev/null +++ b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt @@ -0,0 +1,17 @@ +// WITH_RUNTIME +fun main() { + outer@ + 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 + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt.after new file mode 100644 index 00000000000..a20a0e5d6dd --- /dev/null +++ b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithLabel.kt.after @@ -0,0 +1,17 @@ +// WITH_RUNTIME +fun main() { + outer@ + (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 + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt new file mode 100644 index 00000000000..d294c4709a2 --- /dev/null +++ b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt @@ -0,0 +1,13 @@ +// WITH_RUNTIME +fun main() { + for (i in 1..100) { + if (i % 2 == 0) continue + inner@ + for (j in 1..100) { + continue@inner + } + for (j in 1..100) { + continue + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt.after new file mode 100644 index 00000000000..ad02e42ab34 --- /dev/null +++ b/idea/testData/intentions/convertToForEachFunctionCall/continueToReturnWithoutLabel.kt.after @@ -0,0 +1,13 @@ +// WITH_RUNTIME +fun main() { + (1..100).forEach { i -> + if (i % 2 == 0) return@forEach + inner@ + for (j in 1..100) { + continue@inner + } + for (j in 1..100) { + continue + } + } +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java index 090184c1a0b..122aa501d45 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java +++ b/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");