IDEA-231237 Loop-to-stream conversion produces incorrect code if pattern variable is introduced

Also: pattern variable can be deleted via simple .delete()

GitOrigin-RevId: d3523f64d86de28d90b82d7bd98c1fd70567652d

java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/PostHighlightingVisitor.java

@@ -397,7 +397,7 @@ class PostHighlightingVisitor {
     else if (parameter instanceof PsiPatternVariable) {
       HighlightInfo highlightInfo = checkUnusedParameter(parameter, identifier);
       if (highlightInfo != null) {
-        QuickFixAction.registerQuickFixAction(highlightInfo, QuickFixFactory.getInstance().createRemoveUnusedVariableFix(parameter));
+        QuickFixAction.registerQuickFixAction(highlightInfo, QuickFixFactory.getInstance().createDeleteFix(parameter));
         return highlightInfo;
       }
     }

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

@@ -72,23 +72,7 @@ public class RemoveUnusedVariableFix implements IntentionAction {
   @Override
   public void invoke(@NotNull Project project, Editor editor, PsiFile file) {
     if (!FileModificationService.getInstance().prepareFileForWrite(myVariable.getContainingFile())) return;
-    if (myVariable instanceof PsiPatternVariable) {
-      removePatternVariable((PsiPatternVariable)myVariable);
-    } else {
-      removeVariableAndReferencingStatements(editor);
-    }
-  }
-
-  private static void removePatternVariable(PsiPatternVariable variable) {
-    Runnable action = () -> {
-      PsiPattern pattern = variable.getPattern();
-      if (pattern instanceof PsiTypeTestPattern) {
-        variable.replace(variable.getTypeElement());
-        return;
-      }
-      throw new UnsupportedOperationException("Unable to remove pattern variable " + variable.getName());
-    };
-    ApplicationManager.getApplication().runWriteAction(action);
+    removeVariableAndReferencingStatements(editor);
   }
 
   private void removeVariableAndReferencingStatements(Editor editor) {

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

@@ -51,7 +51,7 @@ public class SplitFilterAction extends PsiElementBaseIntentionAction {
     }
     for (PsiElement child = expression.getFirstChild(); child != token; child = child.getNextSibling()) {
       if (child instanceof PsiExpression) {
-        for (PsiPatternVariable variable : JavaPsiPatternUtil.getPatternVariablesVisibleOutsideOf((PsiExpression)child)) {
+        for (PsiPatternVariable variable : JavaPsiPatternUtil.getExposedPatternVariables((PsiExpression)child)) {
           for (PsiExpression operand : afterOperands) {
             if (VariableAccessUtils.variableIsUsed(variable, operand)) return true;
           }

java/java-impl/src/com/intellij/codeInspection/streamMigration/MigrateToStreamFix.java

@@ -18,13 +18,16 @@ package com.intellij.codeInspection.streamMigration;
 import com.intellij.codeInspection.*;
 import com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.StreamSource;
 import com.intellij.openapi.project.Project;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiLoopStatement;
-import com.intellij.psi.PsiStatement;
+import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.JavaCodeStyleManager;
+import com.intellij.psi.util.JavaPsiPatternUtil;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.siyeh.ig.psiutils.VariableAccessUtils;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.List;
+
 import static com.intellij.util.ObjectUtils.tryCast;
 
 class MigrateToStreamFix implements LocalQuickFix {
@@ -66,6 +69,24 @@ class MigrateToStreamFix implements LocalQuickFix {
     LambdaCanBeMethodReferenceInspection.replaceAllLambdasWithMethodReferences(result);
     RemoveRedundantTypeArgumentsUtil.removeRedundantTypeArguments(result);
     result = SimplifyStreamApiCallChainsInspection.simplifyStreamExpressions(result, true);
+    removeRedundantPatternVariables(result);
     JavaCodeStyleManager.getInstance(project).shortenClassReferences(result);
   }
+
+  private static void removeRedundantPatternVariables(PsiElement element) {
+    for (PsiLambdaExpression lambda : PsiTreeUtil.collectElementsOfType(element, PsiLambdaExpression.class)) {
+      PsiElement body = lambda.getBody();
+      if (body instanceof PsiExpression) {
+        PsiExpression expression = (PsiExpression)body;
+        if (PsiType.BOOLEAN.equals(expression.getType())) {
+          List<PsiPatternVariable> variables = JavaPsiPatternUtil.getExposedPatternVariablesIgnoreParent(expression);
+          for (PsiPatternVariable variable : variables) {
+            if (!VariableAccessUtils.variableIsUsed(variable, expression)) {
+              variable.delete();
+            }
+          }
+        }
+      }
+    }
+  }
 }

java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java

@@ -7,19 +7,19 @@ import com.intellij.psi.controlFlow.ControlFlow;
 import com.intellij.psi.controlFlow.ControlFlowUtil;
 import com.intellij.psi.search.LocalSearchScope;
 import com.intellij.psi.search.searches.ReferencesSearch;
+import com.intellij.psi.util.JavaPsiPatternUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.ArrayUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.IntArrayList;
 import com.siyeh.ig.psiutils.*;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.NoSuchElementException;
-import java.util.Objects;
+import java.util.*;
+import java.util.stream.Stream;
 
 import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.*;
 import static com.intellij.util.ObjectUtils.tryCast;
@@ -105,10 +105,11 @@ class TerminalBlock {
     PsiStatement single = getSingleStatement();
     if (single instanceof PsiIfStatement) {
       PsiIfStatement ifStatement = (PsiIfStatement)single;
-      if(ifStatement.getElseBranch() == null && ifStatement.getCondition() != null) {
+      PsiExpression condition = ifStatement.getCondition();
+      if(ifStatement.getElseBranch() == null && condition != null) {
         PsiStatement thenBranch = ifStatement.getThenBranch();
         if(thenBranch != null) {
-          return new TerminalBlock(this, new FilterOp(ifStatement.getCondition(), myVariable, false), myVariable, thenBranch);
+          return fromCondition(condition, false, thenBranch);
         }
       }
     }
@@ -117,7 +118,8 @@ class TerminalBlock {
       // extract filter with negation
       if(first instanceof PsiIfStatement) {
         PsiIfStatement ifStatement = (PsiIfStatement)first;
-        if(ifStatement.getCondition() == null) return null;
+        PsiExpression condition = ifStatement.getCondition();
+        if(condition == null) return null;
         PsiStatement branch = ifStatement.getThenBranch();
         if(branch instanceof PsiBlockStatement) {
           PsiStatement[] statements = ((PsiBlockStatement)branch).getCodeBlock().getStatements();
@@ -132,12 +134,32 @@ class TerminalBlock {
         } else {
           statements = Arrays.copyOfRange(myStatements, 1, myStatements.length);
         }
-        return new TerminalBlock(this, new FilterOp(ifStatement.getCondition(), myVariable, true), myVariable, statements);
+        return fromCondition(condition, true, statements);
       }
     }
     return null;
   }
 
+  @Nullable
+  private TerminalBlock fromCondition(PsiExpression condition, boolean negated, PsiStatement... statements) {
+    TerminalBlock result = new TerminalBlock(this, new FilterOp(condition, myVariable, negated), myVariable, statements);
+    List<PsiPatternVariable> vars = JavaPsiPatternUtil.getExposedPatternVariables(condition);
+    if (!vars.isEmpty()) {
+      List<PsiPatternVariable> used =
+        ContainerUtil.filter(vars, var -> Stream.of(statements).anyMatch(st -> VariableAccessUtils.variableIsUsed(var, st)));
+      if (used.size() > 1) return null;
+      if (!used.isEmpty()) {
+        PsiPatternVariable var = used.get(0);
+        String text = JavaPsiPatternUtil.getEffectiveInitializerText(var);
+        if (text == null) return null;
+        if (Stream.of(statements).anyMatch(st -> VariableAccessUtils.variableIsUsed(myVariable, st))) return null;
+        PsiExpression mappingExpression = JavaPsiFacade.getElementFactory(condition.getProject()).createExpressionFromText(text, var);
+        result = new TerminalBlock(result, new MapOp(mappingExpression, myVariable, var.getType()), var, statements);
+      }
+    }
+    return result;
+  }
+
   /**
    * Returns an equivalent {@code TerminalBlock} with one more intermediate operation extracted
    * or null if extraction is not possible.

java/java-psi-impl/src/com/intellij/psi/impl/source/tree/java/PsiPatternVariableImpl.java

@@ -157,6 +157,16 @@ public class PsiPatternVariableImpl extends CompositePsiElement implements PsiPa
     return new LocalSearchScope(getDeclarationScope());
   }
 
+  @Override
+  public void delete() throws IncorrectOperationException {
+    PsiPattern pattern = getPattern();
+    if (pattern instanceof PsiTypeTestPattern) {
+      replace(getTypeElement());
+      return;
+    }
+    super.delete();
+  }
+
   @Override
   public String toString() {
     return "PsiPatternVariable:" + getName();

java/java-psi-impl/src/com/intellij/psi/util/JavaPsiPatternUtil.java

@@ -19,7 +19,7 @@ public class JavaPsiPatternUtil {
    * @return list of pattern variables declared within an expression that could be visible outside of given expression.
    */
   @Contract(pure = true)
-  public static @NotNull List<PsiPatternVariable> getPatternVariablesVisibleOutsideOf(@NotNull PsiExpression expression) {
+  public static @NotNull List<PsiPatternVariable> getExposedPatternVariables(@NotNull PsiExpression expression) {
     PsiElement parent = PsiUtil.skipParenthesizedExprUp(expression.getParent());
     boolean parentMayAccept =
       parent instanceof PsiPrefixExpression && ((PsiPrefixExpression)parent).getOperationTokenType().equals(JavaTokenType.EXCL) ||
@@ -30,7 +30,19 @@ public class JavaPsiPatternUtil {
       return Collections.emptyList();
     }
     List<PsiPatternVariable> list = new ArrayList<>();
-    collectPatternVariableCandidates(expression, expression, list);
+    collectPatternVariableCandidates(expression, expression, list, false);
+    return list;
+  }
+
+  /**
+   * @param expression expression to search pattern variables in
+   * @return list of pattern variables declared within an expression that could be visible outside of given expression
+   * under some other parent (e.g. under PsiIfStatement).
+   */
+  @Contract(pure = true)
+  public static @NotNull List<PsiPatternVariable> getExposedPatternVariablesIgnoreParent(@NotNull PsiExpression expression) {
+    List<PsiPatternVariable> list = new ArrayList<>();
+    collectPatternVariableCandidates(expression, expression, list, true);
     return list;
   }
 
@@ -50,7 +62,7 @@ public class JavaPsiPatternUtil {
   }
 
   private static void collectPatternVariableCandidates(@NotNull PsiExpression scope, @NotNull PsiExpression expression,
-                                                       Collection<PsiPatternVariable> candidates) {
+                                                       Collection<PsiPatternVariable> candidates, boolean strict) {
     while (true) {
       if (expression instanceof PsiParenthesizedExpression) {
         expression = ((PsiParenthesizedExpression)expression).getExpression();
@@ -67,7 +79,7 @@ public class JavaPsiPatternUtil {
       PsiPattern pattern = ((PsiInstanceOfExpression)expression).getPattern();
       if (pattern instanceof PsiTypeTestPattern) {
         PsiPatternVariable variable = ((PsiTypeTestPattern)pattern).getPatternVariable();
-        if (variable != null && !PsiTreeUtil.isAncestor(scope, variable.getDeclarationScope(), false)) {
+        if (variable != null && !PsiTreeUtil.isAncestor(scope, variable.getDeclarationScope(), strict)) {
           candidates.add(variable);
         }
       }
@@ -77,7 +89,7 @@ public class JavaPsiPatternUtil {
       IElementType tokenType = polyadicExpression.getOperationTokenType();
       if (tokenType.equals(JavaTokenType.ANDAND) || tokenType.equals(JavaTokenType.OROR)) {
         for (PsiExpression operand : polyadicExpression.getOperands()) {
-          collectPatternVariableCandidates(scope, operand, candidates);
+          collectPatternVariableCandidates(scope, operand, candidates, strict);
         }
       }
     }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/removeUnusedVariable/afterPatternVariable.java

@@ -1,4 +1,4 @@
-// "Remove pattern variable 'string'" "true"
+// "Remove pattern variable" "true"
 class X {
   public void test(Object object) {
     if (object instanceof String) {}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/removeUnusedVariable/beforePatternVariable.java

@@ -1,4 +1,4 @@
-// "Remove pattern variable 'string'" "true"
+// "Remove pattern variable" "true"
 class X {
   public void test(Object object) {
     if (object instanceof String str<caret>ing) {}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/afterPattern.java

@@ -0,0 +1,10 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+
+class X {
+  void test(List<Object> list) {
+    List<String> result = list.stream().filter(o -> o instanceof String).map(o -> (String) o).collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/afterPatternOriginalUsed.java

@@ -0,0 +1,12 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+
+class X {
+  void test(List<Object> list) {
+    List<Object> result = list.stream().filter(o -> getObject(o) instanceof String s && !s.isEmpty()).collect(Collectors.toList());
+  }
+  
+  native Object getObject(Object obj);
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/afterPatternUsed.java

@@ -0,0 +1,10 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+
+class X {
+  void test(List<Object> list) {
+    List<String> result = list.stream().filter(o -> o instanceof String s && !s.isEmpty()).map(o -> (String) o).collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/beforePattern.java

@@ -0,0 +1,14 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+
+class X {
+  void test(List<Object> list) {
+    List<String> result = new ArrayList<>();
+    f<caret>or (Object o : list) {
+      if (o instanceof String s) {
+        result.add(s);
+      }
+    }
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/beforePatternOriginalUsed.java

@@ -0,0 +1,16 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+
+class X {
+  void test(List<Object> list) {
+    List<Object> result = new ArrayList<>();
+    f<caret>or (Object o : list) {
+      if (getObject(o) instanceof String s && !s.isEmpty()) {
+        result.add(o);
+      }
+    }
+  }
+  
+  native Object getObject(Object obj);
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/java14/beforePatternUsed.java

@@ -0,0 +1,14 @@
+// "Fix all 'Loop can be collapsed with Stream API' problems in file" "true"
+import java.util.List;
+import java.util.ArrayList;
+
+class X {
+  void test(List<Object> list) {
+    List<String> result = new ArrayList<>();
+    f<caret>or (Object o : list) {
+      if (o instanceof String s && !s.isEmpty()) {
+        result.add(s);
+      }
+    }
+  }
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/daemon/quickFix/StreamApiMigrationInspectionTestSuite.java

@@ -50,6 +50,7 @@ import org.junit.runners.Suite;
   StreamApiMigrationInspectionTestSuite.SummingTest.class,
   StreamApiMigrationInspectionTestSuite.Java9Test.class,
   StreamApiMigrationInspectionTestSuite.Java10Test.class,
+  StreamApiMigrationInspectionTestSuite.Java14Test.class,
 })
 public class StreamApiMigrationInspectionTestSuite {
   public static abstract class StreamApiMigrationInspectionBaseTest extends LightQuickFixParameterizedTestCase {
@@ -228,4 +229,11 @@ public class StreamApiMigrationInspectionTestSuite {
       return "java10";
     }
   }
+
+  public static class Java14Test extends StreamApiMigrationInspectionBaseTest {
+    @Override
+    String getFolder() {
+      return "java14";
+    }
+  }
 }