From a042720e74a4f4dec631617c86e997e93da8356c Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Thu, 6 Mar 2014 14:52:08 +0100
Subject: [PATCH] stream migration inspection: preserve comments (IDEA-121679)

---
 .../StreamApiMigrationInspection.java          | 14 ++++++++++++++
 .../afterCollectArrayListComment.java          | 16 ++++++++++++++++
 .../afterNormalComments.java                   | 11 +++++++++++
 .../beforeCollectArrayListComment.java         | 18 ++++++++++++++++++
 .../beforeNormalComments.java                  | 13 +++++++++++++
 5 files changed, 72 insertions(+)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectArrayListComment.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterNormalComments.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectArrayListComment.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeNormalComments.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java b/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java
index fccc77a59b06..b06ba2a5358b 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java
+++ b/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();
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectArrayListComment.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectArrayListComment.java
new file mode 100644
index 000000000000..a9b24a5a2005
--- /dev/null
+++ b/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
+  }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterNormalComments.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterNormalComments.java
new file mode 100644
index 000000000000..c1bbf0a0924b
--- /dev/null
+++ b/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);
+  }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectArrayListComment.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectArrayListComment.java
new file mode 100644
index 000000000000..c26fa6954bd6
--- /dev/null
+++ b/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());
+    }
+  }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeNormalComments.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeNormalComments.java
new file mode 100644
index 000000000000..f394dd7d9696
--- /dev/null
+++ b/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);
+    }
+  }
+}