Make JavaTemplate engine extensible

rewrite-java/src/main/java/org/openrewrite/java/JavaParser.java

@@ -347,6 +347,15 @@ public B classpath(Collection<Path> classpath) {
             return (B) this;
         }
 
+        B addClasspath(Path classpath) {
+            if (this.classpath.isEmpty()) {
+                this.classpath = Collections.singletonList(classpath);
+            } else {
+                this.classpath.add(classpath);
+            }
+            return (B) this;
+        }
+
         public B classpath(String... classpath) {
             this.classpath = dependenciesFromClasspath(classpath);
             return (B) this;

rewrite-java/src/main/java/org/openrewrite/java/JavaTemplate.java

@@ -27,25 +27,53 @@
 import org.openrewrite.java.tree.JavaCoordinates;
 import org.openrewrite.template.SourceTemplate;
 
+import java.io.IOException;
+import java.io.InputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
-import java.util.HashSet;
-import java.util.Set;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
 import java.util.function.Consumer;
 
 @SuppressWarnings("unused")
 public class JavaTemplate implements SourceTemplate<J, JavaCoordinates> {
+    protected static final Path TEMPLATE_CLASSPATH_DIR;
+
+    static {
+        try {
+            TEMPLATE_CLASSPATH_DIR = Files.createTempDirectory("java-template");
+            Path templateDir = Files.createDirectories(TEMPLATE_CLASSPATH_DIR.resolve("org/openrewrite/java/internal/template"));
+            try (InputStream in = JavaTemplateParser.class.getClassLoader().getResourceAsStream("org/openrewrite/java/internal/template/__M__.class")) {
+                Files.copy(in, templateDir.resolve("__M__.class"));
+            }
+            try (InputStream in = JavaTemplateParser.class.getClassLoader().getResourceAsStream("org/openrewrite/java/internal/template/__P__.class")) {
+                Files.copy(in, templateDir.resolve("__P__.class"));
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private final String code;
     private final int parameterCount;
     private final Consumer<String> onAfterVariableSubstitution;
     private final JavaTemplateParser templateParser;
 
-    private JavaTemplate(boolean contextSensitive, JavaParser.Builder<?, ?> javaParser, String code, Set<String> imports,
+    private JavaTemplate(boolean contextSensitive, JavaParser.Builder<?, ?> parser, String code, Set<String> imports,
                          Consumer<String> onAfterVariableSubstitution, Consumer<String> onBeforeParseTemplate) {
+        this(code, StringUtils.countOccurrences(code, "#{"), onAfterVariableSubstitution, new JavaTemplateParser(contextSensitive, augmentClasspath(parser), onAfterVariableSubstitution, onBeforeParseTemplate, imports));
+    }
+
+    private static JavaParser.Builder<?,?> augmentClasspath(JavaParser.Builder<?,?> parserBuilder) {
+        return parserBuilder.addClasspath(TEMPLATE_CLASSPATH_DIR);
+    }
+
+    protected JavaTemplate(String code, int parameterCount, Consumer<String> onAfterVariableSubstitution, JavaTemplateParser templateParser) {
         this.code = code;
+        this.parameterCount = parameterCount;
         this.onAfterVariableSubstitution = onAfterVariableSubstitution;
-        this.parameterCount = StringUtils.countOccurrences(code, "#{");
-        this.templateParser = new JavaTemplateParser(contextSensitive, javaParser, onAfterVariableSubstitution, onBeforeParseTemplate, imports);
+        this.templateParser = templateParser;
     }
 
     public String getCode() {
@@ -63,7 +91,7 @@ public <J2 extends J> J2 apply(Cursor scope, JavaCoordinates coordinates, Object
             throw new IllegalArgumentException("This template requires " + parameterCount + " parameters.");
         }
 
-        Substitutions substitutions = new Substitutions(code, parameters);
+        Substitutions substitutions = substitutions(parameters);
         String substitutedTemplate = substitutions.substitute();
         onAfterVariableSubstitution.accept(substitutedTemplate);
 
@@ -73,6 +101,10 @@ public <J2 extends J> J2 apply(Cursor scope, JavaCoordinates coordinates, Object
                 .visit(scope.getValue(), 0, scope.getParentOrThrow());
     }
 
+    protected Substitutions substitutions(Object[] parameters) {
+        return new Substitutions(code, parameters);
+    }
+
     @Incubating(since = "8.0.0")
     public static boolean matches(String template, Cursor cursor) {
         return JavaTemplate.builder(template).build().matches(cursor);
@@ -126,14 +158,14 @@ public static class Builder {
 
         private boolean contextSensitive;
 
-        private JavaParser.Builder<?, ?> javaParser = JavaParser.fromJavaVersion();
+        private JavaParser.Builder<?, ?> parser = JavaParser.fromJavaVersion();
 
         private Consumer<String> onAfterVariableSubstitution = s -> {
         };
         private Consumer<String> onBeforeParseTemplate = s -> {
         };
 
-        Builder(String code) {
+        protected Builder(String code) {
             this.code = code.trim();
         }
 
@@ -168,8 +200,8 @@ private void validateImport(String typeName) {
             }
         }
 
-        public Builder javaParser(JavaParser.Builder<?, ?> javaParser) {
-            this.javaParser = javaParser;
+        public Builder javaParser(JavaParser.Builder<?, ?> parser) {
+            this.parser = parser.clone();
             return this;
         }
 
@@ -184,7 +216,7 @@ public Builder doBeforeParseTemplate(Consumer<String> beforeParseTemplate) {
         }
 
         public JavaTemplate build() {
-            return new JavaTemplate(contextSensitive, javaParser, code, imports,
+            return new JavaTemplate(contextSensitive, parser, code, imports,
                     onAfterVariableSubstitution, onBeforeParseTemplate);
         }
     }

rewrite-java/src/main/java/org/openrewrite/java/ReplaceConstantWithAnotherConstant.java

@@ -94,13 +94,17 @@ private J replaceFieldAccess(Expression fieldAccess, JavaType.Variable fieldType
 
             JavaTemplate.Builder templateBuilder;
             if (fieldAccess instanceof J.Identifier) {
-                maybeAddImport(owningType, newConstantName, false);
-                templateBuilder = JavaTemplate.builder(newConstantName)
-                        .staticImports(owningType + '.' + newConstantName);
+                templateBuilder = JavaTemplate.builder(newConstantName);
+                if (owningType.indexOf('.') != -1) {
+                    templateBuilder.staticImports(owningType + '.' + newConstantName);
+                    maybeAddImport(owningType, newConstantName, false);
+                }
             } else {
-                maybeAddImport(owningType, false);
-                templateBuilder = JavaTemplate.builder(owningType.substring(owningType.lastIndexOf('.') + 1) + '.' + newConstantName)
-                        .imports(owningType);
+                templateBuilder = JavaTemplate.builder(owningType.substring(owningType.lastIndexOf('.') + 1) + '.' + newConstantName);
+                if (owningType.indexOf('.') != -1) {
+                    templateBuilder.imports(owningType);
+                    maybeAddImport(owningType, false);
+                }
             }
 
             return templateBuilder.contextSensitive().build()

rewrite-java/src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java

@@ -45,23 +45,6 @@
 public class BlockStatementTemplateGenerator {
     private static final String TEMPLATE_COMMENT = "__TEMPLATE__";
     private static final String STOP_COMMENT = "__TEMPLATE_STOP__";
-    static final String EXPR_STATEMENT_PARAM = "class __P__ {" +
-                                               "  static native <T> T p();" +
-                                               "  static native <T> T[] arrp();" +
-                                               "  static native boolean booleanp();" +
-                                               "  static native byte bytep();" +
-                                               "  static native char charp();" +
-                                               "  static native double doublep();" +
-                                               "  static native int intp();" +
-                                               "  static native long longp();" +
-                                               "  static native short shortp();" +
-                                               "  static native float floatp();" +
-                                               "}";
-    private static final String METHOD_INVOCATION_STUBS = "class __M__ {" +
-                                                          "  static native Object any(Object o);" +
-                                                          "  static native Object any(java.util.function.Predicate<Boolean> o);" +
-                                                          "  static native <T> Object anyT();" +
-                                                          "}";
 
     private final Set<String> imports;
     private final boolean contextSensitive;
@@ -219,7 +202,7 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
     }
 
     @SuppressWarnings("DataFlowIssue")
-    private void contextFreeTemplate(Cursor cursor, J j, StringBuilder before, StringBuilder after, J insertionPoint, JavaCoordinates.Mode mode) {
+    protected void contextFreeTemplate(Cursor cursor, J j, StringBuilder before, StringBuilder after, J insertionPoint, JavaCoordinates.Mode mode) {
         if (j instanceof J.Lambda) {
             throw new IllegalArgumentException(
                     "Templating a lambda requires a cursor so that it can be properly parsed and type-attributed. " +
@@ -229,10 +212,10 @@ private void contextFreeTemplate(Cursor cursor, J j, StringBuilder before, Strin
                     "Templating a method reference requires a cursor so that it can be properly parsed and type-attributed. " +
                     "Mark this template as context-sensitive by calling JavaTemplate.Builder#contextSensitive().");
         } else if (j instanceof Expression && !(j instanceof J.Assignment)) {
-            before.insert(0, "class Template {{\n");
+            before.insert(0, "class Template {\n");
             before.append("Object o = ");
             after.append(";");
-            after.append("\n}}");
+            after.append("\n}");
         } else if ((j instanceof J.MethodDeclaration || j instanceof J.VariableDeclarations || j instanceof J.Block || j instanceof J.ClassDeclaration)
                    && cursor.getValue() instanceof J.Block
                    && (cursor.getParent().getValue() instanceof J.ClassDeclaration || cursor.getParent().getValue() instanceof J.NewClass)) {
@@ -252,7 +235,7 @@ private void contextFreeTemplate(Cursor cursor, J j, StringBuilder before, Strin
             after.append("\n}}");
         }
 
-        before.insert(0, EXPR_STATEMENT_PARAM + METHOD_INVOCATION_STUBS);
+        before.insert(0, "import org.openrewrite.java.internal.template.__M__;\nimport org.openrewrite.java.internal.template.__P__;\n");
         for (String anImport : imports) {
             before.insert(0, anImport);
         }
@@ -262,7 +245,7 @@ private void contextFreeTemplate(Cursor cursor, J j, StringBuilder before, Strin
     private void contextTemplate(Cursor cursor, J prior, StringBuilder before, StringBuilder after, J insertionPoint, JavaCoordinates.Mode mode) {
         J j = cursor.getValue();
         if (j instanceof JavaSourceFile) {
-            before.insert(0, EXPR_STATEMENT_PARAM + METHOD_INVOCATION_STUBS);
+            before.insert(0, "import org.openrewrite.java.internal.template.__M__;\nimport org.openrewrite.java.internal.template.__P__;\n");
 
             JavaSourceFile cu = (JavaSourceFile) j;
             for (J.Import anImport : cu.getImports()) {

rewrite-java/src/main/java/org/openrewrite/java/internal/template/JavaTemplateParser.java

@@ -21,8 +21,10 @@
 import org.openrewrite.Cursor;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.InMemoryExecutionContext;
+import org.openrewrite.Parser;
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.internal.PropertyPlaceholderHelper;
+import org.openrewrite.internal.lang.NonNull;
 import org.openrewrite.java.JavaParser;
 import org.openrewrite.java.RandomizeIdVisitor;
 import org.openrewrite.java.tree.*;
@@ -49,23 +51,35 @@ public class JavaTemplateParser {
     @Language("java")
     private static final String SUBSTITUTED_ANNOTATION = "@java.lang.annotation.Documented public @interface SubAnnotation { int value(); }";
 
-    private final JavaParser.Builder<?, ?> parser;
+    private final Parser.Builder parser;
     private final Consumer<String> onAfterVariableSubstitution;
     private final Consumer<String> onBeforeParseTemplate;
     private final Set<String> imports;
     private final boolean contextSensitive;
     private final BlockStatementTemplateGenerator statementTemplateGenerator;
     private final AnnotationTemplateGenerator annotationTemplateGenerator;
 
-    public JavaTemplateParser(boolean contextSensitive, JavaParser.Builder<?, ?> parser, Consumer<String> onAfterVariableSubstitution,
+    public JavaTemplateParser(boolean contextSensitive, Parser.Builder parser, Consumer<String> onAfterVariableSubstitution,
                               Consumer<String> onBeforeParseTemplate, Set<String> imports) {
+        this(
+                parser,
+                onAfterVariableSubstitution,
+                onBeforeParseTemplate,
+                imports,
+                contextSensitive,
+                new BlockStatementTemplateGenerator(imports, contextSensitive),
+                new AnnotationTemplateGenerator(imports)
+        );
+    }
+
+    protected JavaTemplateParser(Parser.Builder parser, Consumer<String> onAfterVariableSubstitution, Consumer<String> onBeforeParseTemplate, Set<String> imports, boolean contextSensitive, BlockStatementTemplateGenerator statementTemplateGenerator, AnnotationTemplateGenerator annotationTemplateGenerator) {
         this.parser = parser;
         this.onAfterVariableSubstitution = onAfterVariableSubstitution;
         this.onBeforeParseTemplate = onBeforeParseTemplate;
         this.imports = imports;
         this.contextSensitive = contextSensitive;
-        this.statementTemplateGenerator = new BlockStatementTemplateGenerator(imports, contextSensitive);
-        this.annotationTemplateGenerator = new AnnotationTemplateGenerator(imports);
+        this.statementTemplateGenerator = statementTemplateGenerator;
+        this.annotationTemplateGenerator = annotationTemplateGenerator;
     }
 
     public List<Statement> parseParameters(Cursor cursor, String template) {
@@ -235,7 +249,7 @@ private JavaSourceFile compileTemplate(@Language("java") String stub) {
         ExecutionContext ctx = new InMemoryExecutionContext();
         ctx.putMessage(JavaParser.SKIP_SOURCE_SET_TYPE_GENERATION, true);
         ctx.putMessage(ExecutionContext.REQUIRE_PRINT_EQUALS_INPUT, false);
-        JavaParser jp = parser.clone().build();
+        Parser jp = newParser();
         return (stub.contains("@SubAnnotation") ?
                 jp.reset().parse(ctx, stub, SUBSTITUTED_ANNOTATION) :
                 jp.reset().parse(ctx, stub))
@@ -245,6 +259,11 @@ private JavaSourceFile compileTemplate(@Language("java") String stub) {
                 .orElseThrow(() -> new IllegalArgumentException("Could not parse as Java"));
     }
 
+    @NonNull
+    private Parser newParser() {
+        return parser.build();
+    }
+
     /**
      * Return the result of parsing the stub.
      * Cache the LST elements parsed from stub only if the stub is context free.

rewrite-java/src/main/java/org/openrewrite/java/internal/template/Substitutions.java

@@ -75,20 +75,12 @@ public String substitute() {
                             }
                         }
 
-                        s = "(/*__p" + i + "__*/new ";
-
-                        StringBuilder extraDim = new StringBuilder();
+                        int dimensions = 1;
                         for (; arrayType.getElemType() instanceof JavaType.Array; arrayType = (JavaType.Array) arrayType.getElemType()) {
-                            extraDim.append("[0]");
-                        }
-
-                        if (arrayType.getElemType() instanceof JavaType.Primitive) {
-                            s += ((JavaType.Primitive) arrayType.getElemType()).getKeyword();
-                        } else if (arrayType.getElemType() instanceof JavaType.FullyQualified) {
-                            s += ((JavaType.FullyQualified) arrayType.getElemType()).getFullyQualifiedName().replace("$", ".");
+                            dimensions++;
                         }
 
-                        s += "[0]" + extraDim + ")";
+                        s = "(" + newArrayParameter(arrayType.getElemType(), dimensions, i) + ")";
                     } else if ("any".equals(matcherName)) {
                         String fqn;
 
@@ -114,10 +106,9 @@ public String substitute() {
                         fqn = fqn.replace("$", ".");
 
                         JavaType.Primitive primitive = JavaType.Primitive.fromKeyword(fqn);
-                        s = "__P__." + (primitive == null || primitive.equals(JavaType.Primitive.String) ?
-                                "<" + fqn + ">/*__p" + i + "__*/p()" :
-                                "/*__p" + i + "__*/" + fqn + "p()"
-                        );
+                        s = primitive == null || primitive.equals(JavaType.Primitive.String) ?
+                                newObjectParameter(fqn, i) :
+                                newPrimitiveParameter(fqn, i);
 
                         parameters[i] = ((J) parameter).withPrefix(Space.EMPTY);
                     } else {
@@ -138,6 +129,26 @@ public String substitute() {
         return substituted;
     }
 
+    protected String newObjectParameter(String fqn, int index) {
+        return "__P__." + "<" + fqn + ">/*__p" + index + "__*/p()";
+    }
+    protected String newPrimitiveParameter(String fqn, int index) {
+        return "__P__./*__p" + index + "__*/" + fqn + "p()";
+    }
+
+    protected String newArrayParameter(JavaType elemType, int dimensions, int index) {
+        StringBuilder builder = new StringBuilder("/*__p" + index + "__*/" + "new ");
+        if (elemType instanceof JavaType.Primitive) {
+            builder.append(((JavaType.Primitive) elemType).getKeyword());
+        } else if (elemType instanceof JavaType.FullyQualified) {
+            builder.append(((JavaType.FullyQualified) elemType).getFullyQualifiedName().replace("$", "."));
+        }
+        for (int i = 0; i < dimensions; i++) {
+            builder.append("[0]");
+        }
+        return builder.toString();
+    }
+
     private String getTypeName(@Nullable JavaType type) {
         if (type instanceof JavaType.Parameterized) {
             StringJoiner joiner = new StringJoiner(",", "<", ">");
@@ -288,7 +299,7 @@ private static class ThrowingErrorListener extends BaseErrorListener {
         public void syntaxError(Recognizer<?, ?> recognizer, Object offendingSymbol,
                                 int line, int charPositionInLine, String msg, RecognitionException e) {
             throw new IllegalArgumentException(
-                    String.format("Syntax error at line %d:%d %s.", line, charPositionInLine, msg), e);
+                    java.lang.String.format("Syntax error at line %d:%d %s.", line, charPositionInLine, msg), e);
         }
     }
 }

rewrite-java/src/main/java/org/openrewrite/java/internal/template/__M__.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2023 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.internal.template;
+
+public class __M__ {
+    public static native Object any(Object o);
+
+    public static native Object any(java.util.function.Predicate<Boolean> o);
+
+    public static native <T> Object anyT();
+}