Java: Fixed the intention "Invert 'if' condition" in the case of nested 'if' without braces (IDEA-167957)

java/java-impl/src/com/intellij/codeInsight/intention/impl/InvertIfConditionAction.java

@@ -84,7 +84,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
 
     PsiExpression condition = (PsiExpression) ifStatement.getCondition().copy();
 
-    setupBranches(ifStatement, controlFlow);
+    ifStatement = setupBranches(ifStatement, controlFlow);
     if (condition != null) {
       ifStatement.getCondition().replace(CodeInsightServicesUtil.invertCondition(condition));
     }
@@ -156,10 +156,6 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     return null;
   }
 
-  private static PsiElement findNearestCodeBlock(PsiIfStatement ifStatement) {
-    return PsiTreeUtil.getParentOfType(ifStatement, PsiCodeBlock.class);
-  }
-
   private static ControlFlow buildControlFlow(PsiElement element) {
     try {
       return ControlFlowFactory.getInstance(element.getProject()).getControlFlow(element, LocalsOrMyInstanceFieldsControlFlowPolicy.getInstance(), false);
@@ -169,7 +165,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     }
   }
 
-  private static void setupBranches(PsiIfStatement ifStatement, ControlFlow flow) throws IncorrectOperationException {
+  private static PsiIfStatement setupBranches(PsiIfStatement ifStatement, ControlFlow flow) throws IncorrectOperationException {
     PsiElementFactory factory = JavaPsiFacade.getInstance(ifStatement.getProject()).getElementFactory();
     Project project = ifStatement.getProject();
 
@@ -180,7 +176,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
       elseBranch = (PsiStatement) elseBranch.copy();
       setElseBranch(ifStatement, thenBranch, flow);
       ifStatement.getThenBranch().replace(elseBranch);
-      return;
+      return ifStatement;
     }
 
     final CodeStyleManager codeStyle = CodeStyleManager.getInstance(project);
@@ -190,7 +186,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
       statement = (PsiStatement) codeStyle.reformat(statement);
       statement = (PsiStatement) ifStatement.getThenBranch().replace(statement);
       codeStyle.reformat(statement);
-      return;
+      return ifStatement;
     }
 
     int endOffset = calcEndOffset(flow, ifStatement);
@@ -211,11 +207,11 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         }
       } else {
         if (!(thenBranch instanceof PsiReturnStatement)) {
-          addAfter(ifStatement, thenBranch);
+          ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch);
         }
       }
       ifStatement.getThenBranch().replace(statement);
-      return;
+      return ifStatement;
     }
     PsiElement element = flow.getElement(endOffset);
     while (element != null && !(element instanceof PsiStatement)) element = element.getParent();
@@ -225,27 +221,38 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         element instanceof PsiForeachStatement && flow.getStartOffset(element) + 1 == endOffset) {
       PsiStatement statement = factory.createStatementFromText("continue;", null);
       statement = (PsiStatement)codeStyle.reformat(statement);
-      addAfter(ifStatement, thenBranch);
+      ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch);
       ifStatement.getThenBranch().replace(statement);
-      return;
+      return ifStatement;
     }
 
     if (element instanceof PsiReturnStatement) {
       PsiReturnStatement returnStatement = (PsiReturnStatement) element;
-      addAfter(ifStatement, thenBranch);
+      ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch);
       ifStatement.getThenBranch().replace(returnStatement.copy());
 
       ControlFlow flow2 = buildControlFlow(findCodeBlock(ifStatement));
       if (!ControlFlowUtil.isInstructionReachable(flow2, flow2.getStartOffset(returnStatement), 0)) returnStatement.delete();
-      return;
+      return ifStatement;
     }
 
     boolean nextUnreachable = flow.getEndOffset(ifStatement) == flow.getSize();
     if (!nextUnreachable) {
-      PsiElement nearestCodeBlock = findNearestCodeBlock(ifStatement);
-      if (nearestCodeBlock != null) {
-        ControlFlow flow2 = buildControlFlow(nearestCodeBlock);
-        nextUnreachable = !ControlFlowUtil.isInstructionReachable(flow2, flow2.getEndOffset(ifStatement), getThenOffset(flow2, ifStatement));
+      PsiElement parent = ifStatement.getParent();
+      if (parent != null) {
+        if (!(parent instanceof PsiCodeBlock)) {
+          PsiCodeBlock codeBlock = factory.createCodeBlockFromText("{}", ifStatement);
+          codeBlock = (PsiCodeBlock)codeStyle.reformat(codeBlock);
+          codeBlock.add(ifStatement);
+          codeBlock = (PsiCodeBlock)ifStatement.replace(codeBlock);
+          ifStatement = (PsiIfStatement)codeBlock.getStatements()[0];
+          parent = ifStatement.getParent();
+          thenBranch = ifStatement.getThenBranch();
+        }
+        ControlFlow localFlow = buildControlFlow(parent);
+        int startThenOffset = getThenOffset(localFlow, ifStatement);
+        int afterIfOffset = localFlow.getEndOffset(ifStatement);
+        nextUnreachable = !ControlFlowUtil.isInstructionReachable(localFlow, afterIfOffset, startThenOffset);
       }
     }
     if (nextUnreachable) {
@@ -270,7 +277,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         ifStatement.getThenBranch().replace(codeBlock);
       }
       codeStyle.reformat(ifStatement);
-      return;
+      return ifStatement;
     }
 
     setElseBranch(ifStatement, thenBranch, flow);
