diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt index 67fe21e8527..d7ca813be57 100644 --- a/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt +++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt @@ -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("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("fun foo() {\n$prev$0$next\n}", statement).bodyExpression as KtBlockExpression } fun createComment(text: String): PsiComment { diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt index 9d99829341b..f56b392aab4 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt @@ -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::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::class.ja val elseExpr = `else` if (elseExpr != null && caretLocation >= elseKeyword!!.startOffset) { elseExpr - } - else { + } else { thenExpr } } @@ -78,4 +86,35 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja else -> null } } + + private fun deleteCommentsOnSameLine(expression: KtExpression, element: KtElement): Pair { + 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): List { + 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): String? = + if (comments.isEmpty()) null else comments.joinToString("") { it.text } + } diff --git a/idea/testData/intentions/addBraces/doWhileWithComment.kt b/idea/testData/intentions/addBraces/doWhileWithComment.kt new file mode 100644 index 00000000000..1a0538e8f14 --- /dev/null +++ b/idea/testData/intentions/addBraces/doWhileWithComment.kt @@ -0,0 +1,6 @@ +fun foo() {} + +fun test() { + do /* aaa */ foo() // comment + while(true) +} diff --git a/idea/testData/intentions/addBraces/doWhileWithComment.kt.after b/idea/testData/intentions/addBraces/doWhileWithComment.kt.after new file mode 100644 index 00000000000..8efe570e251 --- /dev/null +++ b/idea/testData/intentions/addBraces/doWhileWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test() { + do { + /* aaa */ foo() // comment + } while(true) +} diff --git a/idea/testData/intentions/addBraces/elseWithComment.kt b/idea/testData/intentions/addBraces/elseWithComment.kt new file mode 100644 index 00000000000..aadad142e1a --- /dev/null +++ b/idea/testData/intentions/addBraces/elseWithComment.kt @@ -0,0 +1,6 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) foo() else + /* aaa */ foo() // bbb +} diff --git a/idea/testData/intentions/addBraces/elseWithComment.kt.after b/idea/testData/intentions/addBraces/elseWithComment.kt.after new file mode 100644 index 00000000000..e3460325312 --- /dev/null +++ b/idea/testData/intentions/addBraces/elseWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) foo() else { + /* aaa */ foo() // bbb + } +} diff --git a/idea/testData/intentions/addBraces/forWithComment.kt b/idea/testData/intentions/addBraces/forWithComment.kt new file mode 100644 index 00000000000..5add69eb88e --- /dev/null +++ b/idea/testData/intentions/addBraces/forWithComment.kt @@ -0,0 +1,5 @@ +fun foo() {} + +fun test() { + for (i in 1..4) /* aaa */ foo() // bbb +} diff --git a/idea/testData/intentions/addBraces/forWithComment.kt.after b/idea/testData/intentions/addBraces/forWithComment.kt.after new file mode 100644 index 00000000000..5e9de9b37da --- /dev/null +++ b/idea/testData/intentions/addBraces/forWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test() { + for (i in 1..4) { + /* aaa */ foo() // bbb + } +} diff --git a/idea/testData/intentions/addBraces/ifElseWithComment.kt b/idea/testData/intentions/addBraces/ifElseWithComment.kt new file mode 100644 index 00000000000..831632482d5 --- /dev/null +++ b/idea/testData/intentions/addBraces/ifElseWithComment.kt @@ -0,0 +1,5 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) foo() /* if comment */ else foo() // else comment +} diff --git a/idea/testData/intentions/addBraces/ifElseWithComment.kt.after b/idea/testData/intentions/addBraces/ifElseWithComment.kt.after new file mode 100644 index 00000000000..3b7cc5f2b80 --- /dev/null +++ b/idea/testData/intentions/addBraces/ifElseWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) { + foo() /* if comment */ + } else foo() // else comment +} diff --git a/idea/testData/intentions/addBraces/ifWithComment.kt b/idea/testData/intentions/addBraces/ifWithComment.kt new file mode 100644 index 00000000000..8128bbdb114 --- /dev/null +++ b/idea/testData/intentions/addBraces/ifWithComment.kt @@ -0,0 +1,5 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) /* aaa */ /* bbb */ foo() // ccc +} diff --git a/idea/testData/intentions/addBraces/ifWithComment.kt.after b/idea/testData/intentions/addBraces/ifWithComment.kt.after new file mode 100644 index 00000000000..72f3449f686 --- /dev/null +++ b/idea/testData/intentions/addBraces/ifWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) { + /* aaa */ /* bbb */ foo() // ccc + } +} diff --git a/idea/testData/intentions/addBraces/whenWithComment.kt b/idea/testData/intentions/addBraces/whenWithComment.kt new file mode 100644 index 00000000000..5b79c6cff19 --- /dev/null +++ b/idea/testData/intentions/addBraces/whenWithComment.kt @@ -0,0 +1,7 @@ +fun foo() {} + +fun test(a: Int) { + when (a) { + 1 -> /* aaa */ foo() // bbb + } +} diff --git a/idea/testData/intentions/addBraces/whenWithComment.kt.after b/idea/testData/intentions/addBraces/whenWithComment.kt.after new file mode 100644 index 00000000000..718c2b5a8db --- /dev/null +++ b/idea/testData/intentions/addBraces/whenWithComment.kt.after @@ -0,0 +1,9 @@ +fun foo() {} + +fun test(a: Int) { + when (a) { + 1 -> { + /* aaa */ foo() // bbb + } + } +} diff --git a/idea/testData/intentions/addBraces/whileWithComment.kt b/idea/testData/intentions/addBraces/whileWithComment.kt new file mode 100644 index 00000000000..182268a7ad2 --- /dev/null +++ b/idea/testData/intentions/addBraces/whileWithComment.kt @@ -0,0 +1,5 @@ +fun foo() {} + +fun test(b: Boolean) { + while (b) /* aaa */ foo() // bbb +} diff --git a/idea/testData/intentions/addBraces/whileWithComment.kt.after b/idea/testData/intentions/addBraces/whileWithComment.kt.after new file mode 100644 index 00000000000..c1431b17c39 --- /dev/null +++ b/idea/testData/intentions/addBraces/whileWithComment.kt.after @@ -0,0 +1,7 @@ +fun foo() {} + +fun test(b: Boolean) { + while (b) { + /* aaa */ foo() // bbb + } +} diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java index d9d5cd73144..1bded1d5ac6 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java @@ -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")