Creating a Java Compiler Plugin

1. Overview

Java 8 provides an API for creating Javac plugins. Unfortunately, it’s hard to find good documentation for it.

In this article, we’re going to show the whole process of creating a compiler extension which adds custom code to *.class files.

2. Setup

First, we need to add JDK’s tools.jar as a dependency for our project:

<dependency>
    <groupId>com.sun</groupId>
    <artifactId>tools</artifactId>
    <version>1.8.0</version>
    <scope>system</scope>
    <systemPath>${java.home}/../lib/tools.jar</systemPath>
</dependency>

Every compiler extension is a class which implements com.sun.source.util.Plugin interface. Let’s create it in our example:

Let’s create it in our example:

public class SampleJavacPlugin implements Plugin {

    @Override
    public String getName() {
        return "MyPlugin";
    }

    @Override
    public void init(JavacTask task, String... args) {
        Context context = ((BasicJavacTask) task).getContext();
        Log.instance(context)
          .printRawLines(Log.WriterKind.NOTICE, "Hello from " + getName());
    }
}

For now, we’re just printing “Hello” to ensure that our code is successfully picked up and included in the compilation.

Our end goal will be to create a plugin that adds runtime checks for every numeric argument marked with a given annotation, and throw an exception if the argument doesn’t match a condition.

There’s one more necessary step to make the extension discoverable by Javac: it should be exposed through the ServiceLoader framework.

To achieve this, we need to create a file named com.sun.source.util.Plugin with content which is our plugin’s fully qualified class name (com.baeldung.javac.SampleJavacPlugin) and place it in the META-INF/services directory.

After that, we can call Javac with the -Xplugin:MyPlugin switch:

baeldung/tutorials$ javac -cp ./core-java/target/classes -Xplugin:MyPlugin ./core-java/src/main/java/com/baeldung/javac/TestClass.java
Hello from MyPlugin

Note that we must always use a String returned from the plugin’s getName() method as a -Xplugin option value.

3. Plugin Lifecycle

A plugin is called by the compiler only once, through the init() method.

To be notified of subsequent events, we have to register a callback. These arrive before and after every processing stage per source file:

• PARSE – builds an Abstract Syntax Tree (AST)
• ENTER – source code imports are resolved
• ANALYZE – parser output (an AST) is analyzed for errors
• GENERATE – generating binaries for the target source file

There are two more event kinds – ANNOTATION_PROCESSING and ANNOTATION_PROCESSING_ROUND but we’re not interested in them here.

For example, when we want to enhance compilation by adding some checks based on source code info, it’s reasonable to do that at the PARSE finished event handler:

public void init(JavacTask task, String... args) {
    task.addTaskListener(new TaskListener() {
        public void started(TaskEvent e) {
        }

        public void finished(TaskEvent e) {
            if (e.getKind() != TaskEvent.Kind.PARSE) {
                return;
            }
            // Perform instrumentation
        }
    });
}

4. Extract AST Data

We can get an AST generated by the Java compiler through the TaskEvent.getCompilationUnit(). Its details can be examined through the TreeVisitor interface.

Note that only a Tree element, for which the accept() method is called, dispatches events to the given visitor.

For example, when we execute ClassTree.accept(visitor), only visitClass() is triggered; we can’t expect that, say, visitMethod() is also activated for every method in the given class.

We can use TreeScanner to overcome the problem:

