JavaTemplate bug when inserting final var into for-each (#4806)

* JavaTemplate bug when inserting `final var` into for-each * Split variable declarations when they contain stop comment
openrewrite · Dec 27, 2024 · 7444d1c · 7444d1c
1 parent b91499c
commit 7444d1c
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 3 deletions.
diff --git a/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest.java b/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest.java
@@ -72,7 +72,7 @@ void assignmentWithinIfPredicate() {
               @Override
               public J.Assignment visitAssignment(J.Assignment assignment, ExecutionContext ctx) {
                   if ((assignment.getAssignment() instanceof J.Literal) &&
-                      ((J.Literal) assignment.getAssignment()).getValue().equals(1)) {
+                    ((J.Literal) assignment.getAssignment()).getValue().equals(1)) {
                       return JavaTemplate.builder("value = 0")
                         .contextSensitive()
                         .build()
@@ -1351,7 +1351,7 @@ void replaceMethodArgumentsInIfStatementWithoutBraces() {
               public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
                   J.MethodInvocation mi = super.visitMethodInvocation(method, ctx);
                   if (new MethodMatcher("Foo bar(..)").matches(mi) &&
-                      mi.getArguments().get(0) instanceof J.Binary) {
+                    mi.getArguments().get(0) instanceof J.Binary) {
                       return JavaTemplate.builder("\"Hello, {}\", \"World!\"")
                         .contextSensitive()
                         .build()
@@ -1382,4 +1382,47 @@ void foo(boolean condition) {
           )
         );
     }
+
+    @Test
+    void replaceVariableDeclarationWithFinalVar() {
+        rewriteRun(
+          spec -> spec.recipe(toRecipe(() -> new JavaIsoVisitor<>() {
+              @Override
+              public J.VariableDeclarations visitVariableDeclarations(J.VariableDeclarations multiVariable, ExecutionContext ctx) {
+                  J.VariableDeclarations vd = super.visitVariableDeclarations(multiVariable, ctx);
+                  if (TypeUtils.isString(vd.getType()) && "String".equals(((J.Identifier) vd.getTypeExpression()).getSimpleName())) {
+                      JavaCoordinates coordinates = vd.getCoordinates().replace();
+                      return JavaTemplate.builder("final var #{}")
+                        .contextSensitive()
+                        .doBeforeParseTemplate(System.out::println)
+                        .build()
+                        .apply(getCursor(), coordinates, new Object[]{vd.getVariables().get(0).getSimpleName()});
+                  }
+                  return vd;
+              }
+          })),
+          java(
+            """
+              import java.util.List;
+              import java.util.ArrayList;
+              
+              class A {
+                  void bar(List<String> lst) {
+                      for (String s : lst) {}
+                  }
+              }
+              """,
+            """
+              import java.util.List;
+              import java.util.ArrayList;
+              
+              class A {
+                  void bar(List<String> lst) {
+                      for (final var s : lst) {}
+                  }
+              }
+              """
+          )
+        );
+    }
 }
diff --git a/...src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java b/...src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java
@@ -25,6 +25,7 @@
 import org.openrewrite.Cursor;
 import org.openrewrite.SourceFile;
 import org.openrewrite.Tree;
+import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.tree.*;
@@ -409,7 +410,7 @@ private void contextTemplate(Cursor cursor, J prior, StringBuilder before, Strin
         } else if (j instanceof J.ForEachLoop.Control) {
             J.ForEachLoop.Control c = (J.ForEachLoop.Control) j;
             if (referToSameElement(prior, c.getVariable())) {
-                after.append(" = /*" + STOP_COMMENT + "/*").append(c.getIterable().printTrimmed(cursor));
+                after.append(" = /*" + STOP_COMMENT + "*/").append(c.getIterable().printTrimmed(cursor));
             } else if (referToSameElement(prior, c.getIterable())) {
                 before.insert(0, "Object __b" + cursor.getPathAsStream().count() + "__ =");
                 after.append(";");
@@ -818,6 +819,20 @@ public J visitMethodInvocation(J.MethodInvocation method, Integer integer) {
                 }
                 return mi;
             }
+
+            @Override
+            public J visitVariableDeclarations(J.VariableDeclarations multiVariable, Integer integer) {
+                List<J.VariableDeclarations.NamedVariable> variables = multiVariable.getVariables();
+                for (J.VariableDeclarations.NamedVariable variable : variables) {
+                    J.VariableDeclarations.NamedVariable.Padding padding = variable.getPadding();
+                    if (padding.getInitializer() != null && stopCommentExists(padding.getInitializer().getBefore().getComments())) {
+                        // Split the variable declarations at the variable with the `STOP_COMMENT` & trim off initializer
+                        List<J.VariableDeclarations.NamedVariable> vars = variables.subList(0, variables.indexOf(variable) + 1);
+                        return multiVariable.withVariables(ListUtils.mapLast(vars, v -> v.withInitializer(null)));
+                    }
+                }
+                return super.visitVariableDeclarations(multiVariable, integer);
+            }
         }
     }
 }