Fix JavaTemplate expression replacement inside return statements (#7096)

timtebeek · web-flow · commit 76c43213204f · 2026-03-24T11:17:14.000+01:00
* Fix JavaTemplate expression replacement inside return statements

When a method invocation or new class expression is the direct expression
of a J.Return (or other non-block expression context), coordinates.replace()
produces STATEMENT_PREFIX but parseBlockStatements fails because the template
is an expression. Detect expression context by checking if the parent tree
is not a J.Block, and use parseExpression instead.

* Add test for J.NewClass replacement inside return statements
diff --git a/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest8Test.java b/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest8Test.java
@@ -19,6 +19,7 @@
 import org.openrewrite.Cursor;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.ExecutionContext;
+import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.NameTree;
 import org.openrewrite.marker.SearchResult;
 import org.openrewrite.test.RewriteTest;
@@ -28,6 +29,80 @@
 
 class JavaTemplateTest8Test implements RewriteTest {
 
+    @Test
+    void replaceMethodInvocationInsideReturn() {
+        rewriteRun(
+          spec -> spec.recipe(toRecipe(() -> new JavaIsoVisitor<>() {
+              final JavaTemplate t = JavaTemplate.builder("String.valueOf(#{any(String)})")
+                .build();
+
+              @Override
+              public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
+                  method = super.visitMethodInvocation(method, ctx);
+                  if (method.getSimpleName().equals("toString")) {
+                      return t.apply(getCursor(), method.getCoordinates().replace(),
+                        method.getSelect());
+                  }
+                  return method;
+              }
+          })),
+          java(
+            """
+              class Test {
+                  String test(String s) {
+                      return s.toString();
+                  }
+              }
+              """,
+            """
+              class Test {
+                  String test(String s) {
+                      return String.valueOf(s);
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void replaceNewClassInsideReturn() {
+        rewriteRun(
+          spec -> spec.recipe(toRecipe(() -> new JavaIsoVisitor<>() {
+              final JavaTemplate t = JavaTemplate.builder("new StringBuilder(#{any(String)})")
+                .build();
+
+              @Override
+              public J.NewClass visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
+                  newClass = super.visitNewClass(newClass, ctx);
+                  if (newClass.getClazz() != null &&
+                      newClass.getClazz().toString().equals("StringBuffer") &&
+                      newClass.getArguments().size() == 1) {
+                      return t.apply(getCursor(), newClass.getCoordinates().replace(),
+                        newClass.getArguments().get(0));
+                  }
+                  return newClass;
+              }
+          })),
+          java(
+            """
+              class Test {
+                  CharSequence test(String s) {
+                      return new StringBuffer(s);
+                  }
+              }
+              """,
+            """
+              class Test {
+                  CharSequence test(String s) {
+                      return new StringBuilder(s);
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @DocumentExample
     @Test
     void parameterizedMatch() {
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/internal/template/JavaTemplateJavaExtension.java b/rewrite-java/src/main/java/org/openrewrite/java/internal/template/JavaTemplateJavaExtension.java
@@ -435,12 +435,32 @@ public J visitMethodInvocation(J.MethodInvocation method, Integer integer) {
                     }
                     return m;
                 }
+                if (loc == STATEMENT_PREFIX && isScope(method) &&
+                    !(getCursor().getParentTreeCursor().getValue() instanceof J.Block)) {
+                    // Method invocation is used as an expression (e.g., inside return, assignment),
+                    // not as a standalone statement in a block. Parse as expression replacement.
+                    return autoFormat(unsubstitute(templateParser.parseExpression(
+                                    getCursor(),
+                                    substitutedTemplate,
+                                    substitutions.getTypeVariables(),
+                                    loc))
+                            .withPrefix(method.getPrefix()), integer);
+                }
                 return maybeReplaceStatement(method, J.class, 0);
             }
 
             @Override
             public J visitNewClass(J.NewClass newClass, Integer p) {
                 if (isScope(newClass)) {
+                    if (loc == STATEMENT_PREFIX &&
+                        !(getCursor().getParentTreeCursor().getValue() instanceof J.Block)) {
+                        return autoFormat(unsubstitute(templateParser.parseExpression(
+                                        getCursor(),
+                                        substitutedTemplate,
+                                        substitutions.getTypeVariables(),
+                                        loc))
+                                .withPrefix(newClass.getPrefix()), p);
+                    }
                     // allow a `J.NewClass` to also be replaced by an expression
                     return maybeReplaceStatement(newClass, J.class, p);
                 }
