Better comment handling in ternary->if conversion & stream-to-loop

java/java-impl/src/com/intellij/codeInspection/streamToLoop/StreamToLoopInspection.java

@@ -608,7 +608,14 @@ public class StreamToLoopInspection extends AbstractBaseJavaLocalInspectionTool
         PsiElement placeHolderCopy = PsiTreeUtil.releaseMark(returnCopy, mark);
         LOG.assertTrue(placeHolderCopy != null);
         PsiElement replacement = placeHolderCopy.replace(createExpression(conditionalExpression.getTrueBranch()));
-        return (placeHolderCopy == returnCopy ? replacement : returnCopy).getText();
+        if (returnCopy == placeHolderCopy) {
+          returnCopy = replacement;
+        }
+        String text = returnCopy.getText();
+        if (returnCopy.getLastChild() instanceof PsiComment) {
+          text += "\n";
+        }
+        return text;
       }
       PsiElement parent = PsiUtil.skipParenthesizedExprUp(myStreamExpression.getParent());
       if(parent instanceof PsiIfStatement && conditionalExpression instanceof ConditionalExpression.Boolean &&

java/java-impl/src/com/intellij/refactoring/util/EnsureCodeBlockImpl.java

@@ -214,25 +214,31 @@ class EnsureCodeBlockImpl {
         statement = (PsiStatement)assignment.getParent();
       }
     }
+    CommentTracker ct = new CommentTracker();
     PsiIfStatement ifStatement =
-      (PsiIfStatement)factory.createStatementFromText("if(" + ternary.getCondition().getText() + ") {} else {}", statement);
+      (PsiIfStatement)factory.createStatementFromText("if(" + ct.text(ternary.getCondition()) + ") {} else {}", statement);
     Object mark = new Object();
     PsiTreeUtil.mark(ternary, mark);
+    for (PsiElement child : statement.getChildren()) {
+      if (child instanceof PsiComment) {
+        ct.delete(child);
+      }
+    }
     PsiStatement thenStatement = (PsiStatement)statement.copy();
     PsiConditionalExpression thenTernary = Objects.requireNonNull((PsiConditionalExpression)PsiTreeUtil.releaseMark(thenStatement, mark));
-    PsiExpression branch1 = ternary.getThenExpression();
-    if (branch1 != null) {
-      thenTernary.replace(branch1);
+    PsiExpression thenBranch = ternary.getThenExpression();
+    if (thenBranch != null) {
+      thenTernary.replace(ct.markUnchanged(thenBranch));
     }
     PsiStatement elseStatement = (PsiStatement)statement.copy();
     PsiConditionalExpression elseTernary = Objects.requireNonNull((PsiConditionalExpression)PsiTreeUtil.releaseMark(elseStatement, mark));
-    PsiExpression branch = ternary.getElseExpression();
-    if (branch != null) {
-      elseTernary.replace(branch);
+    PsiExpression elseBranch = ternary.getElseExpression();
+    if (elseBranch != null) {
+      elseTernary.replace(ct.markUnchanged(elseBranch));
     }
     ((PsiBlockStatement)Objects.requireNonNull(ifStatement.getThenBranch())).getCodeBlock().add(thenStatement);
     ((PsiBlockStatement)Objects.requireNonNull(ifStatement.getElseBranch())).getCodeBlock().add(elseStatement);
-    return statement.replace(ifStatement);
+    return ct.replaceAndRestoreComments(statement, ifStatement);
   }
 
   private static PsiElement splitIf(PsiIfStatement outerIf, PsiPolyadicExpression andChain, PsiExpression operand) {

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterIfComment.java

@@ -0,0 +1,23 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.Arrays;
+import java.util.List;
+
+public class Main {
+  private static String test(List<String> list) {
+    if (list == null) {
+      return null;
+    } else {
+        for (String str : list) {
+            if (str.contains("x")) {
+                return str; // comment 
+            }
+        }
+        return null; // comment 
+    }
+  }
+
+  public static void main(String[] args) {
+    System.out.println(test(Arrays.asList("a", "b", "syz")));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterTernary.java

@@ -5,6 +5,9 @@ import java.util.List;
 
 public class Main {
   private static String test(List<String> list) {
+      // otherwise not null
+      // if list is null 
+      // return null
       if (list == null) {
           return null;
       } else {

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeIfComment.java

@@ -0,0 +1,18 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.Arrays;
+import java.util.List;
+
+public class Main {
+  private static String test(List<String> list) {
+    if (list == null) {
+      return null;
+    } else {
+      return list.stream().filter(str -> str.contains("x")).find<caret>First().orElse(null); // comment 
+    }
+  }
+
+  public static void main(String[] args) {
+    System.out.println(test(Arrays.asList("a", "b", "syz")));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeTernary.java

@@ -5,7 +5,9 @@ import java.util.List;
 
 public class Main {
   private static String test(List<String> list) {
-    return list == null ? null : list.stream().filter(str -> str.contains("x")).find<caret>First().orElse(null);
+    return list == null ? // if list is null 
+           null :  // return null
+           list.stream().filter(str -> str.contains("x")).find<caret>First().orElse(null); // otherwise not null
   }
 
   public static void main(String[] args) {