Add braces to 'if' statement: save/restore comments correctly

So #KT-16332 Fixed
2018-02-15 14:50:10 +03:00
parent b436ee6e19
commit 4726b44371
17 changed files with 178 additions and 5 deletions
@@ -815,8 +815,10 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
        return createFunction("fun foo() {\n$bodyText\n}").bodyExpression as KtBlockExpression
    }

-    fun createSingleStatementBlock(statement: KtExpression): KtBlockExpression {
-        return createDeclarationByPattern<KtNamedFunction>("fun foo() {\n$0\n}", statement).bodyExpression as KtBlockExpression
+    fun createSingleStatementBlock(statement: KtExpression, prevComment: String? = null, nextComment: String? = null): KtBlockExpression {
+        val prev = if (prevComment == null) "" else " $prevComment "
+        val next = if (nextComment == null) "" else " $nextComment "
+        return createDeclarationByPattern<KtNamedFunction>("fun foo() {\n$prev$0$next\n}", statement).bodyExpression as KtBlockExpression
    }

    fun createComment(text: String): PsiComment {
@@ -17,8 +17,15 @@
 package org.jetbrains.kotlin.idea.intentions

 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.idea.refactoring.getLineNumber
+import org.jetbrains.kotlin.j2k.isInSingleLine
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.nextLeafs
+import org.jetbrains.kotlin.psi.psiUtil.prevLeafs
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import java.lang.IllegalArgumentException

@@ -52,8 +59,10 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
            element.nextSibling!!.delete()
        }

+        val (prevComment, nextComment) = deleteCommentsOnSameLine(expression, element)
+
        val psiFactory = KtPsiFactory(element)
-        expression.replace(psiFactory.createSingleStatementBlock(expression))
+        expression.replace(psiFactory.createSingleStatementBlock(expression, prevComment, nextComment))

        if (element is KtDoWhileExpression) { // remove new line between '}' and while
            (element.body!!.parent.nextSibling as? PsiWhiteSpace)?.delete()
@@ -67,8 +76,7 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
                val elseExpr = `else`
                if (elseExpr != null && caretLocation >= elseKeyword!!.startOffset) {
                    elseExpr
-                }
-                else {
+                } else {
                    thenExpr
                }
            }
@@ -78,4 +86,35 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
            else -> null
        }
    }
+
+    private fun deleteCommentsOnSameLine(expression: KtExpression, element: KtElement): Pair<String?, String?> {
+        val lineNumber = expression.getLineNumber()
+
+        val prevComments = getCommentsOnSameLine(lineNumber, expression.prevLeafs).reversed()
+        val prevCommentText = createCommentText(prevComments)
+
+        val nextLeafs = when {
+            expression.parent.node.elementType == KtNodeTypes.THEN && (element as? KtIfExpression)?.`else` != null -> expression.nextLeafs
+            element is KtDoWhileExpression -> expression.nextLeafs
+            else -> expression.nextLeafs
+        }
+        val nextComments = getCommentsOnSameLine(lineNumber, nextLeafs)
+        val nextCommentText = createCommentText(nextComments)
+
+        (prevComments + nextComments).forEach { (it as? PsiComment)?.delete() }
+
+        return prevCommentText to nextCommentText
+    }
+
+    private fun getCommentsOnSameLine(lineNumber: Int, elements: Sequence<PsiElement>): List<PsiElement> {
+        return elements
+            .takeWhile { (it is PsiWhiteSpace || it is PsiComment) && lineNumber == it.getLineNumber() && it.isInSingleLine() }
+            .dropWhile { it is PsiWhiteSpace }
+            .toList()
+            .dropLastWhile { it is PsiWhiteSpace }
+    }
+
+    private fun createCommentText(comments: List<PsiElement>): String? =
+        if (comments.isEmpty()) null else comments.joinToString("") { it.text }
+
 }
@@ -0,0 +1,6 @@
+fun foo() {}
+
+fun test() {
+    <caret>do /* aaa */ foo() // comment
+    while(true)
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test() {
+    do {
+        /* aaa */ foo() // comment
+    } while(true)
+}
@@ -0,0 +1,6 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) foo() <caret>else
+        /* aaa */ foo() // bbb
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) foo() else {
+        /* aaa */ foo() // bbb
+    }
+}
@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test() {
+    <caret>for (i in 1..4) /* aaa */ foo() // bbb
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test() {
+    for (i in 1..4) {
+        /* aaa */ foo() // bbb
+    }
+}
@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>if (b) foo() /* if comment */ else foo() // else comment
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) {
+        foo() /* if comment */
+    } else foo() // else comment
+}
@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>if (b) /* aaa */ /* bbb */ foo() // ccc
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) {
+        /* aaa */ /* bbb */ foo() // ccc
+    }
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(a: Int) {
+    when (a) {
+        <caret>1 -> /* aaa */ foo() // bbb
+    }
+}
@@ -0,0 +1,9 @@
+fun foo() {}
+
+fun test(a: Int) {
+    when (a) {
+        1 -> {
+            /* aaa */ foo() // bbb
+        }
+    }
+}
@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>while (b) /* aaa */ foo() // bbb
+}
@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    while (b) {
+        /* aaa */ foo() // bbb
+    }
+}
@@ -81,6 +81,36 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/addBraces"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
        }

+        @TestMetadata("doWhileWithComment.kt")
+        public void testDoWhileWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/doWhileWithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("elseWithComment.kt")
+        public void testElseWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/elseWithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forWithComment.kt")
+        public void testForWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/forWithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifElseWithComment.kt")
+        public void testIfElseWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/ifElseWithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifWithComment.kt")
+        public void testIfWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/ifWithComment.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("notInsideElseIfBlock.kt")
        public void testNotInsideElseIfBlock() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/notInsideElseIfBlock.kt");
@@ -98,6 +128,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/whenSimple.kt");
            doTest(fileName);
        }
+
+        @TestMetadata("whenWithComment.kt")
+        public void testWhenWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/whenWithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("whileWithComment.kt")
+        public void testWhileWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/whileWithComment.kt");
+            doTest(fileName);
+        }
    }

    @TestMetadata("idea/testData/intentions/addConstModifier")