public void finished(TaskEvent e) {
    if (e.getKind() != TaskEvent.Kind.PARSE) {
        return;
    }
    e.getCompilationUnit().accept(new TreeScanner<Void, Void>() {
        @Override
        public Void visitClass(ClassTree node, Void aVoid) {
            return super.visitClass(node, aVoid);

        @Override
        public Void visitMethod(MethodTree node, Void aVoid) {
            return super.visitMethod(node, aVoid);
        }
    }, null);
}

In this example, it’s necessary to call super.visitXxx(node, value) to recursively process the current node’s children.

5. Modify AST

To showcase how we can modify the AST, we’ll insert runtime checks for all numeric arguments marked with a @Positive annotation.

This is a simple annotation that can be applied to method parameters:

@Documented
@Retention(RetentionPolicy.CLASS)
@Target({ElementType.PARAMETER})
public @interface Positive { }

Here’s an example of using the annotation:

public void service(@Positive int i) { }

In the end, we want the bytecode to look as if it’s compiled from a source like this:

public void service(@Positive int i) {
    if (i <= 0) {
        throw new IllegalArgumentException("A non-positive argument ("
          + i + ") is given as a @Positive parameter 'i'");
    }
}

What this means is that we want an IllegalArgumentException to be thrown for every argument marked with @Positive which is equal or less than 0. 

5.1. Where to Instrument

Let’s find out how we can locate target places where the instrumentation should be applied:

private static Set<String> TARGET_TYPES = Stream.of(
  byte.class, short.class, char.class, 
  int.class, long.class, float.class, double.class)
 .map(Class::getName)
 .collect(Collectors.toSet());

For simplicity, we’ve only added primitive numeric types here.

Next, let’s define a shouldInstrument() method that checks if the parameter has a type in the TARGET_TYPES set as well as the @Positive annotation:

private boolean shouldInstrument(VariableTree parameter) {
    return TARGET_TYPES.contains(parameter.getType().toString())
      && parameter.getModifiers().getAnnotations().stream()
      .anyMatch(a -> Positive.class.getSimpleName()
        .equals(a.getAnnotationType().toString()));
}

Then we’ll continue the finished() method in our SampleJavacPlugin class with applying a check to all parameters that fulfill our conditions:

public void finished(TaskEvent e) {
    if (e.getKind() != TaskEvent.Kind.PARSE) {
        return;
    }
    e.getCompilationUnit().accept(new TreeScanner<Void, Void>() {
        @Override
        public Void visitMethod(MethodTree method, Void v) {
            List<VariableTree> parametersToInstrument
              = method.getParameters().stream()
              .filter(SampleJavacPlugin.this::shouldInstrument)
              .collect(Collectors.toList());
            
              if (!parametersToInstrument.isEmpty()) {
                Collections.reverse(parametersToInstrument);
                parametersToInstrument.forEach(p -> addCheck(method, p, context));
            }
            return super.visitMethod(method, v);
        }
    }, null);

In this example, we’ve reversed the parameters list because there’s a possible case that more than one argument is marked by @Positive. As every check is added as the very first method instruction, we process them RTL to ensure correct order.

As every check is added as the very first method instruction, we process them RTL to ensure correct order.

5.2. How to Instrument

The problem is that “read AST” lays in the public API area, while “modify AST” operations like “add null-checks” are a private API.

To address this, we’ll create new AST elements through a TreeMaker instance.

First, we need to obtain a Context instance:

@Override
public void init(JavacTask task, String... args) {
    Context context = ((BasicJavacTask) task).getContext();
    // ...
}

Then, we can obtain the TreeMarker object through the TreeMarker.instance(Context) method.

Now we can build new AST elements, e.g., an if expression can be constructed by a call to TreeMaker.If():

private static JCTree.JCIf createCheck(VariableTree parameter, Context context) {
    TreeMaker factory = TreeMaker.instance(context);
    Names symbolsTable = Names.instance(context);
        
    return factory.at(((JCTree) parameter).pos)
      .If(factory.Parens(createIfCondition(factory, symbolsTable, parameter)),
        createIfBlock(factory, symbolsTable, parameter), 
        null);
}

Please note that we want to show the correct stack trace line when an exception is thrown from our check. That’s why we adjust the AST factory position before creating new elements through it with factory.at(((JCTree) parameter).pos).

The createIfCondition() method builds the “parameterId < 0″ if condition:

private static JCTree.JCBinary createIfCondition(TreeMaker factory, 
  Names symbolsTable, VariableTree parameter) {
    Name parameterId = symbolsTable.fromString(parameter.getName().toString());
    return factory.Binary(JCTree.Tag.LE, 
      factory.Ident(parameterId), 
      factory.Literal(TypeTag.INT, 0));
}

Next, the createIfBlock() method builds a block that returns an IllegalArgumentException:

private static JCTree.JCBlock createIfBlock(TreeMaker factory, 
  Names symbolsTable, VariableTree parameter) {
    String parameterName = parameter.getName().toString();
    Name parameterId = symbolsTable.fromString(parameterName);
        
    String errorMessagePrefix = String.format(
      "Argument '%s' of type %s is marked by @%s but got '", 
      parameterName, parameter.getType(), Positive.class.getSimpleName());
    String errorMessageSuffix = "' for it";
        
    return factory.Block(0, com.sun.tools.javac.util.List.of(
      factory.Throw(
        factory.NewClass(null, nil(), 
          factory.Ident(symbolsTable.fromString(
            IllegalArgumentException.class.getSimpleName())),
            com.sun.tools.javac.util.List.of(factory.Binary(JCTree.Tag.PLUS, 
            factory.Binary(JCTree.Tag.PLUS, 
              factory.Literal(TypeTag.CLASS, errorMessagePrefix), 
              factory.Ident(parameterId)), 
              factory.Literal(TypeTag.CLASS, errorMessageSuffix))), null))));
}

Now that we’re able to build new AST elements, we need to insert them into the AST prepared by the parser. We can achieve this by casting public API elements to private API types:

private void addCheck(MethodTree method, VariableTree parameter, Context context) {
    JCTree.JCIf check = createCheck(parameter, context);
    JCTree.JCBlock body = (JCTree.JCBlock) method.getBody();
    body.stats = body.stats.prepend(check);
}

6. Testing the Plugin

We need to be able to test our plugin. It involves the following:

• compile the test source
• run the compiled binaries and ensure that they behave as expected

For this, we need to introduce a few auxiliary classes.

SimpleSourceFile exposes the given source file’s text to the Javac:

public class SimpleSourceFile extends SimpleJavaFileObject {
    private String content;

    public SimpleSourceFile(String qualifiedClassName, String testSource) {
        super(URI.create(String.format(
          "file://%s%s", qualifiedClassName.replaceAll("\\.", "/"),
          Kind.SOURCE.extension)), Kind.SOURCE);
        content = testSource;
    }

    @Override
    public CharSequence getCharContent(boolean ignoreEncodingErrors) {
        return content;
    }
}

SimpleClassFile holds the compilation result as a byte array:

public class SimpleClassFile extends SimpleJavaFileObject {

    private ByteArrayOutputStream out;

    public SimpleClassFile(URI uri) {
        super(uri, Kind.CLASS);
    }

    @Override
    public OutputStream openOutputStream() throws IOException {
        return out = new ByteArrayOutputStream();
    }

    public byte[] getCompiledBinaries() {
        return out.toByteArray();
    }

    // getters
}

SimpleFileManager ensures the compiler uses our bytecode holder:

public class SimpleFileManager
  extends ForwardingJavaFileManager<StandardJavaFileManager> {

    private List<SimpleClassFile> compiled = new ArrayList<>();

    // standard constructors/getters

    @Override
    public JavaFileObject getJavaFileForOutput(Location location,
      String className, JavaFileObject.Kind kind, FileObject sibling) {
        SimpleClassFile result = new SimpleClassFile(
          URI.create("string://" + className));
        compiled.add(result);
        return result;
    }

    public List<SimpleClassFile> getCompiled() {
        return compiled;
    }
}

Finally, all of that is bound to the in-memory compilation:

public class TestCompiler {
    public byte[] compile(String qualifiedClassName, String testSource) {
        StringWriter output = new StringWriter();

        JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
        SimpleFileManager fileManager = new SimpleFileManager(
          compiler.getStandardFileManager(null, null, null));
        List<SimpleSourceFile> compilationUnits 
          = singletonList(new SimpleSourceFile(qualifiedClassName, testSource));
        List<String> arguments = new ArrayList<>();
        arguments.addAll(asList("-classpath", System.getProperty("java.class.path"),
          "-Xplugin:" + SampleJavacPlugin.NAME));
        JavaCompiler.CompilationTask task 
          = compiler.getTask(output, fileManager, null, arguments, null,
          compilationUnits);
        
        task.call();
        return fileManager.getCompiled().iterator().next().getCompiledBinaries();
    }
}

After that, we need only to run the binaries:

public class TestRunner {

    public Object run(byte[] byteCode, String qualifiedClassName, String methodName,
      Class<?>[] argumentTypes, Object... args) throws Throwable {
        ClassLoader classLoader = new ClassLoader() {
            @Override
            protected Class<?> findClass(String name) throws ClassNotFoundException {
                return defineClass(name, byteCode, 0, byteCode.length);
            }
        };
        Class<?> clazz;
        try {
            clazz = classLoader.loadClass(qualifiedClassName);
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("Can't load compiled test class", e);
        }

        Method method;
        try {
            method = clazz.getMethod(methodName, argumentTypes);
        } catch (NoSuchMethodException e) {
            throw new RuntimeException(
              "Can't find the 'main()' method in the compiled test class", e);
        }

        try {
            return method.invoke(null, args);
        } catch (InvocationTargetException e) {
            throw e.getCause();
        }
    }
}

A test might look like this:

public class SampleJavacPluginTest {

    private static final String CLASS_TEMPLATE
      = "package com.baeldung.javac;\n\n" +
        "public class Test {\n" +
        "    public static %1$s service(@Positive %1$s i) {\n" +
        "        return i;\n" +
        "    }\n" +
        "}\n" +
        "";

    private TestCompiler compiler = new TestCompiler();
    private TestRunner runner = new TestRunner();

    @Test(expected = IllegalArgumentException.class)
    public void givenInt_whenNegative_thenThrowsException() throws Throwable {
        compileAndRun(double.class,-1);
    }
    
    private Object compileAndRun(Class<?> argumentType, Object argument) 
      throws Throwable {
        String qualifiedClassName = "com.baeldung.javac.Test";
        byte[] byteCode = compiler.compile(qualifiedClassName, 
          String.format(CLASS_TEMPLATE, argumentType.getName()));
        return runner.run(byteCode, qualifiedClassName, 
        "service", new Class[] {argumentType}, argument);
    }
}

Here we’re compiling a Test class with a service() method that has a parameter annotated with @Positive. Then, we’re running the Test class by setting a double value of -1 for the method parameter.

As a result of running the compiler with our plugin, the test will throw an IllegalArgumentException for the negative parameter.

7. Conclusion

In this article, we’ve shown the full process of creating, testing and running a Java Compiler plugin.

The full source code of the examples can be found over on GitHub.