@@ -278,6 +285,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     statement = (PsiStatement) codeStyle.reformat(statement);
     statement = (PsiStatement) ifStatement.getThenBranch().replace(statement);
     codeStyle.reformat(statement);
+    return ifStatement;
   }
 
   private static void setElseBranch(PsiIfStatement ifStatement, PsiStatement thenBranch, ControlFlow flow)
@@ -312,6 +320,26 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     ifStatement.setElseBranch(thenBranch);
   }
 
+  private static PsiIfStatement addAfterWithinCodeBlock(PsiIfStatement ifStatement, PsiStatement branch) {
+    final PsiElement parent = ifStatement.getParent();
+    if (parent != null && !(parent instanceof PsiCodeBlock)) {
+      final Project project = ifStatement.getProject();
+      final PsiElementFactory factory = JavaPsiFacade.getInstance(project).getElementFactory();
+      final CodeStyleManager codeStyle = CodeStyleManager.getInstance(project);
+      PsiCodeBlock codeBlock = factory.createCodeBlockFromText("{}", ifStatement);
+      codeBlock = (PsiCodeBlock)codeStyle.reformat(codeBlock);
+      final PsiIfStatement wrappedIfStatement = (PsiIfStatement)codeBlock.add(ifStatement);
+
+      addAfter(wrappedIfStatement, branch);
+      codeBlock = (PsiCodeBlock)ifStatement.replace(codeBlock);
+      return (PsiIfStatement)codeBlock.getStatements()[0];
+    }
+    else {
+      addAfter(ifStatement, branch);
+      return ifStatement;
+    }
+  }
+
   static void addAfter(PsiIfStatement ifStatement, PsiStatement branch) throws IncorrectOperationException {
     if (branch instanceof PsiBlockStatement) {
       PsiBlockStatement blockStatement = (PsiBlockStatement) branch;

java/java-tests/testData/codeInsight/invertIfCondition/afterNestedIf1.java

@@ -0,0 +1,12 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        if (a) {
+            if (b) {
+                return false;
+            }
+            return true;
+        }
+        return false;
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/afterNestedIf2.java

@@ -0,0 +1,13 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        for (int i = 1; i < 10; i++)
+            if (a) {
+                if (b)   /* comment 1 */ {
+                    continue;
+                }
+                return true; /* comment 2 */
+            }
+        return false;            /* comment 3 */
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/afterNestedIf3.java

@@ -0,0 +1,14 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        if (a) {
+            if (b) {
+            }
+            else {
+                return true; // comment
+            }
+        }
+        int x = 1;
+        return false;
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/afterNestedIf4.java

@@ -0,0 +1,13 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        for (int i = 1; i < 10; i++)
+        {
+            if (b) {
+                continue;
+            }
+            return true;
+        }
+        return false;
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/beforeNestedIf1.java

@@ -0,0 +1,9 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        if (a)
+            <caret>if (!b)
+                return true;
+        return false;
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/beforeNestedIf2.java

@@ -0,0 +1,10 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        for (int i = 1; i < 10; i++)
+            if (a)
+                <caret>if (!b)   /* comment 1 */
+                    return true; /* comment 2 */
+        return false;            /* comment 3 */
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/beforeNestedIf3.java

@@ -0,0 +1,10 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        if (a)
+            <caret>if (!b)
+                return true; // comment
+        int x = 1;
+        return false;
+    }
+}

java/java-tests/testData/codeInsight/invertIfCondition/beforeNestedIf4.java

@@ -0,0 +1,9 @@
+// "Invert 'if' condition" "true"
+class Main {
+    boolean method(boolean a, boolean b) {
+        for (int i = 1; i < 10; i++)
+            <caret>if (!b)
+            return true;
+        return false;
+    }
+}