ReplaceStringBuilderWithString should keep comments (#511)

* Initial test case for comments

* Additional test cases and logic to retain comments

* Apply formatter

---------

Co-authored-by: Tim te Beek <[email protected]>

Author: Jammy-Louie

src/main/java/org/openrewrite/staticanalysis/ChainStringBuilderAppendCalls.java

@@ -27,7 +27,6 @@
 import java.util.ArrayList;
 import java.util.List;
 
-import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
 public class ChainStringBuilderAppendCalls extends Recipe {
@@ -126,8 +125,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
      */
     public static J.Binary concatAdditionBinary(Expression left, Expression right) {
         J.Binary b = getAdditiveBinaryTemplate();
-        String rightWhiteSpace = right.getPrefix().getWhitespace();
-        Space rightPrefix = rightWhiteSpace.isEmpty() ? Space.SINGLE_SPACE : Space.build(rightWhiteSpace, emptyList());
+        Space rightPrefix = right.getPrefix().isEmpty() ? Space.SINGLE_SPACE : right.getPrefix();
         return b.withPrefix(b.getLeft().getPrefix())
                 .withLeft(left)
                 .withRight(right.withPrefix(rightPrefix));

src/main/java/org/openrewrite/staticanalysis/ReplaceStringBuilderWithString.java

@@ -69,7 +69,15 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx)
                     Collections.reverse(arguments);
                     arguments = adjustExpressions(method, arguments);
 
-                    Expression additive = ChainStringBuilderAppendCalls.additiveExpression(arguments).withPrefix(method.getPrefix());
+                    Expression additive = ChainStringBuilderAppendCalls.additiveExpression(arguments);
+                    if (additive == null) {
+                        return m;
+                    }
+
+                    if (arguments.get(0).getComments().isEmpty() || !arguments.get(0).getPrefix().getWhitespace().startsWith("\n")) {
+                        additive = additive.withPrefix(method.getPrefix());
+                    }
+
                     if (isAMethodSelect(method)) {
                         additive = new J.Parentheses<>(randomId(), Space.EMPTY, Markers.EMPTY, JRightPadded.build(additive));
                     }
@@ -111,7 +119,7 @@ private List<Expression> adjustExpressions(J.MethodInvocation method, List<Expre
                             }
                         }
                     } else if (!(arg instanceof J.Identifier || arg instanceof J.Literal || arg instanceof J.MethodInvocation)) {
-                        return new J.Parentheses<>(randomId(), Space.EMPTY, Markers.EMPTY, JRightPadded.build(arg));
+                        return new J.Parentheses<>(randomId(), arg.getPrefix(), Markers.EMPTY, JRightPadded.build(arg.withPrefix(Space.EMPTY)));
                     }
                     return arg;
                 });
@@ -141,8 +149,12 @@ private boolean flatMethodInvocationChain(J.MethodInvocation method, List<Expres
                     if (args.size() != 1) {
                         return false;
                     } else {
-                        Space selectMethodWhiteSpace = selectMethod.getPadding().getSelect().getAfter();
-                        arguments.add(args.get(0).withPrefix(selectMethodWhiteSpace));
+                        JRightPadded<Expression> jrp = selectMethod.getPadding().getSelect();
+                        if (jrp == null) {
+                            arguments.add(args.get(0));
+                        } else {
+                            arguments.add(args.get(0).withPrefix(jrp.getAfter()));
+                        }
                     }
                 }
 
@@ -153,7 +165,9 @@ private boolean flatMethodInvocationChain(J.MethodInvocation method, List<Expres
                     if (nc.getArguments().size() == 1 && TypeUtils.isString(nc.getArguments().get(0).getType())) {
                         arguments.add(nc.getArguments().get(0));
                     } else if (!arguments.isEmpty()) {
-                        arguments.set(arguments.size() - 1, arguments.get(arguments.size() - 1).withPrefix(Space.EMPTY));
+                        Expression lastArgument = arguments.get(arguments.size() - 1);
+                        Space formattedPrefix = lastArgument.getComments().isEmpty() ? Space.EMPTY : lastArgument.getPrefix();
+                        arguments.set(arguments.size() - 1, lastArgument.withPrefix(formattedPrefix));
                     }
                     return true;
                 }

src/test/java/org/openrewrite/staticanalysis/ReplaceStringBuilderWithStringTest.java

@@ -230,6 +230,107 @@ String name() {
         );
     }
 
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite-static-analysis/issues/88")
+    void retainComments() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              class A {
+                  void foo() {
+                      String scenarioOne = new StringBuilder()
+                          // A
+                          .append("A")
+                          // B
+                          .append("B")
+                          // C
+                          .append("C")
+                          .toString();
+                      String scenarioTwo = new StringBuilder("A")
+                          // B
+                          .append("B")
+                          // C
+                          .append("C")
+                          .toString();
+                      String scenarioThree = new StringBuilder()
+                          // A
+                          .append("A")
+                          // B
+                          .append("B")
+                          // C
+                          .append("C")
+                          // test method
+                          .append(testString())
+                          .toString();
+                      String scenarioFour = new StringBuilder()
+                          // 1 + 1
+                          .append(1 + 1)
+                          // 1
+                          .append(1)
+                          .toString();
+                      String scenarioFive = new StringBuilder()
+                          // A
+                          .append("A")
+                          // 1 + 1
+                          .append(1 + 1)
+                          // 1
+                          .append(1)
+                          .toString();
+                  }
+
+                  String testString() {
+                      return "testString";
+                  }
+              }
+              """,
+            """
+              class A {
+                  void foo() {
+                      String scenarioOne =
+                          // A
+                          "A" +
+                          // B
+                          "B" +
+                          // C
+                          "C";
+                      String scenarioTwo = "A" +
+                          // B
+                          "B" +
+                          // C
+                          "C";
+                      String scenarioThree =
+                          // A
+                          "A" +
+                          // B
+                          "B" +
+                          // C
+                          "C" +
+                          // test method
+                          testString();
+                      String scenarioFour =
+                          // 1 + 1
+                          String.valueOf(1 + 1) +
+                          // 1
+                          1;
+                      String scenarioFive =
+                          // A
+                          "A" +
+                          // 1 + 1
+                          (1 + 1) +
+                          // 1
+                          1;
+                  }
+
+                  String testString() {
+                      return "testString";
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void objectsGrouping() {
         rewriteRun(