Implementing a Custom Lombok Annotation

1. Overview

In this tutorial, we’ll implement a custom annotation using Lombok to remove the boiler-plate around implementing Singletons in an application.

Lombok is a powerful Java library that aims to reduce the boiler-plate code in Java. If you’re not familiar with it, here you can find the introduction to all the features of Lombok.

2. Lombok as an Annotation Processor

Java allows application developers to process annotations during the compilation phase; most importantly, to generate new files based on an annotation. As a result, libraries like Hibernate allow developers to reduce the boiler-plate code and use annotations instead.

Annotation processing is covered in depth in this tutorial.

In the same way, Project Lombok also works as an Annotation Processor. It processes the annotation by delegating it to a specific handler.

When delegating, it sends the compiler’s Abstract Syntax Tree (AST) of the annotated code to the handler. Therefore, it allows the handlers to modify the code by extending the AST.

3. Implementing a Custom Annotation

3.1. Extending Lombok

Surprisingly, Lombok is not easy to extend and add a custom annotation.

In fact, the newer versions of Lombok use Shadow ClassLoader (SCL) to hide the .class files in Lombok as .scl files. Thus, it forces the developers to fork the Lombok source code and implement annotations there.

On the positive side, it simplifies the process of extending custom handlers and AST modification using utility functions.

3.2. Singleton Annotation

Generally, a lot of code is required for implementing a Singleton class. For applications that don’t use a dependency injection framework, this is just boilerplate stuff.

For instance, here’s one way of implementing a Singleton class:

public class SingletonRegistry {
    private SingletonRegistry() {}
    
    private static class SingletonRegistryHolder {
        private static SingletonRegistry registry = new SingletonRegistry();
    }
    
    public static SingletonRegistry getInstance() {
        return SingletonRegistryHolder.registry;
    }
	
    // other methods
}

In contrast, this is how it would look like if we implement an annotation version of it:

@Singleton
public class SingletonRegistry {}

And, the Singleton annotation :

@Target(ElementType.TYPE)
public @interface Singleton {}

It is important to emphasize here that a Lombok Singleton handler would generate the implementation code we saw above by modifying the AST.

Since the AST is different for every compiler, a custom Lombok handler is needed for each. Lombok allows custom handlers for javac (used by Maven/Gradle and Netbeans) and the Eclipse compiler.

In the following sections, we’ll implement our Annotation handler for each compiler.

4. Implementing a Handler for javac

4.1. Maven Dependency

Let’s pull the required dependencies for Lombok first:

<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
    <version>1.18.4</version>
</dependency>

Additionally, we would also need the tools.jar shipped with Java for accessing and modifying the javac AST. However, there is no Maven repository for it. The easiest way to include this in a Maven project is by adding it to Profile:

<profiles>
    <profile>
        <id>default-tools.jar</id>
            <activation>
                <property>
                    <name>java.vendor</name>
                    <value>Oracle Corporation</value>
                </property>
            </activation>
            <dependencies>
                <dependency>
                    <groupId>com.sun</groupId>
                    <artifactId>tools</artifactId>
                    <version>${java.version}</version>
                    <scope>system</scope>
                    <systemPath>${java.home}/../lib/tools.jar</systemPath>
                </dependency>
            </dependencies>
    </profile>
</profiles>

4.2. Extending JavacAnnotationHandler

In order to implement a custom javac handler, we need to extend Lombok’s JavacAnnotationHandler:

public class SingletonJavacHandler extends JavacAnnotationHandler<Singleton> {
    public void handle(
      AnnotationValues<Singleton> annotation,
      JCTree.JCAnnotation ast,
      JavacNode annotationNode) {}
}

Next, we’ll implement the handle() method. Here, the annotation AST is made available as a parameter by Lombok.

4.3. Modifying the AST

This is where things get tricky. Generally, changing an existing AST is not as straightforward.

Luckily, Lombok provides many utility functions in JavacHandlerUtil and JavacTreeMaker for generating code and injecting it in the AST. With this in mind, let’s use these functions and create the code for our SingletonRegistry:

