Java: Add braces when extracting method from then/else branch or loop body (IDEA-181096)

java/java-impl/src/com/intellij/refactoring/extractMethod/ExtractMethodProcessor.java

@@ -62,10 +62,7 @@ import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
 import com.intellij.util.text.UniqueNameGenerator;
 import one.util.streamex.StreamEx;
-import org.jetbrains.annotations.NonNls;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-import org.jetbrains.annotations.TestOnly;
+import org.jetbrains.annotations.*;
 
 import java.util.*;
 
@@ -124,6 +121,7 @@ public class ExtractMethodProcessor implements MatchProvider {
   protected PsiStatement myFirstExitStatementCopy;
   protected PsiMethod myExtractedMethod;
   private PsiMethodCallExpression myMethodCall;
+  private PsiBlockStatement myAddedToMethodCallLocation;
   protected boolean myNullConditionalCheck;
   protected boolean myNotNullConditionalCheck;
   protected Nullness myNullness;
@@ -1364,18 +1362,63 @@ public class ExtractMethodProcessor implements MatchProvider {
     }
   }
 
-  protected PsiElement addToMethodCallLocation(PsiStatement statement) throws IncorrectOperationException {
-    if (myEnclosingBlockStatement == null) {
-      PsiElement containingStatement = myElements[0] instanceof PsiComment ? myElements[0] : PsiTreeUtil.getParentOfType(myExpression != null ? myExpression : myElements[0], PsiStatement.class, false);
-      if (containingStatement == null) {
-        containingStatement = PsiTreeUtil.getParentOfType(myExpression != null ? myExpression : myElements[0], PsiComment.class, false);
-      }
+  protected PsiElement addToMethodCallLocation(PsiStatement newStatement) throws IncorrectOperationException {
+    if (myAddedToMethodCallLocation != null) {
+      PsiCodeBlock block = myAddedToMethodCallLocation.getCodeBlock();
+      return block.addBefore(newStatement, block.getRBrace());
+    }
 
-      return containingStatement.getParent().addBefore(statement, containingStatement);
+    PsiElement location;
+    PsiStatement oldStatement;
+    if (myEnclosingBlockStatement != null) {
+      location = oldStatement = myEnclosingBlockStatement;
     }
     else {
-      return myEnclosingBlockStatement.getParent().addBefore(statement, myEnclosingBlockStatement);
+      PsiElement element = myExpression != null ? myExpression : myElements[0];
+      PsiComment comment = PsiTreeUtil.getParentOfType(element, PsiComment.class, false);
+
+      if (comment == null) {
+        location = oldStatement = PsiTreeUtil.getParentOfType(element, PsiStatement.class, false);
+      }
+      else {
+        location = comment;
+        oldStatement = StreamEx.of(myElements)
+          .filter(e -> !(e instanceof PsiComment) && !(e instanceof PsiWhiteSpace))
+          .map(e -> PsiTreeUtil.getParentOfType(e, PsiStatement.class, false))
+          .nonNull()
+          .findFirst()
+          .orElse(null);
+      }
     }
+
+    LOG.assertTrue(location != null, "Can't find statement/comment at the extracted location");
+
+    PsiElement parent = location.getParent();
+    if (isBranchOrBody(parent, oldStatement)) {
+      // The parent statement will be inconsistent until deleting the extracted part:
+      // the block statement is being added just before a then/else branch or loop body,
+      // but the original then/else branch or loop body will be deleted in the end so it's fine.
+      // Example: "if(..) oldStatement" -> "if(..) { newStatement } oldStatement" -> "if(..) { newStatement }"
+      myAddedToMethodCallLocation = (PsiBlockStatement)myElementFactory.createStatementFromText("{}", oldStatement);
+      myAddedToMethodCallLocation = (PsiBlockStatement)parent.addBefore(myAddedToMethodCallLocation, location);
+      PsiCodeBlock block = myAddedToMethodCallLocation.getCodeBlock();
+      return block.addBefore(newStatement, block.getRBrace());
+    }
+    return parent.addBefore(newStatement, location);
+  }
+
+  @Contract("_,null -> false; null,_ -> false")
+  private static boolean isBranchOrBody(PsiElement parent, PsiElement element) {
+    if (element == null) {
+      return false;
+    }
+    if (parent instanceof PsiIfStatement) {
+      return (((PsiIfStatement)parent).getThenBranch() == element || ((PsiIfStatement)parent).getElseBranch() == element);
+    }
+    if (parent instanceof PsiLoopStatement) {
+      return ((PsiLoopStatement)parent).getBody() == element;
+    }
+    return false;
   }
 
   private void renameInputVariables() throws IncorrectOperationException {

java/java-tests/testData/refactoring/extractMethod/ExtractBareDoWhileBody.java

@@ -0,0 +1,14 @@
+class DoIfWhile {
+    String foo(int a, boolean b) {
+        int x = 0;
+        do <selection>/*comment*/ if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }</selection>
+        while (++x < a);
+
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareDoWhileBody_after.java

@@ -0,0 +1,26 @@
+import org.jetbrains.annotations.Nullable;
+
+class DoIfWhile {
+    String foo(int a, boolean b) {
+        int x = 0;
+        do {
+            String s = newMethod(b, x);
+            if (s != null) return s;
+        }
+        while (++x < a);
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b, int x) {
+        /*comment*/
+        if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareElseBranch.java

@@ -0,0 +1,14 @@
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) {
+
+        } else <selection>if (b) {
+            String s = bar();
+            if (s != null) return s;
+        }</selection>
+
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareElseBranch_after.java

@@ -0,0 +1,25 @@
+import org.jetbrains.annotations.Nullable;
+
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) {
+
+        } else {
+            String s = newMethod(b);
+            if (s != null) return s;
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b) {
+        if (b) {
+            String s = bar();
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareForBody.java

@@ -0,0 +1,12 @@
+class ForIf {
+    String foo(int[] a, boolean b) {
+        for (int x : a) <selection>if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }</selection>
+
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareForBody_after.java

@@ -0,0 +1,23 @@
+import org.jetbrains.annotations.Nullable;
+
+class ForIf {
+    String foo(int[] a, boolean b) {
+        for (int x : a) {
+            String s = newMethod(b, x);
+            if (s != null) return s;
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b, int x) {
+        if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareThenBranch.java

@@ -0,0 +1,12 @@
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) <selection>if (b) {
+            String s = bar();
+            if (s != null) return s;
+        }</selection>
+
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBareThenBranch_after.java

@@ -0,0 +1,23 @@
+import org.jetbrains.annotations.Nullable;
+
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) {
+            String s = newMethod(b);
+            if (s != null) return s;
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b) {
+        if (b) {
+            String s = bar();
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBracedDoWhileBody.java

@@ -0,0 +1,16 @@
+class DoIfWhile {
+    String foo(int a, boolean b) {
+        int x = 0;
+        do <selection>{
+        /*comment*/ if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }
+        }</selection>
+        while (++x < a);
+
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBracedDoWhileBody_after.java

@@ -0,0 +1,26 @@
+import org.jetbrains.annotations.Nullable;
+
+class DoIfWhile {
+    String foo(int a, boolean b) {
+        int x = 0;
+        do {
+            String s = newMethod(b, x);
+            if (s != null) return s;
+        }
+        while (++x < a);
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b, int x) {
+        /*comment*/
+        if (b) {
+            String s = bar(x);
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar(int x) { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBracedElseBranch.java

@@ -0,0 +1,16 @@
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) {
+
+        } else <selection>{
+            if (b) {
+                String s = bar();
+                if (s != null) return s;
+            }
+        }</selection>
+
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethod/ExtractBracedElseBranch_after.java

@@ -0,0 +1,25 @@
+import org.jetbrains.annotations.Nullable;
+
+class ElseIf {
+    String foo(boolean a, boolean b) {
+        if (a) {
+
+        } else {
+            String s = newMethod(b);
+            if (s != null) return s;
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private String newMethod(boolean b) {
+        if (b) {
+            String s = bar();
+            if (s != null) return s;
+        }
+        return null;
+    }
+
+    String bar() { return "";}
+}

java/java-tests/testData/refactoring/extractMethodObject/multipleExitPoints/AssignReturnValueToForeachParameter.java.after

@@ -20,8 +20,10 @@ class Foo {
   }
 
   public Foo getFoo(String name) {
-    if (name != null)
-        Inner inner = new Inner(name).invoke();if (inner.is()) return inner.getFoo();
+    if (name != null) {
+        Inner inner = new Inner(name).invoke();
+        if (inner.is()) return inner.getFoo();
+    }
 
 
       return null;

java/java-tests/testSrc/com/intellij/java/refactoring/ExtractMethodTest.java

@@ -1137,6 +1137,30 @@ public class ExtractMethodTest extends LightCodeInsightTestCase {
     doTest();
   }
 
+  public void testExtractBareThenBranch() throws Exception {
+    doTest();
+  }
+
+  public void testExtractBareElseBranch() throws Exception {
+    doTest();
+  }
+
+  public void testExtractBareForBody() throws Exception {
+    doTest();
+  }
+
+  public void testExtractBareDoWhileBody() throws Exception {
+    doTest();
+  }
+
+  public void testExtractBracedElseBranch() throws Exception {
+    doTest();
+  }
+
+  public void testExtractBracedDoWhileBody() throws Exception {
+    doTest();
+  }
+
   private void doTestDisabledParam() throws PrepareFailedException {
     final CodeStyleSettings settings = CodeStyleSettingsManager.getSettings(getProject());
     settings.ELSE_ON_NEW_LINE = true;