Replace with for-each: add new-line if necessary #KT-15858 Fixed

2018-02-16 08:39:07 +03:00
parent ae7f60a96e
commit 7f880bf58c
10 changed files with 102 additions and 3 deletions
@@ -17,8 +17,11 @@
 package org.jetbrains.kotlin.idea.intentions

 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.contentRange
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
@@ -26,7 +29,10 @@ 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") {
+class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpression>(
+    KtForExpression::class.java,
+    "Replace with a 'forEach' function call"
+) {
    override fun isApplicableTo(element: KtForExpression, caretOffset: Int): Boolean {
        val rParen = element.rightParenthesis ?: return false
        if (caretOffset > rParen.endOffset) return false // available only on the loop header, not in the body
@@ -41,11 +47,14 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
        val body = element.body!!
        val loopParameter = element.loopParameter!!

-        val functionBodyArgument: Any = (body as? KtBlockExpression)?.contentRange() ?: body
+        val blockExpression = body as? KtBlockExpression
+        val functionBodyArgument: Any = blockExpression?.contentRange() ?: body
+        val pattern = if (needLineBreaks(blockExpression)) "$0.forEach{$1->\n$2\n}" else "$0.forEach{$1->$2}"

        val psiFactory = KtPsiFactory(element)
        val foreachExpression = psiFactory.createExpressionByPattern(
-                "$0.forEach{$1->$2}", element.loopRange!!, loopParameter, functionBodyArgument)
+            pattern, element.loopRange!!, loopParameter, functionBodyArgument
+        )
        val result = element.replace(foreachExpression) as KtElement

        result.findDescendantOfType<KtFunctionLiteral>()!!.getContinuesWithLabel(labelName).forEach {
@@ -55,6 +64,14 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
        commentSaver.restore(result)
    }

+    private fun needLineBreaks(blockExpression: KtBlockExpression?): Boolean {
+        val contentRange = blockExpression?.contentRange() ?: return false
+        if ((contentRange.last as? PsiComment)?.tokenType != KtTokens.EOL_COMMENT) return false
+        val statements = blockExpression.statements
+        val fistStatementLine = statements.firstOrNull()?.getLineNumber()
+        return statements.all { it.getLineNumber() == fistStatementLine }
+    }
+
    private fun KtElement.getContinuesWithLabel(labelName: String?): List<KtContinueExpression> {
        val continueElements = ArrayList<KtContinueExpression>()

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        /* comment */
+    }
+}
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l -> /* comment */ }
+}
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        // comment
+    }
+}
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    listOf(1, 2).forEach { l ->
+        // comment
+    }
+}
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        foo() // comment
+    }
+}
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l ->
+        foo() // comment
+    }
+}
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        foo(); foo() // comment
+    }
+}
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l ->
+        foo(); foo() // comment
+    }
+}
@@ -6525,6 +6525,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
            doTest(fileName);
        }

+        @TestMetadata("blockCommentOnly.kt")
+        public void testBlockCommentOnly() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("commentsInBody.kt")
        public void testCommentsInBody() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/commentsInBody.kt");
@@ -6561,6 +6567,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
            doTest(fileName);
        }

+        @TestMetadata("endOfLineComment1.kt")
+        public void testEndOfLineComment1() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("endOfLineComment2.kt")
+        public void testEndOfLineComment2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("endOfLineComment3.kt")
+        public void testEndOfLineComment3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("iterativeElementTypeSpecified.kt")
        public void testIterativeElementTypeSpecified() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt");