Split if intention preserves comments

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

@@ -27,6 +27,16 @@ public val PsiElement.startOffset: Int
 public val PsiElement.endOffset: Int
     get() = getTextRange().getEndOffset()
 
+public fun PsiElement.getStartOffsetIn(ancestor: PsiElement): Int {
+    var offset = 0
+    var parent = this
+    while (parent != this) {
+        offset += parent.getStartOffsetInParent()
+        parent = parent.getParent()
+    }
+    return offset
+}
+
 public data class PsiChildRange(public val first: PsiElement?, public val last: PsiElement?) : Sequence<PsiElement> {
     init {
         if (first == null) {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/CommentSaver.kt

@@ -17,13 +17,18 @@
 package org.jetbrains.kotlin.idea.core
 
 import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiRecursiveElementVisitor
 import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.psi.JetFile
 import org.jetbrains.kotlin.psi.JetPsiFactory
 import org.jetbrains.kotlin.psi.psiUtil.*
-import java.util.*
+import java.util.ArrayList
+import java.util.HashMap
+import java.util.LinkedHashMap
 
 public class CommentSaver(originalElements: PsiChildRange) {
     public constructor(originalElement: PsiElement) : this(PsiChildRange.singleElement(originalElement))
@@ -60,6 +65,8 @@ public class CommentSaver(originalElements: PsiChildRange) {
 
         val nextElements: Sequence<TreeElement>
             get() = nextSiblings.flatMap { it.withDescendants(leftToRight = true) }
+
+//        var debugText: String? = null
     }
 
     private data class CommentData(
@@ -98,7 +105,8 @@ public class CommentSaver(originalElements: PsiChildRange) {
             assert(child.savedTreeElement == null)
 
             val savedChild = TreeElement(parent = parentTreeElement, prev = last)
-            child.putCopyableUserData(SAVED_TREE_KEY, savedChild)
+//            savedChild.debugText = child.getText()
+            child.savedTreeElement = savedChild
             last?.next = savedChild
             last = savedChild
 
@@ -116,8 +124,9 @@ public class CommentSaver(originalElements: PsiChildRange) {
     private val PsiElement.shouldSave: Boolean
         get() = this !is PsiWhiteSpace
 
-    private val PsiElement.savedTreeElement: TreeElement?
+    private var PsiElement.savedTreeElement: TreeElement?
         get() = getCopyableUserData(SAVED_TREE_KEY)
+        set(value) = putCopyableUserData(SAVED_TREE_KEY, value)
 
     public fun deleteCommentsInside(element: PsiElement) {
         element.accept(object : PsiRecursiveElementVisitor() {
@@ -130,6 +139,32 @@ public class CommentSaver(originalElements: PsiChildRange) {
         })
     }
 
+    public fun elementCreatedByText(createdElement: PsiElement, original: PsiElement, rangeInOriginal: TextRange) {
+        assert(createdElement.getTextLength() == rangeInOriginal.getLength())
+        assert(createdElement.getText() == original.getText().substring(rangeInOriginal.getStartOffset(), rangeInOriginal.getEndOffset()))
+
+        createdElement.accept(object : PsiRecursiveElementVisitor() {
+            override fun visitElement(element: PsiElement) {
+                if (!element.shouldSave) return
+
+                val token = original.findElementAt(element.getStartOffsetIn(createdElement) + rangeInOriginal.getStartOffset())
+                if (token != null) {
+                    val elementLength = element.getTextLength()
+                    for (originalElement in token.parents()) {
+                        val length = originalElement.getTextLength()
+                        if (length < elementLength) continue
+                        if (length == elementLength) {
+                            element.savedTreeElement = originalElement.savedTreeElement
+                        }
+                        break
+                    }
+                }
+
+                super.visitElement(element)
+            }
+        })
+    }
+
     public fun restoreComments(resultElement: PsiElement) {
         restoreComments(PsiChildRange.singleElement(resultElement))
     }
@@ -141,7 +176,7 @@ public class CommentSaver(originalElements: PsiChildRange) {
         resultElements.forEach {
             it.accept(object : PsiRecursiveElementVisitor() {
                 override fun visitElement(element: PsiElement) {
-                    element.putCopyableUserData(SAVED_TREE_KEY, null)
+                    element.savedTreeElement = null
                     super.visitElement(element)
                 }
             })

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

@@ -55,7 +55,7 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
         val elseBranch = ifExpression.getElse() ?: psiFactory.createEmptyBody()
 
         val newThen = if (elseBranch is JetIfExpression)
-            psiFactory.wrapInABlock(elseBranch)
+            psiFactory.createSingleStatementBlock(elseBranch)
         else
             elseBranch
 

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

@@ -17,10 +17,16 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
-import org.jetbrains.kotlin.lexer.JetTokens
+import com.intellij.openapi.util.TextRange
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.idea.conversion.copy.length
+import org.jetbrains.kotlin.idea.conversion.copy.range
+import org.jetbrains.kotlin.idea.conversion.copy.start
+import org.jetbrains.kotlin.idea.core.CommentSaver
+import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
@@ -40,8 +46,11 @@ public class SplitIfIntention : JetSelfTargetingIntention<JetExpression>(javaCla
         }
 
         val ifExpression = operator.getNonStrictParentOfType<JetIfExpression>()
+
+        val commentSaver = CommentSaver(ifExpression!!)
+
         val expression = operator.getParent() as JetBinaryExpression
-        val rightExpression = JetPsiUtil.safeDeparenthesize(getRight(expression, ifExpression!!.getCondition()!!))
+        val rightExpression = JetPsiUtil.safeDeparenthesize(getRight(expression, ifExpression.getCondition()!!, commentSaver))
         val leftExpression = JetPsiUtil.safeDeparenthesize(expression.getLeft()!!)
         val thenBranch = ifExpression.getThen()!!
         val elseBranch = ifExpression.getElse()
@@ -50,32 +59,41 @@ public class SplitIfIntention : JetSelfTargetingIntention<JetExpression>(javaCla
 
         val innerIf = psiFactory.createIf(rightExpression, thenBranch, elseBranch)
 
-        when (operator.getReferencedNameElementType()) {
+        val newIf = when (operator.getReferencedNameElementType()) {
-            JetTokens.ANDAND -> ifExpression.replace(psiFactory.createIf(leftExpression, psiFactory.wrapInABlock(innerIf), elseBranch))
+            JetTokens.ANDAND -> psiFactory.createIf(leftExpression, psiFactory.createSingleStatementBlock(innerIf), elseBranch)
 
             JetTokens.OROR -> {
                 val container = ifExpression.getParent()
 
                 if (container is JetBlockExpression && elseBranch == null && thenBranch.lastBlockStatementOrThis().isExitStatement()) { // special case
-                    container.addAfter(innerIf, ifExpression)
+                    val secondIf = container.addAfter(innerIf, ifExpression)
                     container.addAfter(psiFactory.createNewLine(), ifExpression)
-                    ifExpression.replace(psiFactory.createIf(leftExpression, thenBranch))
+                    val firstIf = ifExpression.replace(psiFactory.createIf(leftExpression, thenBranch))
+                    commentSaver.restoreComments(PsiChildRange(firstIf, secondIf))
                     return
                 }
-
+                else {
-                ifExpression.replace(psiFactory.createIf(leftExpression, thenBranch, innerIf))
+                    psiFactory.createIf(leftExpression, thenBranch, innerIf)
+                }
             }
 
             else -> throw IllegalArgumentException()
         }
+
+        val result = ifExpression.replace(newIf)
+        commentSaver.restoreComments(result)
     }
 
-    private fun getRight(element: JetBinaryExpression, condition: JetExpression): JetExpression {
+    private fun getRight(element: JetBinaryExpression, condition: JetExpression, commentSaver: CommentSaver): JetExpression {
         //gets the textOffset of the right side of the JetBinaryExpression in context to condition
-        val startOffset = element.getRight()!!.startOffset - condition.startOffset
+        val conditionRange = condition.range
-        val rightString = condition.getText()!!.substring(startOffset, condition.getTextLength())
+        val startOffset = element.getRight()!!.startOffset - conditionRange.start
+        val endOffset = conditionRange.length
+        val rightString = condition.getText().substring(startOffset, endOffset)
 
-        return JetPsiFactory(element).createExpression(rightString)
+        val expression = JetPsiFactory(element).createExpression(rightString)
+        commentSaver.elementCreatedByText(expression, condition, TextRange(startOffset, endOffset))
+        return expression
     }
 
     private fun getFirstValidOperator(element: JetIfExpression): JetSimpleNameExpression? {

idea/testData/intentions/splitIf/keepComments/ifOrReturn.kt

@@ -0,0 +1,4 @@
+fun foo(p: Int): Boolean {
+    if (p < 0 /* p < 0 */ <caret>|| p == 5 /* or 5 */) return false
+    return true
+}

idea/testData/intentions/splitIf/keepComments/ifOrReturn.kt.after

@@ -0,0 +1,5 @@
+fun foo(p: Int): Boolean {
+    if (p < 0 /* p < 0 */) return false
+    if (p == 5 /* or 5 */) return false
+    return true
+}

idea/testData/intentions/splitIf/keepComments/twoOperators.kt

@@ -0,0 +1,10 @@
+fun doSomething<T>(a: T) {}
+
+fun foo() {
+    val a = true
+    val b = false
+    val c = true
+    if (a /*a*/ && b /*b*/ &<caret>& c /*c*/) {
+        doSomething("test")
+    }
+}

idea/testData/intentions/splitIf/keepComments/twoOperators.kt.after

@@ -0,0 +1,12 @@
+fun doSomething<T>(a: T) {}
+
+fun foo() {
+    val a = true
+    val b = false
+    val c = true
+    if (a /*a*/ && b /*b*/) {
+        if (c /*c*/) {
+            doSomething("test")
+        }
+    }
+}

idea/testData/intentions/splitIf/keepComments/withAnd.kt

@@ -0,0 +1,7 @@
+fun doSomething<T>(a: T) {}
+
+fun foo(p: Int) {
+    <caret>if (0 < p /* > 0 */ && p < 100 /* not too much */) {
+        doSomething("test")
+    }
+}

idea/testData/intentions/splitIf/keepComments/withAnd.kt.after

@@ -0,0 +1,9 @@
+fun doSomething<T>(a: T) {}
+
+fun foo(p: Int) {
+    <caret>if (0 < p /* > 0 */) {
+        if (p < 100 /* not too much */) {
+            doSomething("test")
+        }
+    }
+}

idea/testData/intentions/splitIf/keepComments/withOR.kt

@@ -0,0 +1,7 @@
+fun doSomething<T>(a: T) {}
+
+fun foo(p: Int) {
+    <caret>if (p < 0 /* p < 0 */ || p > 100 /* too much */) {
+        doSomething("test")
+    }
+}

idea/testData/intentions/splitIf/keepComments/withOR.kt.after

@@ -0,0 +1,9 @@
+fun doSomething<T>(a: T) {}
+
+fun foo(p: Int) {
+    <caret>if (p < 0 /* p < 0 */) {
+        doSomething("test")
+    } else if (p > 100 /* too much */) {
+        doSomething("test")
+    }
+}

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

@@ -3906,7 +3906,7 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             }
 
             public void testAllFilesPresentInKeepComments() throws Exception {
-                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/convertToExpressionBody/keepComments"), Pattern.compile("^(.+)\\.kt$"), true);
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/convertToExpressionBody/keepComments"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
             }
         }
     }
@@ -6629,6 +6629,39 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/splitIf/wrongCaretLocation.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("idea/testData/intentions/splitIf/keepComments")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class KeepComments extends AbstractIntentionTest {
+            public void testAllFilesPresentInKeepComments() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/splitIf/keepComments"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
+            }
+
+            @TestMetadata("ifOrReturn.kt")
+            public void testIfOrReturn() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/splitIf/keepComments/ifOrReturn.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("twoOperators.kt")
+            public void testTwoOperators() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/splitIf/keepComments/twoOperators.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("withAnd.kt")
+            public void testWithAnd() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/splitIf/keepComments/withAnd.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("withOR.kt")
+            public void testWithOR() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/splitIf/keepComments/withOR.kt");
+                doTest(fileName);
+            }
+        }
     }
 
     @TestMetadata("idea/testData/intentions/swapBinaryExpression")