support return in forEach replacement

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

@@ -84,6 +84,8 @@ class ForEachMigration extends BaseStreamApiMigration {
       return loopStatement.replace(factory.createStatementFromText(callText, loopStatement));
     }
 
+    tb.replaceContinueWithReturn(factory);
+
     String stream = tb.generate(true) + "." + getReplacement() + "(";
     PsiElement block = tb.convertToElement(factory);
 

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

@@ -41,7 +41,6 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.psi.util.TypeConversionUtil;
 import com.intellij.util.ArrayUtil;
-import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.ig.psiutils.*;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
@@ -509,8 +508,9 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       }
       Collection<PsiStatement> exitPoints = tb.findExitPoints(controlFlow);
       if (exitPoints == null) return null;
-      if ((exitPoints.isEmpty()) &&
-          nonFinalVariables.isEmpty()) {
+      boolean onlyNonLabeledContinue = StreamEx.of(exitPoints).allMatch(statement -> statement instanceof PsiContinueStatement &&
+                                                                ((PsiContinueStatement)statement).getLabelIdentifier() == null);
+      if (onlyNonLabeledContinue && nonFinalVariables.isEmpty()) {
         boolean shouldWarn = SUGGEST_FOREACH &&
                              (REPLACE_TRIVIAL_FOREACH ||
                               tb.hasOperations() ||

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

@@ -463,6 +463,24 @@ class TerminalBlock {
     return block;
   }
 
+
+  /**
+   * method replaces continue statement (without labels) to return statement
+   * @param factory factory to use to create new element if necessary
+   */
+  void replaceContinueWithReturn(PsiElementFactory factory) {
+    for (int i = 0, length = myStatements.length; i < length; i++) {
+      PsiStatement statement = myStatements[i];
+      if(statement instanceof PsiContinueStatement) {
+        myStatements[i] = factory.createStatementFromText("return;", null);
+        continue;
+      }
+      StreamEx.ofTree(statement, (PsiElement s) -> StreamEx.of(s.getChildren()))
+        .select(PsiContinueStatement.class)
+        .forEach(stmt -> stmt.replace(factory.createStatementFromText("return;", null)));
+    }
+  }
+
   String generate() {
     return generate(false);
   }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/continue/afterContinueInside.java

@@ -0,0 +1,14 @@
+// "Replace with forEach" "true"
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+class Sample {
+  List<String> foo = new ArrayList<>();
+  {
+      foo.stream().filter(Objects::isNull).forEach(s -> {
+          return;
+      });
+
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/continue/afterForEachContinue.java

@@ -0,0 +1,14 @@
+// "Replace with forEach" "true"
+
+import java.util.Arrays;
+import java.util.Collection;
+
+public class Test {
+  void test(int[] arr) {
+      Arrays.stream(arr).forEach(x -> {
+          int y = x * 2;
+          if (x > y) return;
+          System.out.println(x);
+      });
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/continue/beforeContinueInside.java

@@ -1,4 +1,4 @@
-// "Replace with forEach" "false"
+// "Replace with forEach" "true"
 import java.util.ArrayList;
 import java.util.List;
 
@@ -12,4 +12,4 @@ class Sample {
     }
 
   }
-}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/foreach/afterForEachContinue.java

@@ -1,13 +0,0 @@
-// "Replace with forEach" "true"
-
-import java.util.Collection;
-
-public class Test {
-  void test(int[] arr) {
-    for<caret>(int x : arr) {
-      int y = x*2;
-      if(x > y) continue;
-      System.out.println(x);
-    }
-  }
-}