It’s finally mature: a framework to define nullability in Java. There have been so many attempts, but JSpecify is here to stay. Now that this has been addressed, another challenge comes to mind: How can we eliminate the boilerplate of handling incoming null arguments?

With the rise of a proper nullability type system, there’s less need to use the Optional type. For example, one might write a null-safe method to uppercase a string value, like this:

public static @Nullable String upper(@Nullable String value) {
    if (value == null) {
        return null;
    }
    return value.toUpperCase();
}

Such a method either returns null or returns the value in uppercase. This pattern is commonly used, especially for utility functions where you don’t control the caller side[1].

When I found myself writing more and more of these methods, I came to the conclusion I needed an annotation to get rid of the boilerplate. Naturally, I checked if the Lombok project had something for this, but sadly there was no annotation that fit my needs[2]. The thing that came closest was the @NonNull annotation, which turns:

public static String upper(@NonNull String value) {
    return value.toUpperCase();
}

into:

public static String upper(@NonNull String value) {
    if (value == null) {
        throw new NullPointerException("value is marked non-null but is null");
    }
    return value.toUpperCase();
}

This was almost exactly what I wanted, except it needed to return null instead of throwing an exception. Luckily, I had already started a project to learn how Lombok modifies the AST with an annotation processor, so I knew how to do it myself[3]. It turned out not to be that hard at all!

We can start by defining an annotation that indicates a formal parameter declaration should cause the enclosing method to return null whenever the argument is null:

@Target(ElementType.PARAMETER)
@Retention(RetentionPolicy.SOURCE)
public @interface NullPropagates {
}

Next, we create the processor boilerplate:

@SupportedSourceVersion(RELEASE_21)
@SupportedAnnotationTypes("org.jdriven.NullPropagates")
public class NullPropagatesProcessor extends AbstractProcessor {
    private Trees treeUtils;      // query information from the AST
    private TreeMaker treeMaker;  // create and modify AST elements

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.treeUtils = instance(processingEnv);
        this.treeMaker = TreeMaker.instance(((JavacProcessingEnvironment) processingEnv).getContext());
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        roundEnv.getElementsAnnotatedWith(NullPropagates.class).stream()
                .forEach( /* do something */ );

        return true;
    }
}

Now for the actual implementation. Because our annotation targets only ElementType.PARAMETER, we know it can only be placed on method and constructor parameters. When processing the annotations, we simply group the annotated parameters by their enclosing method and add null checks at the beginning of each method body. That’s as easy as:

public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
    roundEnv.getElementsAnnotatedWith(NullPropagates.class).stream()
                // The enclosing element of a parameter is its method or constructor.
                .collect(groupingBy(it -> it.getEnclosingElement(), mapping(Function.identity(), toSet())))
                // At this point, each method/constructor is mapped to its @NullPropagates parameters
                .forEach(this::addNullChecks);

    return true;
}

To generate an if block like:

if (<paramName> == null) {
    return null;
}

we can use treeMaker to generate the AST elements:

private JCStatement nullCheck(JCVariableDecl parameter) {
    // Reassign current position to parameter position
    treeMaker.at(parameter.pos);
    // `<paramName> == null`
    var isNull = treeMaker.Binary(EQ, treeMaker.Ident(parameter.getName()), treeMaker.Literal(BOT, null));
    // `return null`
    var returnNull = treeMaker.Return(treeMaker.Literal(BOT, null));
    // { <returnNull> }
    var returnNullBlock = treeMaker.Block(0, of(returnNull));

    // if ( <isNull> ) { <returnNull> }
    return treeMaker.If(isNull, returnNullBlock, null);
}

All that’s left now is to connect the dots:

private void addNullChecks(ExecutableElement method, Set<? extends Element> annotatedParameters) {
    var nullChecks = annotatedParameters.stream()
             // Map each parameter to the javac internal representation
            .map(it -> (JCVariableDecl) treeUtils.getTree(it))
            // ... and create a nullCheck statement
            .map(this::nullCheck)
            .toList();

    // Use the javac internal method representation to add the statements to the method body
    var methodDecl = (JCMethodDecl) treeUtils.getTree(method);
    methodDecl.body.stats = from(nullChecks).appendList(methodDecl.body.stats);
}

And that’s it! Now we can use the annotation as desired:

public static @Nullable String upper(@NullPropagates String value) {
    return value.toUpperCase();
}

And the annotation processor will turn it into:

public static @Nullable String upper(@NullPropagates String value) {
    if (value == null) {
        return null;
    }
    return value.toUpperCase();
}

For a production-ready implementation, we would need to add a few guards: for instance, methods returning void, methods with primitive return types, or primitive parameters should be ignored or flagged with a compiler error. Constructors should also be skipped, because they do not return anything[4]! Adding these guards is a straightforward addition to the core logic shown here.

You can check out the full implementation on GitHub: NullPropagatesProcessor.java.
Happy coding!

1. See, for example, the API of the StringUtils class from the Apache Commons project
2. Interestingly, years ago there was a feature request to introduce a @ReturnNull annotation, whose intended behavior was literally what I am proposing here. The request was rejected because it could lead to null pointer exceptions in unexpected places. That may have been true back in 2018, but now that we have a sound nullability system, I don’t really agree with that premise.
3. There are already many great blog posts out there explaining how annotation processors work. If you are new to the topic, Annotation Processing 101 by Hannes Dorfmann is a great starting point. Annotation processors are officially meant for generating new code, but by tapping into compiler internals, which is Lombok’s core trick, you can also modify the AST.
4. Which conveniently spares us the headache in older Java versions where statements cannot precede a super(…​) or this(…​) invocation. Starting with Java 25, JEP 513: Flexible Constructor Bodies was released to address this issue.
shadow-left