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

So #KT-16332 Fixed

Author: t-kameyama

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<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 {

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>(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 }
+
 }

idea/testData/intentions/addBraces/doWhileWithComment.kt

@@ -0,0 +1,6 @@
+fun foo() {}
+
+fun test() {
+    <caret>do /* aaa */ foo() // comment
+    while(true)
+}

idea/testData/intentions/addBraces/doWhileWithComment.kt.after

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test() {
+    do {
+        /* aaa */ foo() // comment
+    } while(true)
+}

idea/testData/intentions/addBraces/elseWithComment.kt

@@ -0,0 +1,6 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) foo() <caret>else
+        /* aaa */ foo() // bbb
+}

idea/testData/intentions/addBraces/elseWithComment.kt.after

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) foo() else {
+        /* aaa */ foo() // bbb
+    }
+}

idea/testData/intentions/addBraces/forWithComment.kt

@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test() {
+    <caret>for (i in 1..4) /* aaa */ foo() // bbb
+}

idea/testData/intentions/addBraces/forWithComment.kt.after

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test() {
+    for (i in 1..4) {
+        /* aaa */ foo() // bbb
+    }
+}

idea/testData/intentions/addBraces/ifElseWithComment.kt

@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>if (b) foo() /* if comment */ else foo() // else comment
+}

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
+}

idea/testData/intentions/addBraces/ifWithComment.kt

@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>if (b) /* aaa */ /* bbb */ foo() // ccc
+}

idea/testData/intentions/addBraces/ifWithComment.kt.after

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) {
+        /* aaa */ /* bbb */ foo() // ccc
+    }
+}

idea/testData/intentions/addBraces/whenWithComment.kt

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(a: Int) {
+    when (a) {
+        <caret>1 -> /* aaa */ foo() // bbb
+    }
+}

idea/testData/intentions/addBraces/whenWithComment.kt.after

@@ -0,0 +1,9 @@
+fun foo() {}
+
+fun test(a: Int) {
+    when (a) {
+        1 -> {
+            /* aaa */ foo() // bbb
+        }
+    }
+}

idea/testData/intentions/addBraces/whileWithComment.kt

@@ -0,0 +1,5 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    <caret>while (b) /* aaa */ foo() // bbb
+}

idea/testData/intentions/addBraces/whileWithComment.kt.after

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    while (b) {
+        /* aaa */ foo() // bbb
+    }
+}