RefactoringUtil#ensureCodeBlock enhanced and used in Surround with try-catch

Fixes IDEA-178781 "Surround with try-catch" QuickFix for "Unhandled exception" in a field initializer
Enables stream-to-loop in field initializer
Fixes stream-to-loop in for initializer
Disables stream-to-loop in for update

java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/SurroundWithTryCatchFix.java

@@ -25,10 +25,9 @@ import com.intellij.openapi.editor.ScrollType;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
-import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.refactoring.util.RefactoringChangeUtil;
 import com.intellij.refactoring.util.RefactoringUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.siyeh.ig.psiutils.ControlFlowUtils;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -38,15 +37,15 @@ import org.jetbrains.annotations.NotNull;
 public class SurroundWithTryCatchFix implements IntentionAction {
   private static final Logger LOG = Logger.getInstance("#com.intellij.codeInsight.daemon.impl.quickfix.SurroundWithTryCatchFix");
 
-  private PsiElement myStatement;
+  private PsiElement myElement;
 
   public SurroundWithTryCatchFix(@NotNull PsiElement element) {
-    final PsiFunctionalExpression functionalExpression = PsiTreeUtil.getParentOfType(element, PsiFunctionalExpression.class, false, PsiStatement.class);
-    if (functionalExpression == null) {
-      myStatement = PsiTreeUtil.getNonStrictParentOfType(element, PsiStatement.class);
-    }
-    else if (functionalExpression instanceof PsiLambdaExpression) {
-      myStatement = functionalExpression;
+    if (element instanceof PsiStatement ||
+        element instanceof PsiResourceVariable ||
+        (element instanceof PsiExpression &&
+         !(element instanceof PsiMethodReferenceExpression) &&
+         ControlFlowUtils.canExtractStatement((PsiExpression)element, false))) {
+      myElement = element;
     }
   }
 
@@ -64,10 +63,7 @@ public class SurroundWithTryCatchFix implements IntentionAction {
 
   @Override
   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
-    return myStatement != null &&
-           myStatement.isValid() &&
-           (!(myStatement instanceof PsiExpressionStatement) ||
-            !RefactoringChangeUtil.isSuperOrThisMethodCall(((PsiExpressionStatement)myStatement).getExpression()));
+    return myElement != null && myElement.isValid();
   }
 
   @Override
@@ -76,25 +72,17 @@ public class SurroundWithTryCatchFix implements IntentionAction {
     int line = editor.getCaretModel().getLogicalPosition().line;
     editor.getCaretModel().moveToLogicalPosition(new LogicalPosition(0, 0));
 
-    if (myStatement.getParent() instanceof PsiForStatement) {
-      PsiForStatement forStatement = (PsiForStatement)myStatement.getParent();
-      if (myStatement.equals(forStatement.getInitialization()) || myStatement.equals(forStatement.getUpdate())) {
-        myStatement = forStatement;
-      }
-    }
-
-    if (myStatement instanceof PsiLambdaExpression) {
-      final PsiCodeBlock body = RefactoringUtil.expandExpressionLambdaToCodeBlock(((PsiLambdaExpression)myStatement));
-      final PsiStatement[] statements = body.getStatements();
-      LOG.assertTrue(statements.length == 1);
-      myStatement = statements[0];
+    if (myElement instanceof PsiExpression) {
+      myElement = RefactoringUtil.ensureCodeBlock((PsiExpression)myElement);
     }
+    myElement = RefactoringUtil.getParentStatement(myElement, false);
+    if (myElement == null) return;
 
     TextRange range = null;
 
     try{
       JavaWithTryCatchSurrounder handler = new JavaWithTryCatchSurrounder();
-      range = handler.surroundElements(project, editor, new PsiElement[] {myStatement});
+      range = handler.surroundElements(project, editor, new PsiElement[]{myElement});
     }
     catch(IncorrectOperationException e){
       LOG.error(e);

java/java-impl/src/com/intellij/codeInspection/streamToLoop/StreamToLoopInspection.java

@@ -40,6 +40,7 @@ import com.intellij.psi.util.PsiUtil;
 import com.intellij.psi.util.RedundantCastUtil;
 import com.intellij.refactoring.util.RefactoringUtil;
 import com.intellij.util.ExceptionUtil;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.ig.psiutils.*;
 import one.util.streamex.IntStreamEx;
@@ -277,7 +278,7 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
       }
       TerminalOperation terminal = getTerminal(operations);
       if (terminal == null) return;
-      PsiStatement statement = PsiTreeUtil.getParentOfType(terminalCall, PsiStatement.class);
+      PsiStatement statement = ObjectUtils.tryCast(RefactoringUtil.getParentStatement(terminalCall, false), PsiStatement.class);
       LOG.assertTrue(statement != null);
       CommentTracker ct = new CommentTracker();
       try {

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

@@ -21,6 +21,7 @@ import com.intellij.codeInsight.daemon.impl.analysis.HighlightControlFlowUtil;
 import com.intellij.codeInsight.daemon.impl.analysis.JavaHighlightUtil;
 import com.intellij.codeInsight.highlighting.HighlightManager;
 import com.intellij.lang.java.JavaLanguage;
+import com.intellij.lang.jvm.JvmModifier;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.editor.RangeMarker;
@@ -49,9 +50,11 @@ import com.intellij.refactoring.PackageWrapper;
 import com.intellij.refactoring.introduceField.ElementToWorkOn;
 import com.intellij.refactoring.introduceVariable.IntroduceVariableBase;
 import com.intellij.util.IncorrectOperationException;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.HashMap;
 import com.intellij.util.containers.HashSet;
 import gnu.trove.THashMap;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -210,7 +213,6 @@ public class RefactoringUtil {
       parent = parent.getParent();
     }
     PsiElement parentStatement = parent;
-    parent = parentStatement instanceof PsiStatement ? parentStatement : parentStatement.getParent();
     while (parent instanceof PsiStatement) {
       if (!skipScopingStatements && ((parent instanceof PsiForStatement && parentStatement == ((PsiForStatement)parent).getBody()) || (
         parent instanceof PsiForeachStatement && parentStatement == ((PsiForeachStatement)parent).getBody()) || (
@@ -725,9 +727,10 @@ public class RefactoringUtil {
   }
 
   public static String getNewInnerClassName(PsiClass aClass, String oldInnerClassName, String newName) {
-    if (!oldInnerClassName.endsWith(aClass.getName())) return newName;
+    String className = aClass.getName();
+    if (className == null || !oldInnerClassName.endsWith(className)) return newName;
     StringBuilder buffer = new StringBuilder(oldInnerClassName);
-    buffer.replace(buffer.length() - aClass.getName().length(), buffer.length(), newName);
+    buffer.replace(buffer.length() - className.length(), buffer.length(), newName);
     return buffer.toString();
   }
 
@@ -738,7 +741,9 @@ public class RefactoringUtil {
     final PsiMethod[] constructors = subClass.getConstructors();
     if (constructors.length > 0) {
       for (PsiMethod constructor : constructors) {
-        final PsiStatement[] statements = constructor.getBody().getStatements();
+        PsiCodeBlock body = constructor.getBody();
+        if (body == null) continue;
+        final PsiStatement[] statements = body.getStatements();
         if (statements.length < 1 || !JavaHighlightUtil.isSuperOrThisCall(statements[0], true, true)) {
           implicitConstructorUsageVistor.visitConstructor(constructor, baseDefaultConstructor);
         }
@@ -971,6 +976,11 @@ public class RefactoringUtil {
    * {@link PsiIfStatement}, etc.
    * </p>
    *
+   * <p>
+   * This method works if {@link com.siyeh.ig.psiutils.ControlFlowUtils#canExtractStatement(PsiExpression)}
+   * returns true on the same expression.
+   * </p>
+   *
    * @param expression an expression which should be located inside the code block
    * @return a passed expression if it's already surrounded by code block and no changes are necessary;
    *         a replacement expression (which is equivalent to the passed expression) if a new code block was created;
@@ -979,18 +989,62 @@ public class RefactoringUtil {
   @Nullable
   public static <T extends PsiExpression> T ensureCodeBlock(@NotNull T expression) {
     PsiElement parent = getParentStatement(expression, false);
-    if (parent == null) return null;
-    if (parent instanceof PsiStatement && parent.getParent() instanceof PsiCodeBlock) return expression;
+    if (parent == null) {
+      parent = PsiTreeUtil.getParentOfType(expression, PsiField.class, true, PsiClass.class);
+      if (parent == null) return null;
+    }
+    PsiElement grandParent = parent.getParent();
+    if (parent instanceof PsiStatement && grandParent instanceof PsiCodeBlock) {
+      if (!(parent instanceof PsiForStatement) ||
+          !PsiTreeUtil.isAncestor(((PsiForStatement)parent).getInitialization(), expression, true) ||
+          !hasNameCollision(((PsiForStatement)parent).getInitialization(), grandParent)) {
+        return expression;
+      }
+    }
     Object marker = new Object();
     PsiTreeUtil.mark(expression, marker);
     PsiElement copy = parent.copy();
     PsiElement newParent;
     PsiElementFactory factory = JavaPsiFacade.getElementFactory(expression.getProject());
     if (parent instanceof PsiExpression) {
-      PsiLambdaExpression lambda = (PsiLambdaExpression)parent.getParent();
+      PsiLambdaExpression lambda = (PsiLambdaExpression)grandParent;
       String replacement = PsiType.VOID.equals(LambdaUtil.getFunctionalInterfaceReturnType(lambda)) ? "{a;}" : "{return a;}";
       PsiElement block = parent.replace(factory.createCodeBlockFromText(replacement, lambda));
       newParent = LambdaUtil.extractSingleExpressionFromBody(block).replace(copy);
+    }
+    else if (parent instanceof PsiField) {
+      PsiField field = (PsiField)parent;
+      PsiClassInitializer initializer =
+        ObjectUtils.tryCast(PsiTreeUtil.skipWhitespacesAndCommentsForward(field), PsiClassInitializer.class);
+      boolean isStatic = field.hasModifier(JvmModifier.STATIC);
+      if (initializer == null || initializer.hasModifier(JvmModifier.STATIC) != isStatic) {
+        initializer = factory.createClassInitializer();
+        if (isStatic) {
+          Objects.requireNonNull(initializer.getModifierList()).setModifierProperty(PsiModifier.STATIC, true);
+        }
+        initializer = (PsiClassInitializer)field.getParent().addAfter(initializer, field);
+      }
+      PsiCodeBlock body = initializer.getBody();
+      // There are at least two children: open and close brace
+      // we will insert an initializer after the first brace and any whitespace which follow it
+      PsiElement anchor = PsiTreeUtil.skipWhitespacesForward(body.getFirstChild());
+      assert anchor != null;
+      anchor = anchor.getPrevSibling();
+      assert anchor != null;
+
+      PsiExpressionStatement assignment =
+        (PsiExpressionStatement)factory.createStatementFromText(field.getName() + "=null;", initializer);
+      assignment = (PsiExpressionStatement)body.addAfter(assignment, anchor);
+      PsiExpression fieldInitializer = ((PsiField)copy).getInitializer();
+      if (fieldInitializer instanceof PsiArrayInitializerExpression) {
+        fieldInitializer = createNewExpressionFromArrayInitializer((PsiArrayInitializerExpression)fieldInitializer, field.getType());
+      }
+      PsiExpression rExpression = ((PsiAssignmentExpression)assignment.getExpression()).getRExpression();
+      assert fieldInitializer != null;
+      assert rExpression != null;
+      rExpression.replace(fieldInitializer);
+      Objects.requireNonNull(field.getInitializer()).delete();
+      newParent = assignment;
     } else {
       PsiBlockStatement blockStatement = (PsiBlockStatement)parent.replace(factory.createStatementFromText("{}", parent));
       newParent = blockStatement.getCodeBlock().add(copy);
@@ -999,6 +1053,18 @@ public class RefactoringUtil {
     return (T)PsiTreeUtil.releaseMark(newParent, marker);
   }
 
+  private static boolean hasNameCollision(PsiElement declaration, PsiElement context) {
+    if (declaration instanceof PsiDeclarationStatement) {
+      PsiResolveHelper helper = JavaPsiFacade.getInstance(context.getProject()).getResolveHelper();
+      return StreamEx.of(((PsiDeclarationStatement)declaration).getDeclaredElements())
+            .select(PsiLocalVariable.class)
+            .map(PsiLocalVariable::getName)
+            .nonNull()
+            .anyMatch(name -> helper.resolveReferencedVariable(name, context) != null);
+    }
+    return false;
+  }
+
   public interface ImplicitConstructorUsageVisitor {
     void visitConstructor(PsiMethod constructor, PsiMethod baseConstructor);
 

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterFieldInitializer.java

@@ -0,0 +1,62 @@
+// "Fix all 'Stream API call chain can be replaced with loop' problems in file" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+    static final List<String> STR;
+
+    static {
+        List<String> list = new ArrayList<>();
+        for (String s : Arrays.asList("foo", "bar", "baz")) {
+            String toUpperCase = s.toUpperCase();
+            list.add(toUpperCase);
+        }
+        STR = list;
+    }
+
+    final String field;
+
+    {
+        StringJoiner joiner = new StringJoiner(",");
+        for (String s : STR) {
+            if (!s.isEmpty()) {
+                joiner.add(s);
+            }
+        }
+        field = joiner.toString();
+    }
+
+    static {
+        System.out.println("static initializer already exists");
+    }
+
+    final long count;
+
+    {
+        long result = 0L;
+        for (Integer i : Arrays.asList(1, 2, 3, 4)) {
+            if (i % 2 == 0) {
+                result++;
+            }
+        }
+        count = result;
+        System.out.println("initializer already exists");
+    }
+
+    final long x = 0, count2 = Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count();
+
+    final long count3 = Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count(), y = 0;
+
+    final long[] countArray;
+
+    {
+        long result = 0L;
+        for (Integer i : Arrays.asList(1, 2, 3, 4)) {
+            if (i % 2 == 0) {
+                result++;
+            }
+        }
+        countArray = new long[]{result};
+    }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterStreamInFor.java

@@ -0,0 +1,50 @@
+// "Fix all 'Stream API call chain can be replaced with loop' problems in file" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+    String j = "foo";
+
+    public void test(List<String> list) {
+        long i = 0L;
+        for (String s : list) {
+            if (s.isEmpty()) {
+                i++;
+            }
+        }
+        for(;
+            i<10;
+            i+=list.stream().filter(String::isEmpty).count()) {
+            System.out.println(i);
+        }
+
+        {
+            long j = 0L;
+            for (String s : list) {
+                if (s.isEmpty()) {
+                    j++;
+                }
+            }
+            for(;
+                j<10;
+                j+=list.stream().filter(String::isEmpty).count()) {
+                System.out.println(j);
+            }
+        }
+
+        System.out.println(j);
+
+        StringJoiner joiner = new StringJoiner(",");
+        for (String s1 : list) {
+            joiner.add(s1);
+        }
+        for(String s = joiner.toString(); !s.isEmpty(); s = s.substring(1)) {
+            System.out.println(s);
+        }
+    }
+
+    public static void main(String[] args) {
+        new Main().test(Arrays.asList("", "", "foo"));
+    }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeFieldInitializer.java

@@ -0,0 +1,26 @@
+// "Fix all 'Stream API call chain can be replaced with loop' problems in file" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+    static final List<String> STR = Stream.of("foo","bar","baz").ma<caret>p(String::toUpperCase).collect(Collectors.toList());
+
+    final String field = STR.stream().filter(x -> !x.isEmpty()).collect(Collectors.joining(","));
+
+    static {
+        System.out.println("static initializer already exists");
+    }
+
+    final long count = Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count();
+
+    {
+        System.out.println("initializer already exists");
+    }
+
+    final long x = 0, count2 = Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count();
+
+    final long count3 = Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count(), y = 0;
+
+    final long[] countArray = {Stream.of(1,2,3,4).filter(i -> i % 2 == 0).count()};
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeStreamInFor.java

@@ -0,0 +1,32 @@
+// "Fix all 'Stream API call chain can be replaced with loop' problems in file" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+    String j = "foo";
+
+    public void test(List<String> list) {
+        for(long i = list.stream().filter(String::isEmpty).cou<caret>nt();
+            i<10;
+            i+=list.stream().filter(String::isEmpty).count()) {
+            System.out.println(i);
+        }
+
+        for(long j = list.stream().filter(String::isEmpty).count();
+            j<10;
+            j+=list.stream().filter(String::isEmpty).count()) {
+            System.out.println(j);
+        }
+
+        System.out.println(j);
+
+        for(String s = list.stream().collect(Collectors.joining(",")); !s.isEmpty(); s = s.substring(1)) {
+            System.out.println(s);
+        }
+    }
+
+    public static void main(String[] args) {
+        new Main().test(Arrays.asList("", "", "foo"));
+    }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/surroundWithTry/afterField.java

@@ -0,0 +1,19 @@
+// "Surround with try/catch" "true"
+import java.io.IOException;
+
+class C {
+    static final String S;
+
+    static {
+        try {
+            S = getString();
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    static String getString() throws IOException {
+        if(Math.random() > 0.5) throw new IOException();
+        return "foo";
+    }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/surroundWithTry/beforeField.java

@@ -0,0 +1,11 @@
+// "Surround with try/catch" "true"
+import java.io.IOException;
+
+class C {
+    static final String S = getSt<caret>ring();
+
+    static String getString() throws IOException {
+        if(Math.random() > 0.5) throw new IOException();
+        return "foo";
+    }
+}

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ControlFlowUtils.java

@@ -881,16 +881,26 @@ public class ControlFlowUtils {
 
   /**
    * @param expression expression to check
-   * @return true if given expression can be converted to statement without semantics change
+   * @return true if given expression is always executed and can be converted to a statement
    */
   public static boolean canExtractStatement(PsiExpression expression) {
+    return canExtractStatement(expression, true);
+  }
+
+  /**
+   * @param expression    expression to check
+   * @param checkExecuted if true, expression will be considered non-extractable if it is not always executed within its topmost expression
+   *                      (e.g. appears in then/else branches in ?: expression)
+   * @return true if given expression can be converted to a statement
+   */
+  public static boolean canExtractStatement(PsiExpression expression, boolean checkExecuted) {
     PsiElement cur = expression;
     PsiElement parent = cur.getParent();
     while(parent instanceof PsiExpression || parent instanceof PsiExpressionList) {
       if(parent instanceof PsiLambdaExpression) {
         return true;
       }
-      if(parent instanceof PsiPolyadicExpression) {
+      if (checkExecuted && parent instanceof PsiPolyadicExpression) {
         PsiPolyadicExpression polyadicExpression = (PsiPolyadicExpression)parent;
         IElementType type = polyadicExpression.getOperationTokenType();
         if ((type.equals(JavaTokenType.ANDAND) || type.equals(JavaTokenType.OROR)) && polyadicExpression.getOperands()[0] != cur) {
@@ -898,7 +908,7 @@ public class ControlFlowUtils {
           return false;
         }
       }
-      if(parent instanceof PsiConditionalExpression && ((PsiConditionalExpression)parent).getCondition() != cur) {
+      if (checkExecuted && parent instanceof PsiConditionalExpression && ((PsiConditionalExpression)parent).getCondition() != cur) {
         return false;
       }
       if(parent instanceof PsiMethodCallExpression) {
@@ -910,6 +920,12 @@ public class ControlFlowUtils {
       cur = parent;
       parent = cur.getParent();
     }
+    if (parent instanceof PsiStatement) {
+      PsiElement grandParent = parent.getParent();
+      if (checkExecuted && grandParent instanceof PsiForStatement && ((PsiForStatement)grandParent).getUpdate() == parent) {
+        return false;
+      }
+    }
     if(parent instanceof PsiReturnStatement || parent instanceof PsiExpressionStatement) return true;
     if(parent instanceof PsiLocalVariable) {
       PsiElement grandParent = parent.getParent();
@@ -917,6 +933,13 @@ public class ControlFlowUtils {
         return true;
       }
     }
+    if (parent instanceof PsiField) {
+      PsiElement prev = PsiTreeUtil.skipWhitespacesAndCommentsBackward(parent);
+      PsiElement next = PsiTreeUtil.skipWhitespacesAndCommentsForward(parent);
+      boolean multipleFieldsDeclaration = prev instanceof PsiJavaToken && ((PsiJavaToken)prev).getTokenType() == JavaTokenType.COMMA ||
+                                          next instanceof PsiJavaToken && ((PsiJavaToken)next).getTokenType() == JavaTokenType.COMMA;
+      return !multipleFieldsDeclaration;
+    }
     if(parent instanceof PsiForeachStatement && ((PsiForeachStatement)parent).getIteratedValue() == cur) return true;
     if(parent instanceof PsiIfStatement && ((PsiIfStatement)parent).getCondition() == cur) return true;
     return false;