stream migration inspection: preserve comments (IDEA-121679)

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

@@ -195,6 +195,13 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         PsiStatement body = foreachStatement.getBody();
         final PsiExpression iteratedValue = foreachStatement.getIteratedValue();
         if (body != null && iteratedValue != null) {
+          final Collection<PsiComment> comments = PsiTreeUtil.findChildrenOfType(body, PsiComment.class);
+
+          final PsiElement parent = foreachStatement.getParent();
+          for (PsiElement comment : PsiTreeUtil.findChildrenOfType(body, PsiComment.class)) {
+            parent.addBefore(comment, foreachStatement);
+          }
+
           final PsiParameter parameter = foreachStatement.getIterationParameter();
           final PsiIfStatement ifStmt = extractIfStatement(body);
 
@@ -314,6 +321,13 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
             variableName = "";
           }
 
+          if (variableName != null) {
+            final PsiElement parent = foreachStatement.getParent();
+            for (PsiElement comment : PsiTreeUtil.findChildrenOfType(body, PsiComment.class)) {
+              parent.addBefore(comment, foreachStatement);
+            }
+          }
+
           PsiElement result = null;
           if (primitiveInitializer != null) {
             final PsiType initializerType = primitiveInitializer.getType();

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

@@ -0,0 +1,16 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    List<String> names = persons.stream().map(Person::getName).collect(Collectors.toList());
+      //some comment
+  }
+}

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

@@ -0,0 +1,11 @@
+// "Replace with forEach" "true"
+import java.util.ArrayList;
+import java.util.List;
+
+class Sample {
+  List<String> foo = new ArrayList<>();
+  {
+      //some comment
+      foo.forEach(System.out::println);
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    List<String> names = new ArrayList<>();
+    for (Person person : pers<caret>ons) {
+      //some comment
+      names.add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,13 @@
+// "Replace with forEach" "true"
+import java.util.ArrayList;
+import java.util.List;
+
+class Sample {
+  List<String> foo = new ArrayList<>();
+  {
+    for (String s : fo<caret>o) {
+      //some comment
+      System.out.println(s);
+    }
+  }
+}