Correctly handle end-of-line comment in "Add braces to if"

#KT-27408 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt

@@ -123,9 +123,9 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
     private class LineBreakTreeElement() : TreeElement()
 
     private class CommentTreeElement(
-            val commentText: String,
-            val spaceBefore: String,
-            val spaceAfter: String
+        val commentText: String,
+        val spaceBefore: String,
+        val spaceAfter: String
     ) : TreeElement() {
         companion object {
             fun create(comment: PsiComment): CommentTreeElement {
@@ -234,11 +234,15 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
         putNewElementIntoMap(newPsiElement, treeElement)
     }
 
+    fun restore(resultElement: PsiElement, isCommentBeneathSingleLine: Boolean, forceAdjustIndent: Boolean = false) {
+        restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent, isCommentBeneathSingleLine)
+    }
+
     fun restore(resultElement: PsiElement, forceAdjustIndent: Boolean = false) {
         restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent)
     }
 
-    fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false) {
+    fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false, isCommentBeneathSingleLine: Boolean = false) {
         assert(!isFinished)
         assert(!resultElements.isEmpty)
 
@@ -260,7 +264,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                     })
                 }
 
-                restoreComments(resultElements)
+                restoreComments(resultElements, isCommentBeneathSingleLine)
 
                 restoreLineBreaks()
 
@@ -291,7 +295,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
         isFinished = true
     }
 
-    private fun restoreComments(resultElements: PsiChildRange) {
+    private fun restoreComments(resultElements: PsiChildRange, isCommentBeneathSingleLine: Boolean = false) {
         var putAbandonedCommentsAfter = resultElements.last!!
 
         for (commentTreeElement in commentsToRestore) {
@@ -316,25 +320,30 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                         val whiteSpace = restored.nextLeaf(skipEmptyElements = true) as? PsiWhiteSpace
                         if (whiteSpace == null) {
                             parent.addAfter(psiFactory.createWhiteSpace("\n"), restored)
-                        }
-                        else if (!whiteSpace.textContains('\n')) {
+                        } else if (!whiteSpace.textContains('\n')) {
                             val newWhiteSpace = psiFactory.createWhiteSpace("\n" + whiteSpace.text)
                             whiteSpace.replace(newWhiteSpace)
                         }
                     }
-                }
-                else {
+                } else {
                     restored = parent.addBefore(comment, anchorElement) as PsiComment
                     if (commentTreeElement.spaceAfter.isNotEmpty()) {
                         parent.addBefore(psiFactory.createWhiteSpace(commentTreeElement.spaceAfter), anchorElement)
                     }
                 }
-            }
-            else {
+            } else {
                 restored = putAbandonedCommentsAfter.parent.addBefore(comment, putAbandonedCommentsAfter) as PsiComment
                 putAbandonedCommentsAfter = restored
             }
 
+            // shift (possible contained) comment in expression underneath braces
+            if (isCommentBeneathSingleLine && resultElements.count() == 1) {
+                val element = resultElements.first
+                element?.add(psiFactory.createWhiteSpace("\n"))
+                element?.add(restored)
+                restored.delete()
+            }
+
             bindNewElement(restored, commentTreeElement) // will be used when restoring line breaks
 
             if (restored.tokenType == KtTokens.EOL_COMMENT) {
@@ -352,9 +361,9 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
             fun findRestored(leaf: TreeElement): PsiElement? {
                 if (leaf is LineBreakTreeElement) return null
                 return leaf.parentsWithSelf
-                        .takeWhile { it != lineBreakParent }
-                        .mapNotNull { toNewPsiElementMap[it]?.first() } //TODO: what about multiple?
-                        .firstOrNull()
+                    .takeWhile { it != lineBreakParent }
+                    .mapNotNull { toNewPsiElementMap[it]?.first() } //TODO: what about multiple?
+                    .firstOrNull()
             }
 
 
@@ -365,8 +374,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                     psiElement = skipTokensForward(psiElement, tokensToMatch.asReversed())
                     psiElement?.restoreLineBreakAfter()
                     break
-                }
-                else {
+                } else {
                     if (leaf !is TokenTreeElement) break
                     tokensToMatch.add(leaf.tokenType)
                 }
