Fix comment placement on if/else in if or when condition #KT-28224 Fixed

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

@@ -28,7 +28,10 @@ import java.util.*
 import kotlin.properties.Delegates
 
 class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: Boolean = false/*TODO?*/) {
-    constructor(originalElement: PsiElement, saveLineBreaks: Boolean = false/*TODO?*/) : this(PsiChildRange.singleElement(originalElement), saveLineBreaks)
+    constructor(originalElement: PsiElement, saveLineBreaks: Boolean = false/*TODO?*/) : this(
+        PsiChildRange.singleElement(originalElement),
+        saveLineBreaks
+    )
 
     private val SAVED_TREE_KEY = Key<TreeElement>("SAVED_TREE")
     private val psiFactory = KtPsiFactory(originalElements.first!!)
@@ -234,15 +237,25 @@ 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,
+        isCommentBeneathSingleLine: Boolean,
+        isCommentInside: Boolean,
+        forceAdjustIndent: Boolean
+    ) {
+        restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent, isCommentBeneathSingleLine, isCommentInside)
     }
 
     fun restore(resultElement: PsiElement, forceAdjustIndent: Boolean = false) {
         restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent)
     }
 
-    fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false, isCommentBeneathSingleLine: Boolean = false) {
+    fun restore(
+        resultElements: PsiChildRange,
+        forceAdjustIndent: Boolean = false,
+        isCommentBeneathSingleLine: Boolean = false,
+        isCommentInside: Boolean = false
+    ) {
         assert(!isFinished)
         assert(!resultElements.isEmpty)
 
@@ -264,7 +277,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                     })
                 }
 
-                restoreComments(resultElements, isCommentBeneathSingleLine)
+                restoreComments(resultElements, isCommentBeneathSingleLine, isCommentInside)
 
                 restoreLineBreaks()
 
@@ -295,7 +308,11 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
         isFinished = true
     }
 
-    private fun restoreComments(resultElements: PsiChildRange, isCommentBeneathSingleLine: Boolean = false) {
+    private fun restoreComments(
+        resultElements: PsiChildRange,
+        isCommentBeneathSingleLine: Boolean = false,
+        isCommentInside: Boolean = false
+    ) {
         var putAbandonedCommentsAfter = resultElements.last!!
 
         for (commentTreeElement in commentsToRestore) {
@@ -333,6 +350,14 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                 }
             } else {
                 restored = putAbandonedCommentsAfter.parent.addBefore(comment, putAbandonedCommentsAfter) as PsiComment
+
+                if (isCommentInside) {
+                    val element = resultElements.first
+                    val innerExpression = element?.lastChild?.getPrevSiblingIgnoringWhitespace()
+                    innerExpression?.add(psiFactory.createWhiteSpace())
+                    innerExpression?.add(restored)
+                    restored.delete()
+                }
                 putAbandonedCommentsAfter = restored
             }
 

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

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.*
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.nextStatement
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.previousStatement
 import org.jetbrains.kotlin.idea.refactoring.getLineCount
 import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
@@ -53,6 +55,7 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
         if (editor == null) throw IllegalArgumentException("This intention requires an editor")
         val expression = element.getTargetExpression(editor.caretModel.offset)!!
         var isCommentBeneath = false
+        var isCommentInside = false
         val psiFactory = KtPsiFactory(element)
 
         val semicolon = element.getNextSiblingIgnoringWhitespaceAndComments()?.takeIf { it.node.elementType == KtTokens.SEMICOLON }
@@ -79,12 +82,31 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
             }
         }
 
+        // Check for nested if/else
+        if (element is KtIfExpression && expression.isInSingleLine() && element.`else` != null &&
+            element.parent.nextSibling is PsiWhiteSpace &&
+            element.parent.nextSibling.nextSibling is PsiComment
+        ) {
+            isCommentInside = true
+        }
+
         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
             else -> element.getNextSiblingIgnoringWhitespace().takeIf { it is PsiComment }
         }
-        val saver = if (nextComment == null) CommentSaver(element) else CommentSaver(PsiChildRange(element, nextComment))
+
+        val saver = when {
+            isCommentInside -> {
+                CommentSaver(element.parent.nextSibling.nextSibling)
+            }
+            else -> if (nextComment == null) CommentSaver(element) else CommentSaver(PsiChildRange(element, nextComment))
+        }
+
+        if (isCommentInside) {
+            element.parent.nextSibling.nextSibling.delete()
+        }
+
         element.allChildren.filterIsInstance<PsiComment>().toList().forEach {
             it.delete()
         }
@@ -101,7 +123,7 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
         }
 
         // Check for single line expression with comment beneath.
-        saver.restore(result, isCommentBeneath, forceAdjustIndent = false)
+        saver.restore(result, isCommentBeneath, isCommentInside, forceAdjustIndent = false)
     }
 
     private fun KtElement.getTargetExpression(caretLocation: Int): KtExpression? {

idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt

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

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

@@ -0,0 +1,7 @@
+fun test(x: Int, b: Boolean) {
+    val foo = if (x == 1)
+        if (b) 1 else {
+            2 // comment
+        }
+    else 0
+}

idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt

@@ -0,0 +1,8 @@
+fun test(x: Int, b: Boolean) {
+    val bar = when (x) {
+        1 ->
+            if (b) 1 else<caret> 2 // comment
+        else ->
+            0
+    }
+}

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

@@ -0,0 +1,10 @@
+fun test(x: Int, b: Boolean) {
+    val bar = when (x) {
+        1 ->
+            if (b) 1 else {
+                2 // comment
+            }
+        else ->
+            0
+    }
+}

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

@@ -656,6 +656,16 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/addBraces/addBracesForElse.kt");
         }
 
+        @TestMetadata("addBracesForElseWithCommentInsideIf.kt")
+        public void testAddBracesForElseWithCommentInsideIf() throws Exception {
+            runTest("idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt");
+        }
+
+        @TestMetadata("addBracesForElseWithCommentInsideWhen.kt")
+        public void testAddBracesForElseWithCommentInsideWhen() throws Exception {
+            runTest("idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt");
+        }
+
         @TestMetadata("addBracesForFor.kt")
         public void testAddBracesForFor() throws Exception {
             runTest("idea/testData/intentions/addBraces/addBracesForFor.kt");