public void handle(
  AnnotationValues<Singleton> annotation,
  JCTree.JCAnnotation ast,
  JavacNode annotationNode) {
    Context context = annotationNode.getContext();
    Javac8BasedLombokOptions options = Javac8BasedLombokOptions
      .replaceWithDelombokOptions(context);
    options.deleteLombokAnnotations();
    JavacHandlerUtil
      .deleteAnnotationIfNeccessary(annotationNode, Singleton.class);
    JavacHandlerUtil
      .deleteImportFromCompilationUnit(annotationNode, "lombok.AccessLevel");
    JavacNode singletonClass = annotationNode.up();
    JavacTreeMaker singletonClassTreeMaker = singletonClass.getTreeMaker();
    addPrivateConstructor(singletonClass, singletonClassTreeMaker);

    JavacNode holderInnerClass = addInnerClass(singletonClass, singletonClassTreeMaker);
    addInstanceVar(singletonClass, singletonClassTreeMaker, holderInnerClass);
    addFactoryMethod(singletonClass, singletonClassTreeMaker, holderInnerClass);
}

It’s important to point out that the deleteAnnotationIfNeccessary() and the deleteImportFromCompilationUnit() methods provided by Lombok are used for removing annotations and any imports for them.

Now, let’s see how other private methods are implemented for generating the code. First, we’ll generate the private constructor:

private void addPrivateConstructor(
  JavacNode singletonClass,
  JavacTreeMaker singletonTM) {
    JCTree.JCModifiers modifiers = singletonTM.Modifiers(Flags.PRIVATE);
    JCTree.JCBlock block = singletonTM.Block(0L, nil());
    JCTree.JCMethodDecl constructor = singletonTM
      .MethodDef(
        modifiers,
        singletonClass.toName("<init>"),
        null, nil(), nil(), nil(), block, null);

    JavacHandlerUtil.injectMethod(singletonClass, constructor);
}

Next, the inner SingletonHolder class:

private JavacNode addInnerClass(
  JavacNode singletonClass,
  JavacTreeMaker singletonTM) {
    JCTree.JCModifiers modifiers = singletonTM
      .Modifiers(Flags.PRIVATE | Flags.STATIC);
    String innerClassName = singletonClass.getName() + "Holder";
    JCTree.JCClassDecl innerClassDecl = singletonTM
      .ClassDef(modifiers, singletonClass.toName(innerClassName),
      nil(), null, nil(), nil());
    return JavacHandlerUtil.injectType(singletonClass, innerClassDecl);
}

Now, we’ll add an instance variable in the holder class:

private void addInstanceVar(
  JavacNode singletonClass,
  JavacTreeMaker singletonClassTM,
  JavacNode holderClass) {
    JCTree.JCModifiers fieldMod = singletonClassTM
      .Modifiers(Flags.PRIVATE | Flags.STATIC | Flags.FINAL);

    JCTree.JCClassDecl singletonClassDecl
      = (JCTree.JCClassDecl) singletonClass.get();
    JCTree.JCIdent singletonClassType
      = singletonClassTM.Ident(singletonClassDecl.name);

    JCTree.JCNewClass newKeyword = singletonClassTM
      .NewClass(null, nil(), singletonClassType, nil(), null);

    JCTree.JCVariableDecl instanceVar = singletonClassTM
      .VarDef(
        fieldMod,
        singletonClass.toName("INSTANCE"),
        singletonClassType,
        newKeyword);
    JavacHandlerUtil.injectField(holderClass, instanceVar);
}

Finally, let’s add a factory method for accessing the singleton object:

private void addFactoryMethod(
  JavacNode singletonClass,
  JavacTreeMaker singletonClassTreeMaker,
  JavacNode holderInnerClass) {
    JCTree.JCModifiers modifiers = singletonClassTreeMaker
      .Modifiers(Flags.PUBLIC | Flags.STATIC);

    JCTree.JCClassDecl singletonClassDecl
      = (JCTree.JCClassDecl) singletonClass.get();
    JCTree.JCIdent singletonClassType
      = singletonClassTreeMaker.Ident(singletonClassDecl.name);

    JCTree.JCBlock block
      = addReturnBlock(singletonClassTreeMaker, holderInnerClass);

    JCTree.JCMethodDecl factoryMethod = singletonClassTreeMaker
      .MethodDef(
        modifiers,
        singletonClass.toName("getInstance"),
        singletonClassType, nil(), nil(), nil(), block, null);
    JavacHandlerUtil.injectMethod(singletonClass, factoryMethod);
}

