Inline method: automatically convert to "continue"; automatically transform method body to single return when possible

Fixes IDEA-37432 "Inline Method refactoring is not supported when return statement interrupts the execution flow" is wrong
Fixes IDEA-158665 Support Inline Method refactoring when return statement interrupts the execution flow
Fixes IDEA-180007 Inline method with returns should work when inlining point is the only expression in a loop or lambda body

java/java-impl/src/com/intellij/codeInsight/intention/impl/singlereturn/ConvertToSingleReturnAction.java

@@ -15,13 +15,14 @@ import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiTypesUtil;
+import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import static com.intellij.util.ObjectUtils.tryCast;
@@ -49,7 +50,7 @@ public class ConvertToSingleReturnAction extends PsiElementBaseIntentionAction {
     PsiType returnType = PsiTypesUtil.getMethodReturnType(block);
     if (returnType == null) return null;
 
-    List<PsiReturnStatement> returns = findReturns(block);
+    List<PsiReturnStatement> returns = Arrays.asList(PsiUtil.findReturnStatements(block));
     indicator.checkCanceled();
     indicator.setFraction(0.1);
     FinishMarker marker = FinishMarker.defineFinishMarker(block, returnType, returns);
@@ -62,12 +63,12 @@ public class ConvertToSingleReturnAction extends PsiElementBaseIntentionAction {
     return copy;
   }
 
-  private static void convertReturns(@NotNull Project project,
-                                     PsiCodeBlock block,
-                                     PsiType returnType,
-                                     FinishMarker marker,
-                                     int count,
-                                     ProgressIndicator indicator) {
+  public static PsiLocalVariable convertReturns(@NotNull Project project,
+                                                PsiCodeBlock block,
+                                                PsiType returnType,
+                                                FinishMarker marker,
+                                                int count,
+                                                ProgressIndicator indicator) {
     ExitContext exitContext = new ExitContext(block, returnType, marker);
     int i=0;
 
@@ -80,8 +81,9 @@ public class ConvertToSingleReturnAction extends PsiElementBaseIntentionAction {
       ReturnReplacementContext.replaceSingleReturn(project, block, exitContext, returnStatement);
     }
     indicator.setFraction(0.9);
-    exitContext.declareVariables();
+    PsiLocalVariable resultVariable = exitContext.declareVariables();
     indicator.setFraction(0.92);
+    return resultVariable;
   }
 
   @Override
@@ -139,28 +141,6 @@ public class ConvertToSingleReturnAction extends PsiElementBaseIntentionAction {
     return visitor.myReturnStatement;
   }
 
-  @NotNull
-  private static List<PsiReturnStatement> findReturns(PsiCodeBlock block) {
-    List<PsiReturnStatement> result = new ArrayList<>();
-    block.accept(new JavaRecursiveElementWalkingVisitor() {
-      @Override
-      public void visitReturnStatement(PsiReturnStatement statement) {
-        super.visitReturnStatement(statement);
-        result.add(statement);
-      }
-
-      @Override
-      public void visitExpression(PsiExpression expression) {}
-
-      @Override
-      public void visitLambdaExpression(PsiLambdaExpression expression) {}
-
-      @Override
-      public void visitClass(PsiClass aClass) {}
-    });
-    return result;
-  }
-
   @Nls(capitalization = Nls.Capitalization.Sentence)
   @NotNull
   @Override

java/java-impl/src/com/intellij/codeInsight/intention/impl/singlereturn/ExitContext.java

@@ -3,6 +3,7 @@ package com.intellij.codeInsight.intention.impl.singlereturn;
 
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.VariableKind;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiTypesUtil;
 import com.intellij.refactoring.util.RefactoringUtil;
 import com.intellij.util.containers.ContainerUtil;
@@ -35,7 +36,7 @@ class ExitContext {
     myReturnType = returnType;
     myReturnVariable =
       new VariableNameGenerator(block, VariableKind.LOCAL_VARIABLE).byName("result", "res").byType(returnType).generate(true);
-    if (marker.myDefaultValue != null && marker.myDefaultValue.isPhysical()) {
+    if (marker.myDefaultValue != null && PsiTreeUtil.isAncestor(block, marker.myDefaultValue, true)) {
       myReturnVariableDefaultValue = (PsiExpression)marker.myDefaultValue.copy();
     } else {
       myReturnVariableDefaultValue = marker.myDefaultValue;
@@ -96,7 +97,7 @@ class ExitContext {
     }
   }
 
-  void declareVariables() {
+  PsiLocalVariable declareVariables() {
     if (myFinishedVariable != null) {
       PsiJavaToken start = requireNonNull(myBlock.getLBrace());
       PsiExpression initializer = myFactory.createExpressionFromText("false", null);
@@ -123,7 +124,9 @@ class ExitContext {
       }
       PsiJavaToken end = requireNonNull(myBlock.getRBrace());
       myBlock.addBefore(myFactory.createStatementFromText("return " + myReturnVariable + ";", myBlock), end);
+      return var;
     }
+    return null;
   }
 
   public boolean isFinishCondition(PsiStatement statement) {

java/java-impl/src/com/intellij/codeInsight/intention/impl/singlereturn/FinishMarker.java

@@ -27,7 +27,7 @@ import static java.util.Objects.requireNonNull;
 /**
  * Represents a way to indicate whether method execution is already finished
  */
-class FinishMarker {
+public class FinishMarker {
   /**
    * Type of finish marker
    */
@@ -48,7 +48,7 @@ class FinishMarker {
    * @param returns list of all method returns
    * @return a FinishMarker which is suitable for given method
    */
-  static FinishMarker defineFinishMarker(@NotNull PsiCodeBlock block, @NotNull PsiType returnType, List<PsiReturnStatement> returns) {
+  public static FinishMarker defineFinishMarker(@NotNull PsiCodeBlock block, @NotNull PsiType returnType, List<PsiReturnStatement> returns) {
     boolean mayNeedMarker = mayNeedMarker(returns, block);
     return defineFinishMarker(block, returns, returnType, mayNeedMarker, JavaPsiFacade.getElementFactory(block.getProject()));
   }

java/java-impl/src/com/intellij/refactoring/inline/InlineMethodHandler.java

@@ -60,7 +60,7 @@ class InlineMethodHandler extends JavaInlineActionHandler {
     }
     boolean allowInlineThisOnly = false;
     if (InlineMethodProcessor.checkBadReturns(method) && !InlineUtil.allUsagesAreTailCalls(method)) {
-      if (reference != null && InlineUtil.getTailCallType(reference) != InlineUtil.TailCallType.None) {
+      if (reference != null) {
         allowInlineThisOnly = true;
       }
       else {

java/java-impl/src/com/intellij/refactoring/inline/InlineMethodProcessor.java

@@ -806,7 +806,6 @@ public class InlineMethodProcessor extends BaseRefactoringProcessor {
 
     PsiType returnType = callSubstitutor.substitute(myMethod.getReturnType());
     InlineTransformer transformer = InlineTransformer.getSuitableTransformer(myMethod, ref);
-    assert transformer != null;
 
     PsiLocalVariable[] parmVars = declareParameters(block, argumentList, callSubstitutor);
 

java/java-impl/src/com/intellij/refactoring/inline/InlineTransformer.java

@@ -2,6 +2,9 @@
 package com.intellij.refactoring.inline;
 
 import com.intellij.codeInsight.BlockUtils;
+import com.intellij.codeInsight.intention.impl.singlereturn.ConvertToSingleReturnAction;
+import com.intellij.codeInsight.intention.impl.singlereturn.FinishMarker;
+import com.intellij.openapi.progress.EmptyProgressIndicator;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
@@ -14,7 +17,9 @@ import com.siyeh.ig.psiutils.CommentTracker;
 import com.siyeh.ig.psiutils.SideEffectChecker;
 import com.siyeh.ig.psiutils.StatementExtractor;
 import com.siyeh.ig.psiutils.VariableNameGenerator;
+import org.jetbrains.annotations.NotNull;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 
@@ -76,6 +81,32 @@ public abstract class InlineTransformer {
       return null;
     }
   }
+
+  static class LoopContinueTransformer extends InlineTransformer {
+    @Override
+    public boolean isMethodAccepted(PsiMethod method) {
+      for (PsiReturnStatement statement : PsiUtil.findReturnStatements(method)) {
+        if (PsiTreeUtil.getParentOfType(statement, PsiLoopStatement.class, true, PsiMethod.class) != null) {
+          // We cannot use "continue" without introducing a label if any of returns is inside nested loop.
+          // Introducing a label is ugly, so let's move to the next transformer 
+          return false;
+        }
+      }
+      return true;
+    }
+
+    @Override
+    public boolean isReferenceAccepted(PsiReference reference) {
+      return InlineUtil.getTailCallType(reference) == InlineUtil.TailCallType.Continue;
+    }
+
+    @Override
+    public PsiLocalVariable transformBody(PsiMethod methodCopy, PsiReference reference, PsiType returnType) {
+      extractReturnValues(methodCopy, true);
+      return null;
+    }
+  }
+  
   
   static class SimpleTailCallTransformer extends InlineTransformer {
     @Override
@@ -90,10 +121,16 @@ public abstract class InlineTransformer {
 
     @Override
     public PsiLocalVariable transformBody(PsiMethod methodCopy, PsiReference reference, PsiType returnType) {
-      PsiReturnStatement[] returnStatements = PsiUtil.findReturnStatements(methodCopy);
-      for (PsiReturnStatement returnStatement : returnStatements) {
-        final PsiExpression returnValue = returnStatement.getReturnValue();
-        if (returnValue == null) continue;
+      extractReturnValues(methodCopy, false);
+      return null;
+    }
+  }
+
+  private static void extractReturnValues(PsiMethod methodCopy, boolean replaceWithContinue) {
+    PsiReturnStatement[] returnStatements = PsiUtil.findReturnStatements(methodCopy);
+    for (PsiReturnStatement returnStatement : returnStatements) {
+      final PsiExpression returnValue = returnStatement.getReturnValue();
+      if (returnValue != null) {
         List<PsiExpression> sideEffects = SideEffectChecker.extractSideEffectExpressions(returnValue);
         CommentTracker ct = new CommentTracker();
         sideEffects.forEach(ct::markUnchanged);
@@ -106,7 +143,30 @@ public abstract class InlineTransformer {
         }
         ct.insertCommentsBefore(returnStatement);
       }
-      return null;
+      if (replaceWithContinue) {
+        new CommentTracker().replaceAndRestoreComments(returnStatement, "continue;");
+      }
+    }
+  }
+
+  private static class ConvertToSingleReturnTransformer extends InlineTransformer {
+    @Override
+    public boolean isMethodAccepted(PsiMethod method) {
+      return true;
+    }
+
+    @Override
+    public boolean isReferenceAccepted(PsiReference reference) {
+      return true;
+    }
+
+    @Override
+    public PsiLocalVariable transformBody(PsiMethod methodCopy, PsiReference reference, PsiType returnType) {
+      PsiCodeBlock block = Objects.requireNonNull(methodCopy.getBody());
+      List<PsiReturnStatement> returns = Arrays.asList(PsiUtil.findReturnStatements(block));
+      FinishMarker marker = FinishMarker.defineFinishMarker(block, returnType, returns);
+      return ConvertToSingleReturnAction.convertReturns(methodCopy.getProject(), block, returnType, marker, returns.size(),
+                                                        new EmptyProgressIndicator());
     }
   }
 
@@ -114,16 +174,19 @@ public abstract class InlineTransformer {
     return ContainerUtil.immutableList(
       new TailCallTransformer(),
       new SimpleTailCallTransformer(),
-      new NormalTransformer()
+      new NormalTransformer(),
+      new LoopContinueTransformer(),
+      new ConvertToSingleReturnTransformer()
     );
   }
-  
+
+  @NotNull
   static InlineTransformer getSuitableTransformer(PsiMethod method, PsiReference reference) {
     for (InlineTransformer transformer : getTransformers()) {
       if (transformer.isMethodAccepted(method) && transformer.isReferenceAccepted(reference)) {
         return transformer;
       }
     }
-    return null;
+    throw new InternalError("Transformer is unavailable");
   }
 }

java/java-impl/src/com/intellij/refactoring/util/InlineUtil.java

@@ -340,6 +340,7 @@ public class InlineUtil {
           if (blockParent instanceof PsiMethod || blockParent instanceof PsiLambdaExpression) return TailCallType.Simple;
           if (!(blockParent instanceof PsiBlockStatement)) return TailCallType.None;
           parent = blockParent.getParent();
+          if (parent instanceof PsiLoopStatement) return TailCallType.Continue;
         }
         if (!(parent instanceof PsiLabeledStatement) && !(parent instanceof PsiIfStatement)) return TailCallType.None;
         curElement = (PsiStatement)parent;
@@ -512,6 +513,6 @@ public class InlineUtil {
   }
 
   public enum TailCallType {
-    None, Simple, Return
+    None, Simple, Continue, Return
   }
 }

java/java-tests/testData/refactoring/inlineMethod/ForContinue.java

@@ -0,0 +1,21 @@
+import java.util.*;
+
+class AAA {
+    private void foo(List<String> list) {
+        for (String val : list) {
+            <caret>checkVal(val);
+        }
+    }
+
+    private void checkVal(String message) {
+        if (message == null) return;
+        message = message.trim();
+        if (message.isEmpty()) return;
+        try {
+            Integer.parseInt(message);
+        } catch (NumberFormatException e) {
+            return;
+        }
+        throw new IllegalArgumentException("Should not be a number!");
+    }
+}

java/java-tests/testData/refactoring/inlineMethod/ForContinue.java.after

@@ -0,0 +1,18 @@
+import java.util.*;
+
+class AAA {
+    private void foo(List<String> list) {
+        for (String val : list) {
+            if (val == null) continue;
+            val = val.trim();
+            if (val.isEmpty()) continue;
+            try {
+                Integer.parseInt(val);
+            } catch (NumberFormatException e) {
+                continue;
+            }
+            throw new IllegalArgumentException("Should not be a number!");
+        }
+    }
+
+}

java/java-tests/testData/refactoring/inlineMethod/NotTailCallInsideIf.java.after

@@ -0,0 +1,13 @@
+class A {
+
+    void bar(int x) {
+    if (x > 0) {
+        if(Math.random() > 2) {
+          System.out.println("xyz");
+        } else {
+            System.out.println("oops");
+        }
+    }
+    System.out.println("x < 0");
+  }
+}

java/java-tests/testData/refactoring/inlineMethod/SingleReturn1.java

@@ -0,0 +1,14 @@
+class Tester {
+    // IDEA-37432
+    String callee(String x) {
+        if (x == null) {
+            return null;
+        }
+        return x;
+    }
+
+    void caller(String v) {
+        String g = <caret>callee(v);
+        System.out.println(g);
+    }
+}

java/java-tests/testData/refactoring/inlineMethod/SingleReturn1.java.after

@@ -0,0 +1,11 @@
+class Tester {
+
+    void caller(String v) {
+        String result = null;
+        if (v != null) {
+            result = v;
+        }
+        String g = result;
+        System.out.println(g);
+    }
+}

java/java-tests/testData/refactoring/inlineMethod/SingleReturn2.java

@@ -0,0 +1,17 @@
+class Tester {
+    boolean a;
+
+    native void doB();
+
+    // IDEA-158665
+    void inlinedMethod() {
+        if (a)
+            return;
+        doB();
+    }
+
+    void useInlinedMethod() {
+        <caret>inlinedMethod();
+        System.out.println("ok");
+    }
+}

java/java-tests/testData/refactoring/inlineMethod/SingleReturn2.java.after

@@ -0,0 +1,12 @@
+class Tester {
+    boolean a;
+
+    native void doB();
+
+    void useInlinedMethod() {
+        if (!a) {
+            doB();
+        }
+        System.out.println("ok");
+    }
+}

java/java-tests/testSrc/com/intellij/java/refactoring/inline/InlineMethodTest.java

@@ -248,6 +248,18 @@ public class InlineMethodTest extends LightRefactoringTestCase {
   public void testNotAStatement4() {
     doTest();
   }
+  
+  public void testForContinue() {
+    doTest();
+  }
+  
+  public void testSingleReturn1() {
+    doTestAssertBadReturn();
+  }
+  
+  public void testSingleReturn2() {
+    doTestAssertBadReturn();
+  }
 
   public void testInSuperCall() {
     doTestConflict("Inline cannot be applied to multiline method in constructor call");
@@ -476,6 +488,7 @@ public class InlineMethodTest extends LightRefactoringTestCase {
   private void doTestAssertBadReturn() {
     @NonNls String fileName = configure();
     performAction(new MockInlineMethodOptions(), false, true);
+    checkResultByFile(fileName + ".after");
   }
 
   @NotNull
@@ -507,10 +520,10 @@ public class InlineMethodTest extends LightRefactoringTestCase {
       assertTrue("Bad returns not found", condition);
     } else {
       assertFalse("Bad returns found", condition);
-      final InlineMethodProcessor processor =
-        new InlineMethodProcessor(getProject(), method, refExpr, myEditor, options.isInlineThisOnly(), nonCode, nonCode,
-                                  !options.isKeepTheDeclaration());
-      processor.run();
     }
+    final InlineMethodProcessor processor =
+      new InlineMethodProcessor(getProject(), method, refExpr, myEditor, options.isInlineThisOnly(), nonCode, nonCode,
+                                !options.isKeepTheDeclaration());
+    processor.run();
   }
 }