Write an Annotation Processor for Java
Lombok (not too far from Java) is a popular annotation processor in Java used to add certain commonly-repeated code to your project and reduce repetitive “boiler-plate” code put on the front of many classes. The question is how to take advantage of this pattern and do something similar.
Lombok Does It
Lombok is registered as an annotationProcessor in Gradle, as well as being included as compileOnly for access to the annotations. Normally, annotation processors are used only to generate new source files, not modify existing ones, but Lombok manages to do it, so why can’t we?
The trick that Lombok uses is that the classes obtained by the javac compiler to the annotation processors are the actual internal representations of the compiler’s structure, not copies of them. This means that if you cast the objects to the actual types, you can modify the existing code as well as simply adding it.
The Quick Of It
The Gradle
You will need to set up a gradle project to build this. The key elements are you will need the tools.jar from your JDK. We place it in a flat folder called lib, which is just a convenient way to keep it with our project.
Additionally, we use another annotation processor from Google (Auto Service) to register components for us. This can by done manually by making the entries into src/main/resources/META-INF/services, but the code you don’t have to write is “always right”.
plugins {
id "java-library"
// Possibly add 'maven-publish' if you're pushing the jar to a repository
}dependencies {
annotationProcessor "com.google.auto.service:auto-service:1.0.1"
implementation "com.google.auto.service:auto-service:1.0.1"
annotationProcessor files("lib/tools-1.5.0.jar")
api files("lib/tools-1.5.0.jar") // Add as required...
}
The Annotation
You need an annotation to process. This will be used to decorate something else.
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.SOURCE)
public @interface WrapMethod {
boolean recordParams() default true;
boolean recordResult() default true;
}You should set the retention policy to SOURCE (or greater) so the annotation exists during compilation. Also, set the target to whatever you desire to annotation. We annotate a METHOD, so we set it to that. You can also add any other fields you want to this, as we are for recording parameters and the result of the function. Use as you wish.
Additionally, we will use static methods on a utility class to make the code we are injecting simpler. See the repository at the end for the example.
The Processor
So you have the pieces, now you need to write code to make, well, code. This is the most laborious part, as you have to work with javac within the Abstract Syntax Tree (AST) model to create code as if it was originally part of your project (part of the method, in this instance).
The basic steps you will go through are:
- Loop through the supported annotations
- Loop through elements annotated with that each annotation
- Prepare for Injection
- Inject!
First, you need to extend the public class AnnotationProcessor extends {and override a few things (namely, init and process). In init, we are just gathering the objects we need from javac, and in process, we do the work.
Note the annotations that will be processed are defined by the annotation, as well as the source version you will be generating. In our experimentation, the maximum release version was limited to which version of tools.jar you used. YMMV.
@SupportedAnnotationTypes("com.benjaminhoogterp.*")
@SupportedSourceVersion(SourceVersion.RELEASE_11)
@AutoService(Processor.class)
public class AnnotationProcessor extends AbstractProcessor {
private Trees trees;
private TreeMaker make;
private Name.Table names;
private Symtab syms;
private Types types;
private Resolve rs;
private Parser parser;
private Modules modules;
private JavacElements elements; @Override
public synchronized void init(ProcessingEnvironment env) {
super.init(env);
trees = Trees.instance(env);
Context context = ((JavacProcessingEnvironment) env).getContext();
make = TreeMaker.instance(context);
names = Names.instance(context).table;
syms = Symtab.instance(context);
types = Types.instance(context);
rs = Resolve.instance(context);
modules = Modules.instance(context);
elements = JavacElements.instance(context);
} @Override
public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
annotations.forEach(annotation -> { // All supported annos
roundEnv.getElementsAnnotatedWith(annotation).forEach(element -> {
if (annotation.toString().contains("WrapMethod")) {
WrapMethod anno = element.getAnnotation(WrapMethod.class);
if (anno != null) doWrapMethod(element);
}
}
}
}
}
Now comes the part where you actually inject code into any method decorated with our annotation @WrapMethod. There are certainly any number of ways to do this, but this is one way that works…
To do the injection, we implement the function we referenced above. In there, we first get the actual javac method declaration by using the trees object we setup in the init with trees.getTree(element). Note that we must cast all of these, because they are the actual compiler objects, but function returned it as a different type without the access we need.
Next, we will pull the javac annotation definition. Remember, you are dealing with an AST, not actual code. There is a difference between an instance of an annotation, which is a Java concept, and an instance of an AST JCAnnotation, a compiler one. A quick solution to get what we want is to simply loop over all annotations on the method declaration and match our annotation by name (string). There will always be exactly one since our original source is valid java and it is not repetable, and we only would have reached this point if the method was annotated.
private void doWrapMethod(Element element) {
JCTree.JCMethodDecl decl = (JCTree.JCMethodDecl) trees.getTree(element);
List<JCTree.JCAnnotation> gcpTraces = List.nil();
for (JCTree.JCAnnotation anno : decl.mods.annotations) {
if (anno.annotationType.toString().contains("WrapMethod"))
gcpTraces = gcpTraces.append(anno);
}
decl.body = rewriteFunctionForTrace(decl, gcpTraces.get(0));
}You can also get the class that defines this method with the following, but we did not require it in this example:
JCTree.JCClassDecl classDecl = (JCTree.JCClassDecl) trees.getTree(decl.sym.owner);
The AST
Finally, we are ready to rewrite the original function, effectively rewriting the original code with whatever we want. We will do a try with resources, as well as log both parameters passed in, and the return value out. This should be a simple enough example to extend to many other use cases.
We have the javac element (compiler instance representing the function) and the original annotation. We will parse off any attributes we want to use to customize the injection, and then we will wrap the original code. We will be returning a JCBlock element, as this will replace the original body in the calling function (see above).
There are a few things to note as we go into hacking the javac. First, the List object you see is is the javac list, not the java.util.List one. As such, it has an .append() method, but you must do list = list.append(x), because it does not modify the original. Also, because we are dealing with the AST, you must find or create every name (token) you want to use. You will use names.fromString(...) a lot. There are many nuances to hacking the AST directly which is often best done just with trial and error. Whenever you’re stuck, see what other methods there are. If you want to inspect your modified code, do a .toString() on the JC object and it will serialize to your java code.
New code elements will typically be created using the make.... functions, for example, make.VarDef(...) , as well as each of the elements passed to the make call. For instance, the modifiers, the symbol, and the initializer, as in the following example.
JCTree resource = make.VarDef(
make.Modifiers(0L), // public, private, final, etc
names.fromString("resource"),
make.Type(findSymbol("com.example.model.MyResource").type),
make.App(make.QualIdent(findSymbol("com.example.util.MyUtil.", "createResource()")))This will create the equivalent code in the AST for the Java MyResource resource = com.example.util.MyUtil.createResource();. Many other combinations are possible, but considering how lengthy the AST creation process is, making a no-args public static method certainly was easiest.
Note, we use a different make call for each type of resource. make.App does a no args invocation of a method. make.QualIdent yields a qualified identifier (token). make.Type gives a type declaration. Every type of object that can be represented in the AST can be created, you just have to think about it in terms of what type of symbols the compiler needs to represent a given piece of Java code as.
It will also probably help to have a few helper functions such as these, because this gets repetitive. This was used in the initializer, above.
private Symbol findSymbol(String className) {
Symbol classSymbol = elements.getTypeElement(className);if (classSymbol == null) throw new IllegalStateException("findSymbol: couldn't find symbol " + className);
return classSymbol;
}
private Symbol findSymbol(String className, String symbolToString) {
Symbol classSymbol = findSymbol(className);
for (Symbol symbol : classSymbol.getEnclosedElements()) {
if (symbolToString.equals(symbol.toString())) return symbol;
}
throw new IllegalStateException("findSymbol: couldn't find symbol " + className + "." + symbolToString);
}The Injection
After lengthy prep, we can actually try to perform some injection.
We parse out the fields of our annotation, which we use later to customize our injection based on the annotation values. They come in as strings this way, which means we must compare them to literals to get booleans. You get the idea.
Next, we store our old body. Since we are doing a try with resources, we will be placing this JCBlock as the body of the try {}. The JCBlock contains a list of statements. We will loop over this and place it into a new JCBlock which contains additional code. Think of a JCBlock as anything that {} contains. Then, we create our resource definition, first by creating the initializer, and then defining our resource as type MyResource and named resource.
@SupportedAnnotationTypes("com.benjaminhoogterp.*")
@SupportedSourceVersion(SourceVersion.RELEASE_11)
@AutoService(Processor.class)
public class AnnotationProcessor extends AbstractProcessor {
private Trees trees;
private TreeMaker make;
private Name.Table names;
private Symtab syms;
private Types types;
private Resolve rs;
private Parser parser;
private Modules modules;
private JavacElements elements;
@Override
public synchronized void init(ProcessingEnvironment env) {
super.init(env);
trees = Trees.instance(env);
Context context = ((JavacProcessingEnvironment) env).getContext();
make = TreeMaker.instance(context);
names = Names.instance(context).table;
syms = Symtab.instance(context);
types = Types.instance(context);
rs = Resolve.instance(context);
modules = Modules.instance(context);
elements = JavacElements.instance(context);
}
@Override
public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
annotations.forEach(annotation -> { // Iterate over annotations that this processor supports
roundEnv.getElementsAnnotatedWith(annotation).forEach(element -> {
if (annotation.toString().contains("WrapMethod")) {
WrapMethod anno = element.getAnnotation(WrapMethod.class);
if (anno != null) doWrapMethod(element);
}
});
});
return true;
}
private void doWrapMethod(Element element) {
JCTree.JCMethodDecl decl = (JCTree.JCMethodDecl) trees.getTree(element);
JCTree.JCClassDecl classDecl = (JCTree.JCClassDecl) trees.getTree(decl.sym.owner);
List<JCTree.JCAnnotation> wrapMethod = List.nil();
for (JCTree.JCAnnotation anno : decl.mods.annotations) {
if (anno.annotationType.toString().contains("WrapMethod"))
wrapMethod = wrapMethod.append(anno);
}
decl.body = rewriteFunctionForWrap(decl, wrapMethod.get(0));
// System.out.println(decl.body);
}
private JCTree.JCBlock rewriteFunctionForWrap(JCTree.JCMethodDecl decl, JCTree.JCAnnotation anno) {
final String trueValue = "true"; // not the hardware store
final String util = "com.benjaminhoogterp.MyUtil";
final String create = "createResource()";
final String returnVarName = "return_";
final String resourceName = "resource";
final String resourceType = "com.benjaminhoogterp.MyResource";
final String recordWrapParams = "recordTraceParams(java.lang.Object...)";
final String recordReturnValue = "recordReturnValue(java.lang.Object)";// NOTE: This prevents compiler assertion errors on the `Bits.incl` method.
// Without this, `.sym` and `.sym.adr` are left at zero and the assertion fails.
// Apparently: "But the newVar() is called by visitVarDef() only if tree.sym is trackable"
// https://stackoverflow.com/questions/46874126/java-lang-assertionerror-thrown-by-compiler-when-adding-generated-method-with-pa
make.at(decl.pos);
Attribute value = anno.attribute.member(names.fromString("recordParams"));
boolean annoDoParams = value == null || value.toString().equals(trueValue);
value = anno.attribute.member(names.fromString("recordResult"));
boolean annoDoResut = value == null || value.toString().equals(trueValue);
JCTree.JCBlock oldBody = decl.body;
JCTree.JCExpression initializer = make.App(make.QualIdent(findSymbol(util, create)));
JCTree resource = make.VarDef(
make.Modifiers(0L),
names.fromString(resourceName),
make.Type(findSymbol(resourceType).type),
initializer
);
List<JCTree.JCExpression> params = List.nil();
params = params.append(make.Ident(names.fromString(resourceName)));
for (JCTree.JCVariableDecl p : decl.getParameters()) {
params = params.append(make.Ident(p));
}
// If enabled, calls the utility with all the function arguments.
JCTree.JCStatement recordParams = make.Exec(make.App(make.QualIdent(findSymbol(util, recordWrapParams)), params));
List<JCTree.JCStatement> statements = List.nil();
if (annoDoParams) statements = statements.append(recordParams);
for (JCTree.JCStatement stmt : oldBody.getStatements()) {
if (!(stmt instanceof JCTree.JCReturn)) {
statements = statements.append(stmt);
} else {
String actualReturnVar = returnVarName + new Random().nextInt(10000);
JCTree.JCReturn ret = (JCTree.JCReturn) stmt;
JCTree.JCExpression exp = ret.getExpression();
JCTree.JCVariableDecl origStmtToVar = make.VarDef(make.Modifiers(0L), names.fromString(actualReturnVar), make.Type(decl.getReturnType().type), exp);
statements = statements.append(origStmtToVar);
// Calls a static method on the util with the value, which allows logging or any other process
JCTree.JCStatement recordResult = make.Exec(make.App(make.QualIdent(findSymbol(util, recordReturnValue)), List.of(make.Ident(names.fromString(actualReturnVar)))));
if (annoDoResut) statements = statements.append(recordResult);
statements = statements.append(make.Return(make.Ident(names.fromString(actualReturnVar))));
}
}
JCTree.JCBlock tryContents = make.Block(0L, statements);
JCTree.JCStatement tryBlock = make.Try(List.of(resource), tryContents, List.nil(), null);
JCTree.JCStatement wrapper = make.Block(0L, List.of(tryBlock));
// This is useful to view the outputted code as the JCStatement serializes to the code itself.
// System.out.print("----------------\nWRAPPER:\n" + wrapper);
return make.Block(0L, List.of(wrapper));
}
private Symbol findSymbol(String className) {
Symbol classSymbol = elements.getTypeElement(className);
if (classSymbol == null) throw new IllegalStateException("findSymbol: couldn't find symbol " + className);
return classSymbol;
}
private Symbol findSymbol(String className, String symbolToString) {
Symbol classSymbol = findSymbol(className);
for (Symbol symbol : classSymbol.getEnclosedElements()) {
if (symbolToString.equals(symbol.toString())) return symbol;
}
throw new IllegalStateException("findSymbol: couldn't find symbol " + className + "." + symbolToString);
}
}
After creating a new list of statements, we construct a new code block, create a try with the resource definition and code block (and no catch statements), optionally print the block for debugging (commented out), and then return the new block. This new block has replaced the complier’s definition of the original method.
The Result
So what do you get out of all of this work?
Consider the following class, which has our annotation processor set up in gradle already (both with annotationProcesssor as well as compileOnly). We simply add our annotation @Wrapmethod onto an existing method, and voila!
public class Demo {
@WrapMethod
public int demo() {
System.out.println("Begin Normal Function handler...");
for(int i = 0; i < 10; i++) {
System.out.println("Some normal operations: " + i);
}
System.out.println("Prepare for return...");
return -1;
}
}If the annotation processor is set up correctly, the function will be re-written as:
try (com.benjaminhoogterp.MyResource resource = com.benjaminhoogterp.MyUtil.createResource();) {
com.benjaminhoogterp.MyUtil.recordTraceParams(resource);
System.out.println("Begin Normal Function handler...");
for (int i = 0; i < 10; i++) {
System.out.println("Some normal operations: " + i);
}
System.out.println("Prepare for return...");
int return_3216 = -1;
com.benjaminhoogterp.MyUtil.recordReturnValue(return_3216);
return return_3216;
}
}The entire function body is wrapped in a try with resources (created by a static call to createResource, the params are logged via recordTraceParams, and the return value is intercepted, logged via recordReturnValue before it is actually returned.
All of the original code is still there, but the annotation processor has detected the annotation and injected code into the existing class.
Caveat
For Java 16 and onward, the Java compiler (javac) has begun restricting much of this internal compiler access. As such, you often have to add several java compiler args to add that access back in. This can be done through gradle, or, as I chose to, through a custom gradle plugin.
Resources
The working repo using Java17 and Gradle 7.3 may be available here:
See Also:
http://scg.unibe.ch/archive/projects/Erni08b.pdf
https://github.com/bhoogter/annotationprocessor/blob/master/Erni08b.pdf (mirror)