Clearly, the factory method returns the instance variable from the holder class. Let’s implement that as well:

private JCTree.JCBlock addReturnBlock(
  JavacTreeMaker singletonClassTreeMaker,
  JavacNode holderInnerClass) {

    JCTree.JCClassDecl holderInnerClassDecl
      = (JCTree.JCClassDecl) holderInnerClass.get();
    JavacTreeMaker holderInnerClassTreeMaker
      = holderInnerClass.getTreeMaker();
    JCTree.JCIdent holderInnerClassType
      = holderInnerClassTreeMaker.Ident(holderInnerClassDecl.name);

    JCTree.JCFieldAccess instanceVarAccess = holderInnerClassTreeMaker
      .Select(holderInnerClassType, holderInnerClass.toName("INSTANCE"));
    JCTree.JCReturn returnValue = singletonClassTreeMaker
      .Return(instanceVarAccess);

    ListBuffer<JCTree.JCStatement> statements = new ListBuffer<>();
    statements.append(returnValue);

    return singletonClassTreeMaker.Block(0L, statements.toList());
}

As a result, we have the modified AST for our Singleton class.

4.4. Registering Handler with SPI

Until now, we only implemented a Lombok handler for generating an AST for our SingletonRegistry. Here, it’s important to repeat that Lombok works as an annotation processor.

Usually, Annotation Processors are discovered via META-INF/services. Lombok also maintains a list of handlers in the same way. Additionally, it uses a framework named SPI for automatically updating the handler list.

For our purpose, we’ll use the metainf-services:

<dependency>
    <groupId>org.kohsuke.metainf-services</groupId>
    <artifactId>metainf-services</artifactId>
    <version>1.8</version>
</dependency>

Now, we can register our handler with Lombok:

@MetaInfServices(JavacAnnotationHandler.class)
public class SingletonJavacHandler extends JavacAnnotationHandler<Singleton> {}

This will generate a lombok.javac.JavacAnnotationHandler file at compile time. This behavior is common for all SPI frameworks.

5. Implementing a Handler for Eclipse IDE

5.1. Maven Dependency

Similar to tools.jar we added for accessing the AST for javac, we’ll add eclipse jdt for Eclipse IDE:

<dependency>
    <groupId>org.eclipse.jdt</groupId>
    <artifactId>core</artifactId>
    <version>3.3.0-v_771</version>
</dependency>

5.2. Extending EclipseAnnotationHandler

We’ll now extend EclipseAnnotationHandler for our Eclipse handler:

@MetaInfServices(EclipseAnnotationHandler.class)
public class SingletonEclipseHandler
  extends EclipseAnnotationHandler<Singleton> {
    public void handle(
      AnnotationValues<Singleton> annotation,
      Annotation ast,
      EclipseNode annotationNode) {}
}

Together with the SPI annotation, MetaInfServices, this handler acts as a processor for our Singleton annotation. Hence, whenever a class is compiled in Eclipse IDE, the handler converts the annotated class into a singleton implementation.

5.3. Modifying AST

With our handler registered with SPI, we can now start editing the AST for Eclipse compiler:

