IDEA-162945 Stream API migration: support cases where variable is modified and reassigned

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

@@ -390,31 +390,14 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     return consumerClass != null ? psiFacade.getElementFactory().createType(consumerClass, variable.getType()) : null;
   }
 
-  static boolean isVariableSuitableForStream(PsiVariable variable, PsiStatement statement, StreamSource source) {
-    PsiElement declaration = variable.getParent();
-    // For-loop initializer is not effectively final, but suitable for stream conversion
-    if(declaration instanceof PsiDeclarationStatement) {
-      PsiElement grandParent = declaration.getParent();
-      if (grandParent instanceof PsiForStatement) {
-        PsiForStatement forStatement = (PsiForStatement)grandParent;
-        if (forStatement.getInitialization() == declaration) {
-          PsiStatement body = forStatement.getBody();
-          if(body != null && PsiTreeUtil.isAncestor(statement, body, false)) {
-            return ReferencesSearch.search(variable, new LocalSearchScope(body)).forEach(ref -> {
-              PsiElement element = ref.getElement();
-              return !(element instanceof PsiExpression) || !PsiUtil.isAccessedForWriting((PsiExpression)element);
-            });
-          }
-        }
-      }
-    }
-    if(statement instanceof PsiWhileStatement && source.getVariable() == variable) {
-      return ReferencesSearch.search(variable, variable.getUseScope()).forEach(ref -> {
-        PsiElement element = ref.getElement();
-        return !(element instanceof PsiExpression) ||
-               PsiTreeUtil.isAncestor(((PsiWhileStatement)statement).getCondition(), element, false) ||
-               !PsiUtil.isAccessedForWriting((PsiExpression)element);
-      });
+  static boolean isVariableSuitableForStream(PsiVariable variable, PsiStatement statement, TerminalBlock tb) {
+    if(ReferencesSearch.search(variable, variable.getUseScope()).forEach(ref -> {
+      PsiElement element = ref.getElement();
+      return !(element instanceof PsiExpression) ||
+             !PsiUtil.isAccessedForWriting((PsiExpression)element) ||
+             tb.operations().anyMatch(op -> op.isWriteAllowed(variable, (PsiExpression)element));
+    })) {
+      return true;
     }
     return HighlightControlFlowUtil.isEffectivelyFinal(variable, statement, null);
   }
@@ -590,7 +573,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       int startOffset = controlFlow.getStartOffset(body);
       int endOffset = controlFlow.getEndOffset(body);
       final List<PsiVariable> nonFinalVariables = StreamEx.of(ControlFlowUtil.getUsedVariables(controlFlow, startOffset, endOffset))
-        .remove(variable -> isVariableSuitableForStream(variable, statement, source)).toList();
+        .remove(variable -> isVariableSuitableForStream(variable, statement, tb)).toList();
 
       if (exitPoints.isEmpty()) {
         if(getIncrementedVariable(tb, nonFinalVariables) != null) {
@@ -902,6 +885,10 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     }
 
     abstract String createReplacement();
+
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      return false;
+    }
   }
 
   static class FilterOp extends Operation {
@@ -993,6 +980,11 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       PsiExpression expression = myType == null ? myExpression : RefactoringUtil.convertInitializerToNormalExpression(myExpression, myType);
       return "." + operationName + "(" + LambdaUtil.createLambda(myVariable, expression) + ")";
     }
+
+    @Override
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      return variable == myVariable && reference.getParent() == myExpression.getParent();
+    }
   }
 
   static class FlatMapOp extends Operation {
@@ -1026,6 +1018,11 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       return mySource.createReplacement();
     }
 
+    @Override
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      return mySource.isWriteAllowed(variable, reference);
+    }
+
     boolean breaksMe(PsiBreakStatement statement) {
       return statement.findExitedStatement() == myLoop;
     }
@@ -1070,6 +1067,11 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       myVariable.delete();
     }
 
+    @Override
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      return myVariable == variable && reference.getParent() == PsiTreeUtil.getParentOfType(myExpression, PsiAssignmentExpression.class);
+    }
+
     @Nullable
     public static BufferedReaderLines from(PsiWhileStatement whileLoop) {
       // while ((line = br.readLine()) != null)
@@ -1197,6 +1199,17 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       return className+"."+methodName+"("+myExpression.getText()+", "+myBound.getText()+")";
     }
 
+    @Override
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      if(variable == myVariable) {
+        PsiForStatement forStatement = PsiTreeUtil.getParentOfType(variable, PsiForStatement.class);
+        if(forStatement != null) {
+          return PsiTreeUtil.isAncestor(forStatement.getUpdate(), reference, false);
+        }
+      }
+      return false;
+    }
+
     @Nullable
     public static CountingLoop from(PsiForStatement forStatement) {
       // check that initialization is for(int/long i = <initial_value>;...;...)
@@ -1402,6 +1415,16 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
             }
           }
         }
+        PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(first);
+        if(assignment != null) {
+          PsiExpression lValue = assignment.getLExpression();
+          PsiExpression rValue = assignment.getRExpression();
+          if(rValue != null && lValue instanceof PsiReferenceExpression && ((PsiReferenceExpression)lValue).isReferenceTo(myVariable)) {
+            PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
+            MapOp op = new MapOp(myPreviousOp, rValue, myVariable, myVariable.getType());
+            return new TerminalBlock(op, myVariable, leftOver);
+          }
+        }
       }
       return null;
     }
@@ -1434,7 +1457,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     }
 
     @NotNull
-    private StreamEx<Operation> operations() {
+    StreamEx<Operation> operations() {
       return StreamEx.iterate(myPreviousOp, Objects::nonNull, Operation::getPreviousOp);
     }
 

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterBufferedReaderModifiedLine.java

@@ -0,0 +1,11 @@
+// "Replace with sum()" "true"
+
+import java.io.BufferedReader;
+import java.io.IOException;
+
+public class Main {
+  void test(BufferedReader br) throws IOException {
+      long count = br.lines().map(String::trim).mapToLong(String::length).sum();
+      System.out.println(count);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFilterMapFlatMapArrayMapCollect.java

@@ -0,0 +1,12 @@
+// "Replace with collect" "true"
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class Main {
+  public void test(Map<String, String[]> map) {
+      List<String> result = map.entrySet().stream().filter(entry -> entry.getKey().startsWith("x")).map(Map.Entry::getValue).flatMap(Arrays::stream).map(String::trim).collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeBufferedReaderComplexModifiedLine.java

@@ -0,0 +1,15 @@
+// "Replace with sum()" "false"
+
+import java.io.BufferedReader;
+import java.io.IOException;
+
+public class Main {
+  void test(BufferedReader br) throws IOException {
+    String line = "";
+    long count = 0;
+    wh<caret>ile((line = br.readLine()) != null) {
+      count+=(line = line.trim()).length();
+    }
+    System.out.println(count);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeBufferedReaderModifiedLine.java

@@ -1,4 +1,4 @@
-// "Replace with sum()" "false"
+// "Replace with sum()" "true"
 
 import java.io.BufferedReader;
 import java.io.IOException;

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFilterMapFlatMapArrayMapCollect.java

@@ -0,0 +1,19 @@
+// "Replace with collect" "true"
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  public void test(Map<String, String[]> map) {
+    List<String> result = new ArrayList<>();
+    for(Map.Entry<String, String[]> entry: m<caret>ap.entrySet()) {
+      if(entry.getKey().startsWith("x")) {
+        String[] arr = entry.getValue();
+        for (String str : arr) {
+          str = str.trim();
+          result.add(str);
+        }
+      }
+    }
+  }
+}