StreamApiMigration: register comments from post-loop operations like sorting

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

@@ -71,14 +71,15 @@ class CollectMigration extends BaseStreamApiMigration {
     PsiElement result;
     if (toReplace != null) {
       result = ct.replace(toReplace, stream);
+      terminal.cleanUp(ct);
       removeLoop(ct, loopStatement);
     }
     else {
       PsiVariable variable = terminal.getTargetVariable();
       LOG.assertTrue(variable != null);
+      terminal.cleanUp(ct);
       result = replaceInitializer(loopStatement, variable, variable.getInitializer(), stream, terminal.getStatus(), ct);
     }
-    terminal.cleanUp();
     return result;
   }
 
@@ -216,7 +217,7 @@ class CollectMigration extends BaseStreamApiMigration {
 
     public InitializerUsageStatus getStatus() { return myStatus; }
 
-    void cleanUp() {}
+    void cleanUp(CommentTracker ct) {}
 
     boolean isTrivial() {
       return generateIntermediate(new CommentTracker()).isEmpty();
@@ -670,9 +671,9 @@ class CollectMigration extends BaseStreamApiMigration {
     }
 
     @Override
-    public void cleanUp() {
-      myDownstream.cleanUp();
-      myStatement.delete();
+    public void cleanUp(CommentTracker ct) {
+      myDownstream.cleanUp(ct);
+      ct.delete(myStatement);
     }
 
     @Override
@@ -754,12 +755,12 @@ class CollectMigration extends BaseStreamApiMigration {
     }
 
     @Override
-    public void cleanUp() {
+    public void cleanUp(CommentTracker ct) {
       PsiLocalVariable variable = myUpstream.getTargetVariable();
       if (variable != null && myUpstream.getStatus() != ControlFlowUtils.InitializerUsageStatus.AT_WANTED_PLACE) {
-        variable.delete();
+        ct.delete(variable);
       }
-      myUpstream.cleanUp();
+      myUpstream.cleanUp(ct);
     }
   }
 

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

@@ -205,6 +205,7 @@ public class FuseStreamOperationsInspection extends AbstractBaseJavaLocalInspect
       String stream = terminal.generateIntermediate(ct) + terminal.generateTerminal(ct, myStrictMode);
       PsiElement toReplace = terminal.getElementToReplace();
       PsiElement result;
+      terminal.cleanUp(ct);
       if (toReplace != null) {
         result = ct.replaceAndRestoreComments(toReplace, stream);
       }
@@ -213,7 +214,6 @@ public class FuseStreamOperationsInspection extends AbstractBaseJavaLocalInspect
         PsiExpression initializer = Objects.requireNonNull(variable.getInitializer());
         result = ct.replaceAndRestoreComments(initializer, stream);
       }
-      terminal.cleanUp();
       LambdaCanBeMethodReferenceInspection.replaceAllLambdasWithMethodReferences(result);
     }
   }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/fuseStreamOperations/afterSetListSort.java

@@ -5,6 +5,9 @@ import java.util.stream.Stream;
 
 public class Test {
   public void testSetListSort(String[] args) {
+      // foo
+      // bar
+      /*baz*/
       System.out.println(Arrays.stream(args).distinct().sorted().toArray());
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/fuseStreamOperations/beforeSetListSort.java

@@ -6,8 +6,8 @@ import java.util.stream.Stream;
 public class Test {
   public void testSetListSort(String[] args) {
     Set<String> set = Arrays.stream(args).co<caret>llect(Collectors.toSet());
-    List<String> list = new ArrayList<>(set);
-    list.sort(null);
-    System.out.println(list.toArray());
+    List<String> list = new ArrayList<>(set); // foo
+    list.sort(null); // bar
+    System.out.println(list.toArray(/*baz*/));
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/collect/afterCollectSorted.java

@@ -10,6 +10,7 @@ public class Collect {
   }
 
   void collectNames(List<Person> persons){
-    List<String> names = persons.stream().map(Person::getName).sorted(Comparator.comparing(Person::getName)).collect(Collectors.toList());
+    List<String> names = persons.stream().map(Person::getName).sorted(Comparator.comparing(/*c2*/Person::getName)).collect(Collectors.toList());
+      // comment
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/collect/beforeCollectSorted.java

@@ -13,6 +13,7 @@ public class Collect {
     for (Person person : pers<caret>ons) {
       names.add(person.getName());
     }
-    Collections.sort(names, Comparator.comparing(Person::getName));
+    Collections.// comment
+                 sort(names, Comparator.comparing(/*c2*/Person::getName));
   }
 }

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/CommentTracker.java

@@ -313,6 +313,8 @@ public class CommentTracker {
 
   private void grabCommentsOnDelete(PsiElement element) {
     if (element instanceof PsiExpression && element.getParent() instanceof PsiExpressionStatement ||
+        (element.getParent() instanceof PsiDeclarationStatement &&
+         ((PsiDeclarationStatement)element.getParent()).getDeclaredElements().length == 1) ||
         element.getParent() instanceof PsiJavaCodeReferenceElement) {
       element = element.getParent();
     }