extract method object: process parameters as output variables (IDEA-162642)

java/java-impl/src/com/intellij/refactoring/extractMethodObject/ExtractMethodObjectProcessor.java

@@ -345,6 +345,17 @@ public class ExtractMethodObjectProcessor extends BaseRefactoringProcessor {
         }
       }
 
+      @Override
+      public void visitParameter(PsiParameter parameter) {
+        super.visitParameter(parameter);
+        final PsiElement declarationScope = parameter.getDeclarationScope();
+        for (PsiVariable variable : outputVariables) {
+          if (Comparing.strEqual(variable.getName(), parameter.getName())) {
+            replacementMap.put(parameter, myElementFactory.createStatementFromText(myInnerClassName + ".this." + var2FieldNames.get(variable.getName()) + " = " + parameter.getName() + ";", declarationScope));
+          }
+        }
+      }
+
       @Override
       public void visitReferenceExpression(final PsiReferenceExpression expression) {
         super.visitReferenceExpression(expression);
@@ -368,16 +379,37 @@ public class ExtractMethodObjectProcessor extends BaseRefactoringProcessor {
       }
     }
 
+    Map<PsiStatement, PsiForeachStatement> blocksToReplace = new LinkedHashMap<>();
     for (PsiElement statement : replacementMap.keySet()) {
       final PsiElement replacement = replacementMap.get(statement);
       if (replacement != null) {
-        if (statement instanceof PsiLocalVariable) {
+        if (statement instanceof PsiParameter) {
+          PsiCodeBlock codeBlock = null;
+          final PsiElement declarationScope = ((PsiParameter)statement).getDeclarationScope();
+          if (declarationScope instanceof PsiForeachStatement) {
+            final PsiStatement loopBody = ((PsiForeachStatement)declarationScope).getBody();
+            if (loopBody instanceof PsiBlockStatement) {
+              codeBlock = ((PsiBlockStatement)loopBody).getCodeBlock();
+            }
+            else {
+              blocksToReplace.put((PsiStatement)replacement, (PsiForeachStatement)declarationScope);
+            }
+          }
+          else if (declarationScope instanceof PsiCatchSection){
+            codeBlock = ((PsiCatchSection)declarationScope).getCatchBlock();
+          }
+          if (codeBlock != null) {
+            codeBlock.addBefore(replacement, codeBlock.getFirstBodyElement());
+          }
+        }
+        else if (statement instanceof PsiLocalVariable) {
           PsiLocalVariable variable = (PsiLocalVariable)statement;
           variable.normalizeDeclaration();
           PsiDeclarationStatement declaration = PsiTreeUtil.getParentOfType(statement, PsiDeclarationStatement.class);
           LOG.assertTrue(declaration != null);
           declaration.replace(replacement);
-        } else {
+        }
+        else {
           if (statement instanceof PsiReturnStatement) {
             final PsiExpression returnValue = ((PsiReturnStatement)statement).getReturnValue();
             if (!(returnValue instanceof PsiReferenceExpression || returnValue == null || returnValue instanceof PsiLiteralExpression)) {
@@ -392,6 +424,10 @@ public class ExtractMethodObjectProcessor extends BaseRefactoringProcessor {
       }
     }
 
+    for (PsiStatement statement : blocksToReplace.keySet()) {
+      RefactoringUtil.putStatementInLoopBody(statement, blocksToReplace.get(statement), null);
+    }
+
     myChangeReturnType = true;
   }
 

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

@@ -0,0 +1,34 @@
+
+import java.util.Arrays;
+import java.util.List;
+
+class Foo {
+  private List<Foo> foos = Arrays.asList(new Foo("one"), new Foo("two"));
+
+  private String name;
+
+  public Foo(String name) {
+    this.name = name;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public List<Foo> getFoos() {
+    return foos;
+  }
+
+  public Foo getFoo(String name) {
+    if (name != null)
+      <selection>for (Foo foo : getFoos())
+        if (foo.getName().equals(name)) {
+          return foo;
+        }</selection>
+
+
+
+    return null;
+  }
+
+}

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

@@ -0,0 +1,59 @@
+
+import java.util.Arrays;
+import java.util.List;
+
+class Foo {
+  private List<Foo> foos = Arrays.asList(new Foo("one"), new Foo("two"));
+
+  private String name;
+
+  public Foo(String name) {
+    this.name = name;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public List<Foo> getFoos() {
+    return foos;
+  }
+
+  public Foo getFoo(String name) {
+    if (name != null)
+        Inner inner = new Inner(name).invoke();if (inner.is()) return inner.getFoo();
+
+
+      return null;
+  }
+
+    private class Inner {
+        private boolean myResult;
+        private String name;
+        private Foo foo;
+
+        public Inner(String name) {
+            this.name = name;
+        }
+
+        boolean is() {
+            return myResult;
+        }
+
+        public Foo getFoo() {
+            return foo;
+        }
+
+        public Inner invoke() {
+            for (Foo foo : getFoos()) {
+                Inner.this.foo = foo;
+                if (foo.getName().equals(name)) {
+                    myResult = true;
+                    return this;
+                }
+            }
+            myResult = false;
+            return this;
+        }
+    }
+}

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

@@ -146,4 +146,8 @@ public class ExtractMethodObjectWithMultipleExitPointsTest extends LightRefactor
   public void testFormattingInside() throws Exception {
     doTest();
   }
+
+  public void testAssignReturnValueToForeachParameter() throws Exception {
+    doTest();
+  }
 }