KT-11974 Invert if-condition intention loses comments

#KT-11974 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -163,6 +163,11 @@ tailrec fun PsiElement.getOutermostParentContainedIn(container: PsiElement): Psi
 
 fun PsiElement.isInsideOf(elements: Iterable<PsiElement>): Boolean = elements.any { it.isAncestor(this) }
 
+fun PsiChildRange.trimWhiteSpaces(): PsiChildRange {
+    if (first == null) return this
+    return PsiChildRange(first.siblings().firstOrNull { it !is PsiWhiteSpace }, last!!.siblings(forward = false).firstOrNull { it !is PsiWhiteSpace })
+}
+
 // -------------------- Recursive tree visiting --------------------------------------------------------------------------------------------
 
 inline fun <reified T : PsiElement> PsiElement.forEachDescendantOfType(noinline action: (T) -> Unit) {

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

@@ -17,11 +17,13 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
-import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.moveCaret
+import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.unblockDocument
+import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.idea.util.psi.patternMatching.matches
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
@@ -36,11 +38,24 @@ class InvertIfConditionIntention : SelfTargetingIntention<KtIfExpression>(KtIfEx
     }
 
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
+        val rBrace = parentBlockRBrace(element)
+        val commentSavingRange = if (rBrace != null)
+            PsiChildRange(element, rBrace)
+        else
+            PsiChildRange.singleElement(element)
+        val commentSaver = CommentSaver(commentSavingRange)
+
         val newCondition = element.condition!!.negate()
 
         val newIf = handleSpecialCases(element, newCondition)
                     ?: handleStandardCase(element, newCondition)
 
+        val commentRestoreRange = if (rBrace != null)
+            PsiChildRange(newIf, rBrace)
+        else
+            PsiChildRange(newIf, parentBlockRBrace(newIf) ?: newIf)
+        commentSaver.restore(commentRestoreRange)
+
         val newIfCondition = newIf.condition
         val simplifyIntention = ConvertNegatedExpressionWithDemorgansLawIntention()
         if (newIfCondition is KtPrefixExpression && simplifyIntention.isApplicableTo(newIfCondition)) {
@@ -96,8 +111,10 @@ class InvertIfConditionIntention : SelfTargetingIntention<KtIfExpression>(KtIfEx
                     if (exitStatementAfterIf != null) {
                         val first = afterIfInBlock.first()
                         val last = afterIfInBlock.last()
-                        // build new then branch text from statements after if (we will add exit statement if necessary later)
-                        var newIfBodyText = ifExpression.containingFile.text.substring(first.startOffset, last.endOffset).trim()
+                        // build new then branch from statements after if (we will add exit statement if necessary later)
+                        //TODO: no block if single?
+                        val newThenRange = PsiChildRange(first, last).trimWhiteSpaces()
+                        val newIf = factory.createExpressionByPattern("if ($0) { $1 }", newCondition, newThenRange) as KtIfExpression
 
                         // remove statements after if as they are moving under if
                         block.deleteChildRange(first, last)
@@ -112,12 +129,11 @@ class InvertIfConditionIntention : SelfTargetingIntention<KtIfExpression>(KtIfEx
                             // don't insert the exit statement, if the new if statement placement has the same exit statement executed after it
                             val exitAfterNewIf = exitStatementExecutedAfter(updatedIf)
                             if (exitAfterNewIf == null || !exitAfterNewIf.matches(exitStatementAfterIf)) {
-                                newIfBodyText += "\n" + exitStatementAfterIf.text
+                                val newThen = newIf.then as KtBlockExpression
+                                newThen.addBefore(exitStatementAfterIf, newThen.rBrace)
                             }
                         }
 
-                        //TODO: no block if single?
-                        val newIf = factory.createExpressionByPattern("if ($0) { $1\n}", newCondition, newIfBodyText) // we need to insert '\n' because the text can end with an end-of-line comment
                         return updatedIf.replace(newIf) as KtIfExpression
                     }
                 }
@@ -204,4 +220,8 @@ class InvertIfConditionIntention : SelfTargetingIntention<KtIfExpression>(KtIfEx
         }
         return null
     }
+
+    private fun parentBlockRBrace(element: KtIfExpression): PsiElement? {
+        return (element.parent as? KtBlockExpression)?.rBrace
+    }
 }

idea/testData/intentions/invertIfCondition/addSurroundingBlock_preserveComments.kt

@@ -0,0 +1,9 @@
+fun foo(list: List<String>) {
+    for (s in list)
+        <caret>if (s.length > 0 /* check it */) /* then */ {
+            bar() // bar()
+        }
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/addSurroundingBlock_preserveComments.kt.after

@@ -0,0 +1,9 @@
+fun foo(list: List<String>) {
+    for (s in list) {
+        if (s.length <= 0 /* check it */) /* then */ continue
+        bar() // bar()
+    }
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/assignedToValue.kt

@@ -1,4 +1,4 @@
 fun main() {
     val a = -1
-    val t = <caret>if (a > 0) a else -a
+    val t = <caret>if (a > 0) a /* a */ else -a /* -a */
 }

idea/testData/intentions/invertIfCondition/assignedToValue.kt.after

@@ -1,4 +1,4 @@
 fun main() {
     val a = -1
-    val t = if (a <= 0) -a else a
+    val t = if (a <= 0) -a else a /* a */ /* -a */
 }

idea/testData/intentions/invertIfCondition/simple.kt

@@ -4,6 +4,6 @@ fun foo(): Boolean {
 
 fun main() {
     <caret>if (foo()) {
-
+        //TODO
     }
 }

idea/testData/intentions/invertIfCondition/simple.kt.after

@@ -4,4 +4,5 @@ fun foo(): Boolean {
 
 fun main() {
     <caret>if (!foo()) return
+    //TODO
 }

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

@@ -6227,6 +6227,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
     public static class InvertIfCondition extends AbstractIntentionTest {
+        @TestMetadata("addSurroundingBlock_preserveComments.kt")
+        public void testAddSurroundingBlock_preserveComments() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/addSurroundingBlock_preserveComments.kt");
+            doTest(fileName);
+        }
+
         public void testAllFilesPresentInInvertIfCondition() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/invertIfCondition"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }