Can an `ElementVisitor` be used to traverse the statements in the body of a method?

I'm trying to make a custom annotation that checks to see if a certain method is called in a method's body annotated with it. Something like:

@TypeQualifierDefault(ElementType.METHOD)
@Retention(RetentionPolicy.SOURCE)
@interface MutatingMethod {
}

interface Mutable {
  void preMutate();
  void postMutate();
  // other methods
}

And then within a certain Mutable class we would have:

class Structure<T> implements Mutable {
  @MutatingMethod
  void add(T data) {
    preMutate();
    // actual mutation code
    postMutate();
  }
}

I want to be able to get warnings of some sort if the body of a method like add that is annotated with @MutatingMethod does not include calls to preMutate and postMutate. Can an ElementVisitor (javax.lang.model.element.ElementVisitor) be used to traverse the (possibly obfuscated) statements and method calls in the body of a method? If so what would that look like? If not what else can I use?

Just to clarify, I know this is impossible (or more difficult) to accomplish in runtime via bytecode decompilation, so this annotation is meant to only work during compilation via reflection (java.lang.reflect.* and javax.lang.model.*) and is not retained in class files.

You are free to modify the code however you want to get it to work, for example by introducing a new annotation called @MutableType that Structure and any other Mutable types must be annotated with it for this to work.

A cherry on top would be to assert that preMutate is called before postMutate and not after.

It shouldn't matter but I'm using Gradle and the IntelliJ IDEA IDE.

Any help is greatly appreciated; material on this is strangely scarce and/or inadequate on the web. I have been using publicly available sources to learn about this!

Solution 1:

There are two modules,

• java.compiler which contains the API for annotation processors and the simple abstraction you have already discovered.

  The ElementVisitor abstraction does not support digging into the method’s code.

• The jdk.compiler module, containing an extended API originally not considered to be part of the standard API and hence not included in the official API documentation prior to the introduction of the module system.

  This API allows analyzing the syntax tree of the currently compiled source code.

When your starting point is an annotation processor, you should have a ProcessingEnvironment which was given to your init method. Then, you can invoke Trees.instance(ProcessingEnvironment) to get a helper object which has the method getTree(Element) you can use to get the syntax tree element. Then, you can traverse the syntax tree from there.

Most of these classes documented in the JDK 17 API do already exist in earlier versions (you might notice the “since 1.6”) even when not present in the older documentation. But prior to JDK 9 you have to include the lib/tools.jar of the particular JDK into your classpath when compiling the annotation processor.

_{(when writing a modular annotation processor)}

import javax.annotation.processing.Processor;

module anno.proc.example {
    requires jdk.compiler;
    provides Processor with anno.proc.example.MyProcessor;
}

package anno.proc.example;

import java.util.*;

import javax.annotation.processing.*;
import javax.lang.model.element.TypeElement;
import javax.tools.Diagnostic;

import com.sun.source.tree.*;
import com.sun.source.tree.Tree.Kind;
import com.sun.source.util.Trees;

import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;

@SupportedSourceVersion(SourceVersion.RELEASE_17) // adapt when using older version
@SupportedAnnotationTypes(MyProcessor.ANNOTATION_NAME)
public class MyProcessor extends AbstractProcessor {
    static final String ANNOTATION_NAME = "my.example.MutatingMethod";
    static final String REQUIRED_FIRST = "preMutate", REQUIRED_LAST = "postMutate";

    // the inherited method does already store the processingEnv
    // public void init(ProcessingEnvironment processingEnv) {

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        Optional<? extends TypeElement> o = annotations.stream()
            .filter(e -> ANNOTATION_NAME.contentEquals(e.getQualifiedName())).findAny();
        if(!o.isPresent()) return false;
        TypeElement myAnnotation = o.get();

        roundEnv.getElementsAnnotatedWith(myAnnotation).forEach(this::check);

        return true;
    }

    private void check(Element e) {
        Trees trees = Trees.instance(processingEnv);
        Tree tree = trees.getTree(e);
        if(tree.getKind() != Kind.METHOD) { // should not happen as compiler handles @Target
            processingEnv.getMessager()
                .printMessage(Diagnostic.Kind.ERROR, ANNOTATION_NAME + " only allowed at methods", e);
            return;
        }
        MethodTree m = (MethodTree) tree;
        List<? extends StatementTree> statements = m.getBody().getStatements();
        if(statements.isEmpty() || !isRequiredFirst(statements.get(0))) {
            processingEnv.getMessager()
                .printMessage(Diagnostic.Kind.MANDATORY_WARNING,
                    "Mutating method does not start with " + REQUIRED_FIRST + "();", e);
        }
        // open challenges:
        //   - accept a return statement after postMutate();
        //   - allow a try { body } finally { postMutate(); }
        if(statements.isEmpty() || !isRequiredLast(statements.get(statements.size() - 1))) {
            processingEnv.getMessager()
                .printMessage(Diagnostic.Kind.MANDATORY_WARNING,
                    "Mutating method does not end with " + REQUIRED_LAST + "();", e);
        }
    }

    private boolean isRequiredFirst(StatementTree st) {
        return invokes(st, REQUIRED_FIRST);
    }

    private boolean isRequiredLast(StatementTree st) {
        return invokes(st, REQUIRED_LAST);
    }

    // check whether tree is an invocation of a no-arg method of the given name
    private boolean invokes(Tree tree, String method) {
        if(tree.getKind() != Kind.EXPRESSION_STATEMENT) return false;
        tree = ((ExpressionStatementTree)tree).getExpression();
        if(tree.getKind() != Kind.METHOD_INVOCATION) return false;

        MethodInvocationTree i = (MethodInvocationTree)tree;

        if(!i.getArguments().isEmpty()) return false; // not a no-arg method

        ExpressionTree ms = i.getMethodSelect();
        // TODO add support for explicit this.method()
        return ms.getKind() == Kind.IDENTIFIER
                && method.contentEquals(((IdentifierTree)ms).getName());
    }
}