@@ -391,8 +399,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
 
         if (whitespace == null) {
             addAfter.parent.addAfter(psiFactory.createNewLine(), addAfter)
-        }
-        else {
+        } else {
             whitespace.replace(psiFactory.createWhiteSpace("\n" + whitespace.text))
         }
 
@@ -419,8 +426,8 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
 
     private fun PsiElement.anchorToAddCommentOrSpace(before: Boolean): PsiElement {
         return parentsWithSelf
-                .dropWhile { it.parent !is PsiFile && (if (before) it.prevSibling else it.nextSibling) == null }
-                .first()
+            .dropWhile { it.parent !is PsiFile && (if (before) it.prevSibling else it.nextSibling) == null }
+            .first()
     }
 
     private fun chooseAnchor(anchorBefore: Anchor?, anchorAfter: Anchor?): Anchor? {

idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt

@@ -17,14 +17,15 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
-import com.intellij.psi.PsiComment
-import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.*
+import org.jetbrains.kotlin.idea.refactoring.getLineCount
 import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
+import org.jetbrains.kotlin.j2k.isInSingleLine
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
-import java.lang.IllegalArgumentException
+import org.jetbrains.kotlin.resolve.calls.CallExpressionElement
 
 class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, "Add braces") {
     override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean {
@@ -51,7 +52,7 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
     override fun applyTo(element: KtElement, editor: Editor?) {
         if (editor == null) throw IllegalArgumentException("This intention requires an editor")
         val expression = element.getTargetExpression(editor.caretModel.offset)!!
-
+        var isCommentBeneath = false
         val psiFactory = KtPsiFactory(element)
 
         val semicolon = element.getNextSiblingIgnoringWhitespaceAndComments()?.takeIf { it.node.elementType == KtTokens.SEMICOLON }
@@ -63,6 +64,21 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
                 semicolon.delete()
         }
 
+        // Check for single line if expression
+        if (element is KtIfExpression && expression.isInSingleLine() && element.`else` == null) {
+            // Check if a comment is actually underneath (\n) the expression
+            val allElements = element.siblings(withItself = false).filterIsInstance<PsiElement>()
+            val sibling = allElements.firstOrNull { it is PsiComment }
+            if (sibling is PsiComment) {
+                // Check if \n before first received comment sibling
+                // if false, the normal procedure of adding braces occurs.
+                isCommentBeneath =
+                        sibling.prevSibling is PsiWhiteSpace &&
+                        sibling.prevSibling.textContains('\n') &&
+                        (sibling.prevSibling.prevSibling is PsiComment || sibling.prevSibling.prevSibling is PsiElement)
+            }
+        }
+
         val nextComment = when {
             element is KtDoWhileExpression -> null // bound to the closing while
             element is KtIfExpression && expression === element.then && element.`else` != null -> null // bound to else
@@ -83,7 +99,9 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
             is KtIfExpression ->
                 (result?.parent?.nextSibling as? PsiWhiteSpace)?.delete()
         }
-        saver.restore(result)
+
+        // Check for single line expression with comment beneath.
+        saver.restore(result, isCommentBeneath, forceAdjustIndent = false)
     }
 
     private fun KtElement.getTargetExpression(caretLocation: Int): KtExpression? {

idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt

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

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

@@ -0,0 +1,8 @@
+fun foo() {}
+
+fun test(b: Boolean) {
+    if (b) {
+        foo()
+    }
+    // ccc
+}

idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt

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

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

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

idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt

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

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

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

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -696,6 +696,21 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt");
         }
 
+        @TestMetadata("addBracesForSingleLineIfWithCommentBeneath.kt")
+        public void testAddBracesForSingleLineIfWithCommentBeneath() throws Exception {
+            runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt");
+        }
+
+        @TestMetadata("addBracesForSingleLineIfWithCommentBeneath2.kt")
+        public void testAddBracesForSingleLineIfWithCommentBeneath2() throws Exception {
+            runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt");
+        }
+
+        @TestMetadata("addBracesForSingleLineIfWithCommentBeneath3.kt")
+        public void testAddBracesForSingleLineIfWithCommentBeneath3() throws Exception {
+            runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt");
+        }
+
         @TestMetadata("addBracesForWhile.kt")
         public void testAddBracesForWhile() throws Exception {
             runTest("idea/testData/intentions/addBraces/addBracesForWhile.kt");