"Invert if condition" intention works better for non-symmetric if's

#KT-5009 Fixed

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

@@ -17,9 +17,21 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.tree.IElementType
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.util.isUnit
+import org.jetbrains.kotlin.idea.util.psi.patternMatching.matches
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
+import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
 
 public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpression>(javaClass(), "Invert 'if' condition") {
     override fun isApplicableTo(element: JetIfExpression, caretOffset: Int): Boolean {
@@ -30,10 +42,11 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
     override fun applyTo(element: JetIfExpression, editor: Editor) {
         val psiFactory = JetPsiFactory(element)
 
-        val condition = element.getCondition()!!
+        val newCondition = negate(element.getCondition()!!)
-        val newCondition = negate(condition)
 
-        val thenBranch = element.getThen()
+        if (handleSpecialCases(element, newCondition)) return
+
+        val thenBranch = element.getThen()!!
         val elseBranch = element.getElse() ?: psiFactory.createEmptyBody()
 
         val newThen = if (elseBranch is JetIfExpression)
@@ -49,6 +62,140 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
         element.replace(psiFactory.createIf(newCondition, newThen, newElse))
     }
 
+    private fun handleSpecialCases(ifExpression: JetIfExpression, newCondition: JetExpression): Boolean {
+        val elseBranch = ifExpression.getElse()
+        if (elseBranch != null) return false
+
+        val factory = JetPsiFactory(ifExpression)
+
+        val thenBranch = ifExpression.getThen()!!
+        val lastThenStatement = thenBranch.lastBlockStatementOrThis()
+        if (lastThenStatement.isExitStatement()) {
+            val block = ifExpression.getParent() as? JetBlockExpression
+            if (block != null) {
+                val rBrace = block.getRBrace()
+                val afterIfInBlock = ifExpression.siblings(withItself = false)
+                        .takeWhile { it != rBrace }
+                        .toList()
+                val lastStatementInBlock = afterIfInBlock.lastIsInstanceOrNull<JetExpression>()
+                if (lastStatementInBlock != null) {
+                    val exitStatementAfterIf = if (lastStatementInBlock.isExitStatement())
+                        lastStatementInBlock
+                    else
+                        exitStatementExecutedAfter(lastStatementInBlock)
+                    if (exitStatementAfterIf != null) {
+                        val first = afterIfInBlock.first()
+                        val last = afterIfInBlock.last()
+                        // build new then branch text from statements after if (will add exit statement if necessary later)
+                        var newIfBodyText = ifExpression.getContainingFile().getText().substring(first.startOffset, last.endOffset).trim()
+
+                        // remove statements after if as they are moving under if
+                        block.deleteChildRange(first, last)
+
+                        if (lastThenStatement is JetReturnExpression && lastThenStatement.getReturnedExpression() == null) {
+                            lastThenStatement.delete()
+                        }
+                        copyThenBranchAfter(ifExpression)
+
+                        // check if we need to add exit statement to then branch
+                        if (exitStatementAfterIf != lastStatementInBlock) {
+                            // don't insert the exit statement, if the new if statement placement has the same exit statement executed after it
+                            val exitAfterNewIf = exitStatementExecutedAfter(ifExpression)
+                            if (exitAfterNewIf == null || !exitAfterNewIf.matches(exitStatementAfterIf)) {
+                                newIfBodyText += "\n" + exitStatementAfterIf.getText()
+                            }
+                        }
+
+                        //TODO: no block if single?
+                        ifExpression.replace(factory.createExpressionByPattern("if ($0) { $1 }", newCondition, newIfBodyText))
+                        return true
+                    }
+                }
+            }
+        }
+
+
+        val exitStatement = exitStatementExecutedAfter(ifExpression) ?: return false
+
+        copyThenBranchAfter(ifExpression)
+        ifExpression.replace(factory.createExpressionByPattern("if ($0) $1", newCondition, exitStatement))
+
+        return true
+    }
+
+    private fun copyThenBranchAfter(ifExpression: JetIfExpression) {
+        val factory = JetPsiFactory(ifExpression)
+        val thenBranch = ifExpression.getThen() ?: return
+        if (thenBranch is JetBlockExpression) {
+            val range = thenBranch.contentRange()
+            if (range != null) {
+                ifExpression.getParent().addRangeAfter(range.first, range.second, ifExpression)
+                ifExpression.getParent().addAfter(factory.createNewLine(), ifExpression)
+            }
+        }
+        else {
+            ifExpression.getParent().addAfter(thenBranch, ifExpression)
+            ifExpression.getParent().addAfter(factory.createNewLine(), ifExpression)
+        }
+    }
+
+    private fun JetBlockExpression.contentRange(): Pair<PsiElement, PsiElement>? {
+        val first = getLBrace()?.siblings(withItself = false)?.firstOrNull { it !is PsiWhiteSpace } ?: return null
+        val rBrace = getRBrace()
+        if (first == rBrace) return null
+        val last = rBrace!!.siblings(forward = false, withItself = false).first { it !is PsiWhiteSpace }
+        return Pair(first, last)
+    }
+
+    private fun exitStatementExecutedAfter(expression: JetExpression): JetExpression? {
+        val block = expression.getParent() as? JetBlockExpression ?: return null //TODO?
+
+        val lastStatement = block.getStatements().lastIsInstanceOrNull<JetExpression>()!!
+        if (expression != lastStatement) {
+            if (lastStatement.isExitStatement() && expression.siblings(withItself = false).firstIsInstance<JetExpression>() == lastStatement) {
+                return lastStatement
+            }
+            return null
+        }
+
+        val parent = block.getParent()
+        when (parent) {
+            is JetNamedFunction -> {
+                if (parent.getBodyExpression() == block) {
+                    if (!parent.hasBlockBody()) return null
+                    val returnType = (parent.resolveToDescriptor() as FunctionDescriptor).getReturnType()
+                    if (returnType == null || !returnType.isUnit()) return null
+                    return JetPsiFactory(expression).createExpression("return")
+                }
+            }
+
+            is JetContainerNode -> {
+                val pparent = parent.getParent()
+                when (pparent) {
+                    is JetLoopExpression -> {
+                        if (block == pparent.getBody()) {
+                            return JetPsiFactory(expression).createExpression("continue")
+                        }
+                    }
+
+                    is JetIfExpression -> {
+                        if (block == pparent.getThen() || block == pparent.getElse()) {
+                            return exitStatementExecutedAfter(pparent)
+                        }
+                    }
+                }
+            }
+        }
+        return null
+    }
+
+    private fun JetExpression.isExitStatement(): Boolean {
+        when (this) {
+            is JetContinueExpression, is JetBreakExpression, is JetThrowExpression, is JetReturnExpression -> return true
+            else -> return false
+        }
+    }
+
     companion object {
         private val NEGATABLE_OPERATORS = setOf(JetTokens.EQEQ, JetTokens.EXCLEQ, JetTokens.EQEQEQ,
                 JetTokens.EXCLEQEQEQ, JetTokens.IS_KEYWORD, JetTokens.NOT_IS, JetTokens.IN_KEYWORD,
@@ -91,7 +238,7 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
                 }
 
                 is JetBinaryExpression -> {
-                    val operator = expression.getOperationToken() ?: return null
+                    val operator = expression.getOperationToken()
                     if (operator !in NEGATABLE_OPERATORS) return null
                     val left = expression.getLeft() ?: return null
                     val right = expression.getRight() ?: return null

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

@@ -1,10 +1,8 @@
 fun foo(): Int {
     val x = 0
 
-    if (x != 1) {
+    if (x != 1) return x
-    } else {
+    x + 1
-        x + 1
-    }
 
     return x
 }

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

@@ -3,9 +3,7 @@ fun main() {
     val list = 1..4
 
     for (x in list) {
-        if (x != 2) {
+        if (x != 2) continue
-        } else {
+        list
-            list
-        }
     }
 }

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

@@ -2,7 +2,7 @@ fun foo(): Int {
     val x = 0
 
     if (x != 1) {
-    } else return x + 1
+        return 0
-
+    }
-    return 0
+    return x + 1
 }

idea/testData/intentions/invertIfCondition/ifThenReturn.kt

@@ -0,0 +1,13 @@
+fun foo(): Boolean {
+    val x = 2
+    <caret>if (x <= 1) {
+        bar1()
+        bar2()
+        return false
+    }
+    bar1()
+    return true
+}
+
+fun bar1(){}
+fun bar2(){}

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

@@ -0,0 +1,13 @@
+fun foo(): Boolean {
+    val x = 2
+    <caret>if (x > 1) {
+        bar1()
+        return true
+    }
+    bar1()
+    bar2()
+    return false
+}
+
+fun bar1(){}
+fun bar2(){}

idea/testData/intentions/invertIfCondition/ifThenReturn2.kt

@@ -0,0 +1,12 @@
+fun foo() {
+    val x = 2
+    <caret>if (x <= 1) {
+        bar1()
+        bar2()
+        return
+    }
+    bar1()
+}
+
+fun bar1(){}
+fun bar2(){}

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

@@ -0,0 +1,12 @@
+fun foo() {
+    val x = 2
+    <caret>if (x > 1) {
+        bar1()
+        return
+    }
+    bar1()
+    bar2()
+}
+
+fun bar1(){}
+fun bar2(){}

idea/testData/intentions/invertIfCondition/ifThenReturn3.kt

@@ -0,0 +1,7 @@
+fun foo() {
+    val x = 2
+    <caret>if (x <= 1) return
+    bar()
+}
+
+fun bar(){}

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

@@ -0,0 +1,8 @@
+fun foo() {
+    val x = 2
+    <caret>if (x > 1) {
+        bar()
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/ifThenReturn4.kt

@@ -0,0 +1,14 @@
+fun foo(p: Int): Int {
+    if (p > 0) {
+        val x = p / 2
+        <caret>if (x <= 1) return 1
+        bar1()
+    }
+    else {
+        bar2()
+    }
+    return 2
+}
+
+fun bar1(){}
+fun bar2(){}

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

@@ -0,0 +1,17 @@
+fun foo(p: Int): Int {
+    if (p > 0) {
+        val x = p / 2
+        if (x > 1) {
+            bar1()
+            return 2
+        }
+        return 1
+    }
+    else {
+        bar2()
+    }
+    return 2
+}
+
+fun bar1(){}
+fun bar2(){}

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

@@ -1,4 +1,3 @@
 fun main() {
-    if (true != false) {
+    if (true != false) return
-    }
 }

idea/testData/intentions/invertIfCondition/lastStatement1.kt

@@ -0,0 +1,12 @@
+fun foo() {
+    val x = 2
+    <caret>if (x > 1) {
+        // invoke bar1&bar2
+        bar1()
+        bar2()
+        // done
+    }
+}
+
+fun bar1(){}
+fun bar2(){}

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

@@ -0,0 +1,11 @@
+fun foo() {
+    val x = 2
+    <caret>if (x <= 1) return
+    // invoke bar1&bar2
+    bar1()
+    bar2()
+    // done
+}
+
+fun bar1(){}
+fun bar2(){}

idea/testData/intentions/invertIfCondition/lastStatement2.kt

@@ -0,0 +1,6 @@
+fun foo() {
+    val x = 2
+    <caret>if (x > 1) bar()
+}
+
+fun bar(){}

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

@@ -0,0 +1,7 @@
+fun foo() {
+    val x = 2
+    <caret>if (x <= 1) return
+    bar()
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatement3.kt

@@ -0,0 +1,8 @@
+fun foo(p: Int) {
+    val x = 2
+    if (p > 0) {
+        <caret>if (x > 1) bar()
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,9 @@
+fun foo(p: Int) {
+    val x = 2
+    if (p > 0) {
+        <caret>if (x <= 1) return
+        bar()
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementBeforeBreak.kt

@@ -0,0 +1,13 @@
+fun foo() {
+    for (i in 1..10) {
+        if (i > 3) {
+            <caret>if (i > 5) {
+                bar()
+            }
+            break
+        }
+        bar()
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,12 @@
+fun foo() {
+    for (i in 1..10) {
+        if (i > 3) {
+            <caret>if (i <= 5) break
+            bar()
+            break
+        }
+        bar()
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementBeforeContinue.kt

@@ -0,0 +1,13 @@
+fun foo() {
+    for (i in 1..10) {
+        if (i > 3) {
+            <caret>if (i > 5) {
+                bar()
+            }
+            continue
+        }
+        bar()
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,12 @@
+fun foo() {
+    for (i in 1..10) {
+        if (i > 3) {
+            <caret>if (i <= 5) continue
+            bar()
+            continue
+        }
+        bar()
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementBeforeReturn.kt

@@ -0,0 +1,9 @@
+fun foo(p: Boolean): Boolean {
+    val x = 2
+    <caret>if (x > 1) {
+        bar()
+    }
+    return p
+}
+
+fun bar(){}

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

@@ -0,0 +1,8 @@
+fun foo(p: Boolean): Boolean {
+    val x = 2
+    <caret>if (x <= 1) return p
+    bar()
+    return p
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementInLambda.kt

@@ -0,0 +1,9 @@
+fun foo() {
+    val v: (Int) -> Unit = {
+        <caret>if (it > 1) {
+            bar()
+        }
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,10 @@
+fun foo() {
+    val v: (Int) -> Unit = {
+        <caret>if (it <= 1) {
+        } else {
+            bar()
+        }
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementInLoop.kt

@@ -0,0 +1,9 @@
+fun foo() {
+    for (i in 1..10) {
+        <caret>if (i > 1) {
+            bar()
+        }
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,8 @@
+fun foo() {
+    for (i in 1..10) {
+        <caret>if (i <= 1) continue
+        bar()
+    }
+}
+
+fun bar(){}

idea/testData/intentions/invertIfCondition/lastStatementNonUnitMethod.kt

@@ -0,0 +1,10 @@
+// ERROR: A 'return' expression required in a function with a block body ('{...}')
+
+fun foo(): Int {
+    val x = 2
+    <caret>if (x > 1) {
+        bar()
+    }
+}
+
+fun bar(){}

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

@@ -0,0 +1,11 @@
+// ERROR: A 'return' expression required in a function with a block body ('{...}')
+
+fun foo(): Int {
+    val x = 2
+    <caret>if (x <= 1) {
+    } else {
+        bar()
+    }
+}
+
+fun bar(){}

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

@@ -3,6 +3,5 @@ fun foo(): Boolean {
 }
 
 fun main() {
-    if (foo()) {
+    if (foo()) return
-    }
 }

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

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

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

@@ -4669,12 +4669,90 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("ifThenReturn.kt")
+        public void testIfThenReturn() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/ifThenReturn.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifThenReturn2.kt")
+        public void testIfThenReturn2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/ifThenReturn2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifThenReturn3.kt")
+        public void testIfThenReturn3() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/ifThenReturn3.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifThenReturn4.kt")
+        public void testIfThenReturn4() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/ifThenReturn4.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("invertableOperator.kt")
         public void testInvertableOperator() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/invertableOperator.kt");
             doTest(fileName);
         }
 
+        @TestMetadata("lastStatement1.kt")
+        public void testLastStatement1() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatement1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatement2.kt")
+        public void testLastStatement2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatement2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatement3.kt")
+        public void testLastStatement3() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatement3.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementBeforeBreak.kt")
+        public void testLastStatementBeforeBreak() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementBeforeBreak.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementBeforeContinue.kt")
+        public void testLastStatementBeforeContinue() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementBeforeContinue.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementBeforeReturn.kt")
+        public void testLastStatementBeforeReturn() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementBeforeReturn.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementInLambda.kt")
+        public void testLastStatementInLambda() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementInLambda.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementInLoop.kt")
+        public void testLastStatementInLoop() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementInLoop.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lastStatementNonUnitMethod.kt")
+        public void testLastStatementNonUnitMethod() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/lastStatementNonUnitMethod.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("negatedExpression.kt")
         public void testNegatedExpression() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/invertIfCondition/negatedExpression.kt");