public void handle(
  AnnotationValues<Singleton> annotation,
  Annotation ast,
  EclipseNode annotationNode) {
    EclipseHandlerUtil
      .unboxAndRemoveAnnotationParameter(
        ast,
        "onType",
        "@Singleton(onType=", annotationNode);
    EclipseNode singletonClass = annotationNode.up();
    TypeDeclaration singletonClassType
      = (TypeDeclaration) singletonClass.get();
    
    ConstructorDeclaration constructor
      = addConstructor(singletonClass, singletonClassType);
    
    TypeReference singletonTypeRef 
      = EclipseHandlerUtil.cloneSelfType(singletonClass, singletonClassType);
    
    StringBuilder sb = new StringBuilder();
    sb.append(singletonClass.getName());
    sb.append("Holder");
    String innerClassName = sb.toString();
    TypeDeclaration innerClass
      = new TypeDeclaration(singletonClassType.compilationResult);
    innerClass.modifiers = AccPrivate | AccStatic;
    innerClass.name = innerClassName.toCharArray();
    
    FieldDeclaration instanceVar = addInstanceVar(
      constructor,
      singletonTypeRef,
      innerClass);
    
    FieldDeclaration[] declarations = new FieldDeclaration[]{instanceVar};
    innerClass.fields = declarations;
    
    EclipseHandlerUtil.injectType(singletonClass, innerClass);
    
    addFactoryMethod(
      singletonClass,
      singletonClassType,
      singletonTypeRef,
      innerClass,
      instanceVar);
}

Next, the private constructor:

private ConstructorDeclaration addConstructor(
  EclipseNode singletonClass,
  TypeDeclaration astNode) {
    ConstructorDeclaration constructor
      = new ConstructorDeclaration(astNode.compilationResult);
    constructor.modifiers = AccPrivate;
    constructor.selector = astNode.name;
    
    EclipseHandlerUtil.injectMethod(singletonClass, constructor);
    return constructor;
}

And for the instance variable:

private FieldDeclaration addInstanceVar(
  ConstructorDeclaration constructor,
  TypeReference typeReference,
  TypeDeclaration innerClass) {
    FieldDeclaration field = new FieldDeclaration();
    field.modifiers = AccPrivate | AccStatic | AccFinal;
    field.name = "INSTANCE".toCharArray();
    field.type = typeReference;
    
    AllocationExpression exp = new AllocationExpression();
    exp.type = typeReference;
    exp.binding = constructor.binding;
    
    field.initialization = exp;
    return field;
}

Lastly, the factory method:

private void addFactoryMethod(
  EclipseNode singletonClass,
  TypeDeclaration astNode,
  TypeReference typeReference,
  TypeDeclaration innerClass,
  FieldDeclaration field) {
    
    MethodDeclaration factoryMethod
      = new MethodDeclaration(astNode.compilationResult);
    factoryMethod.modifiers 
      = AccStatic | ClassFileConstants.AccPublic;
    factoryMethod.returnType = typeReference;
    factoryMethod.sourceStart = astNode.sourceStart;
    factoryMethod.sourceEnd = astNode.sourceEnd;
    factoryMethod.selector = "getInstance".toCharArray();
    factoryMethod.bits = ECLIPSE_DO_NOT_TOUCH_FLAG;
    
    long pS = factoryMethod.sourceStart;
    long pE = factoryMethod.sourceEnd;
    long p = (long) pS << 32 | pE;
    
    FieldReference ref = new FieldReference(field.name, p);
    ref.receiver = new SingleNameReference(innerClass.name, p);
    
    ReturnStatement statement
      = new ReturnStatement(ref, astNode.sourceStart, astNode.sourceEnd);
    
    factoryMethod.statements = new Statement[]{statement};
    
    EclipseHandlerUtil.injectMethod(singletonClass, factoryMethod);
}

Additionally, we must plug this handler into the Eclipse boot classpath. Generally, it is done by adding the following parameter to the eclipse.ini:

-Xbootclasspath/a:singleton-1.0-SNAPSHOT.jar

6. Custom Annotation in IntelliJ

Generally speaking, a new Lombok handler is needed for every compiler, like the javac and Eclipse handlers that we implemented before.

Conversely, IntelliJ doesn’t support Lombok handler. It provides Lombok support through a plugin instead.

Due to this, any new annotation must be supported by the plugin explicitly. This also applies to any annotation added to Lombok.

7. Conclusion

In this article, we implemented a custom annotation using Lombok handlers. We also briefly looked at AST modification for our Singleton annotation in different compilers, available in various IDEs.

The full source code is available over on Github.