Dealing with even more kinds of trailing comments in UnwrapElseAfterReturn (#655)

Author: greg-at-moderne

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

@@ -23,12 +23,13 @@
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.tree.Comment;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.Space;
 import org.openrewrite.java.tree.Statement;
 
 import java.time.Duration;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class UnwrapElseAfterReturn extends Recipe {
 
@@ -54,7 +55,7 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             @Override
             public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
                 J.Block b = visitAndCast(block, ctx, super::visitBlock);
-                List<Comment> endComments = new ArrayList<>();
+                AtomicReference<Space> endWhitespace = new AtomicReference<>(null);
                 J.Block alteredBlock = b.withStatements(ListUtils.flatMap(b.getStatements(), statement -> {
                     if (statement instanceof J.If) {
                         J.If ifStatement = (J.If) statement;
@@ -63,7 +64,7 @@ public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
                             Statement elsePart = ifStatement.getElsePart().getBody();
                             if (elsePart instanceof J.Block) {
                                 J.Block elseBlock = (J.Block) elsePart;
-                                endComments.addAll(elseBlock.getEnd().getComments());
+                                endWhitespace.set(elseBlock.getEnd());
                                 return ListUtils.concat(newIf, ListUtils.mapFirst(elseBlock.getStatements(), elseStmt -> {
                                     // Combine comments from the else block itself and the first statement
                                     List<Comment> elseComments = elseBlock.getPrefix().getComments();
@@ -81,9 +82,9 @@ public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
                     return statement;
                 }));
 
-                if (!endComments.isEmpty()) {
-                    List<Comment> mergedComments = ListUtils.concatAll(endComments, b.getEnd().getComments());
-                    alteredBlock = alteredBlock.withEnd(b.getEnd().withComments(mergedComments));
+                if (endWhitespace.get() != null) {
+                    List<Comment> mergedComments = ListUtils.concatAll(endWhitespace.get().getComments(), b.getEnd().getComments());
+                    alteredBlock = alteredBlock.withEnd(b.getEnd().withComments(mergedComments).withWhitespace(endWhitespace.get().getWhitespace()));
                 }
 
                 return maybeAutoFormat(b, alteredBlock, ctx);

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

@@ -597,7 +597,7 @@ void process(String data) {
     }
 
     @Test
-    void comments() {
+    void commentsEverywhere() {
         rewriteRun(
           java(
             """
@@ -625,4 +625,64 @@ int foo(boolean condition) {
           )
         );
     }
+
+    @Test
+    void commentsOnlyInBlocks() {
+        rewriteRun(
+          java(
+            """
+            class Test {
+                int foo(boolean condition) {
+                    if (condition) {
+                        return 1; // end 1
+                    } else {
+                        return 2; // end 2
+                    }
+                }
+            }
+            """,
+            """
+              class Test {
+                  int foo(boolean condition) {
+                      if (condition) {
+                          return 1; // end 1
+                      }
+                      return 2; // end 2
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void commentsOnlyInBlocksWithNewLine() {
+        rewriteRun(
+          java(
+            """
+            class Test {
+                int foo(boolean condition) {
+                    if (condition) {
+                        return 1; // end 1
+                    } else {
+                        return 2;
+                        // next line after 2
+                    }
+                }
+            }
+            """,
+            """
+              class Test {
+                  int foo(boolean condition) {
+                      if (condition) {
+                          return 1; // end 1
+                      }
+                      return 2;
+                      // next line after 2
+                  }
+              }
+              """
+          )
+        );
+    }
 }