StreamApiMigrationInspection: try continue extracting operations when limit is successfully peeled off

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

@@ -652,13 +652,18 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         .remove(variable -> isVariableSuitableForStream(variable, loop, tb)).toList();
 
       TerminalBlock tbWithLimit = tb.tryPeelLimit(loop);
-      if (isCountOperation(loop, nonFinalVariables, tbWithLimit)) {
+      List<PsiVariable> nonFinalVariablesWithLimit = nonFinalVariables;
+      if(tbWithLimit != tb) {
+        nonFinalVariablesWithLimit = StreamEx.of(nonFinalVariables)
+          .remove(variable -> isVariableSuitableForStream(variable, loop, tbWithLimit)).toList();
+      }
+      if (isCountOperation(loop, nonFinalVariablesWithLimit, tbWithLimit)) {
         return new CountMigration();
       }
       if (getAccumulatedVariable(tb, nonFinalVariables) != null) {
         return new SumMigration();
       }
-      if (isCollectCall(tbWithLimit) && nonFinalVariables.isEmpty()) {
+      if (isCollectCall(tbWithLimit) && nonFinalVariablesWithLimit.isEmpty()) {
         return findCollectMigration(loop, tbWithLimit);
       }
       if (isCollectMapCall(loop, tb) && nonFinalVariables.isEmpty() && (REPLACE_TRIVIAL_FOREACH || tb.hasOperations())) {
@@ -999,6 +1004,11 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         LambdaUtil.createLambda(myMatchVariable, myExpression)+")", myExpression);
     }
 
+    @Override
+    boolean isWriteAllowed(PsiVariable variable, PsiExpression reference) {
+      return myFlatMapOp.isWriteAllowed(variable, reference);
+    }
+
     @Override
     StreamEx<PsiExpression> expressions() {
       return StreamEx.of(myExpression, myFlatMapOp.myExpression);
@@ -1572,18 +1582,30 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         default:
           return this;
       }
-      PsiExpression counter = PsiUtil.skipParenthesizedExprDown(flipped ? binOp.getROperand() : binOp.getLOperand());
-      if(counter == null || VariableAccessUtils.variableIsUsed(myVariable, counter)) return this;
+      PsiExpression countExpression = PsiUtil.skipParenthesizedExprDown(flipped ? binOp.getROperand() : binOp.getLOperand());
+      if(countExpression == null || VariableAccessUtils.variableIsUsed(myVariable, countExpression)) return this;
       PsiExpression limit = flipped ? binOp.getLOperand() : binOp.getROperand();
       if(!ExpressionUtils.isSimpleExpression(limit) || VariableAccessUtils.variableIsUsed(myVariable, limit)) return this;
       PsiType type = limit.getType();
       if(!PsiType.INT.equals(type) && !PsiType.LONG.equals(type)) return this;
-      if(counter instanceof PsiPostfixExpression) {
+      if(countExpression instanceof PsiPostfixExpression) {
         delta++;
       }
 
-      LimitOp limitOp = new LimitOp(filter.getPreviousOp(), counter, limit, myVariable, delta);
-      return new TerminalBlock(limitOp, myVariable, statements);
+      Operation prev = filter.getPreviousOp();
+      LOG.assertTrue(prev != null);
+      TerminalBlock block = new TerminalBlock(prev, myVariable, statements);
+      if(extractIncrementedLValue(countExpression) == null) {
+        // when countExpression does not change the counter, we may try to continue extracting ops from the remaining statement
+        // this is helpful to cover cases like for(...) { if(...) list.add(x); if(list.size == limit) break; }
+        while (true) {
+          TerminalBlock newBlock = block.extractOperation();
+          if (newBlock == null || newBlock.getLastOperation() instanceof FlatMapOp) break;
+          block = newBlock;
+        }
+      }
+      LimitOp limitOp = new LimitOp(block.getLastOperation(), countExpression, limit, block.getVariable(), delta);
+      return new TerminalBlock(limitOp, block.getVariable(), block.getStatements());
     }
 
     @NotNull

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

@@ -0,0 +1,13 @@
+// "Replace with collect" "true"
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class Main {
+  List<String> test(Map<String, List<String>> map, int limit) {
+      List<String> list = map.entrySet().stream().filter(entry -> entry.getValue() != null).filter(entry -> entry.getValue().stream().anyMatch(str -> str.contains("foo"))).limit(limit).map(Map.Entry::getKey).collect(Collectors.toList());
+      return list;
+  }
+}

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

@@ -0,0 +1,13 @@
+// "Replace with collect" "true"
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class Main {
+  List<String> test(Map<String, List<String>> map, int limit) {
+      List<String> list = map.entrySet().stream().filter(entry -> entry.getValue() != null).flatMap(entry -> entry.getValue().stream()).filter(str -> str.contains("foo")).limit(limit).collect(Collectors.toList());
+      return list;
+  }
+}

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

@@ -0,0 +1,23 @@
+// "Replace with collect" "true"
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  List<String> test(Map<String, List<String>> map, int limit) {
+    List<String> list = new ArrayList<>();
+    for (Map.Entry<String, List<String>> entry : map.<caret>entrySet()) {
+      if(entry.getValue() != null) {
+        for(String str : entry.getValue()) {
+          if(str.contains("foo")) {
+            list.add(entry.getKey());
+            break;
+          }
+        }
+      }
+      if(list.size() >= limit) break;
+    }
+    return list;
+  }
+}

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

@@ -0,0 +1,22 @@
+// "Replace with collect" "true"
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  List<String> test(Map<String, List<String>> map, int limit) {
+    List<String> list = new ArrayList<>();
+    for (Map.Entry<String, List<String>> entry : map.<caret>entrySet()) {
+      if(entry.getValue() != null) {
+        for(String str : entry.getValue()) {
+          if(str.contains("foo")) {
+            list.add(str);
+          }
+          if(list.size() >= limit) return list;
+        }
+      }
+    }
+    return list;
+  }
+}

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

@@ -0,0 +1,22 @@
+// "Replace with collect" "false"
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  List<String> test(Map<String, List<String>> map, int limit) {
+    List<String> list = new ArrayList<>();
+    for (Map.Entry<String, List<String>> entry : map.<caret>entrySet()) {
+      if(entry.getValue() != null) {
+        for(String str : entry.getValue()) {
+          if(str.contains("foo")) {
+            list.add(str);
+          }
+        }
+        if(list.size() >= limit) return list;
+      }
+    }
+    return list;
+  }
+}