How to resolve SwitchTree Expression type in Java Annotation Processor?

Viewed 231

Let's consider the following code:

switch ( <em>switchTreeExpression</em> ) {
    <em>cases</em>
}

I want to find out, what type for switchTreeExpression is .

I have the following code draft:

...
MethodTree methodTree = trees.getTree(method);
BlockTree blockTree = methodTree.getBody();

for (StatementTree statementTree : blockTree.getStatements()) {
    if (statementTree.getKind() == Tree.Kind.SWITCH) {
        SwitchTree switchTree = (SwitchTree) statementTree;
        ExpressionTree switchTreeExpression = switchTree.getExpression();
        // I need to get the type of *switchTreeExpression* here
    }
}

It is interesting, that I can get the type of switchTreeExpression from .class file. However it seems that there is no way to get byte code of the current class in this phase of annotation processing (if I am wrong, I would be happy just get byte code and analyze it with ObjectWeb ASM library).

1 Answers

Possible solutions

Annotation processor

Let's consider an annotation processor for the type annotations (@Target(ElementType.TYPE)).

Limitation: Processor.process() method: No method bodies

Processing Code:

Annotation processing occurs at a specific point in the timeline of a compilation, after all source files and classes specified on the command line have been read, and analyzed for the types and members they contain, but before the contents of any method bodies have been analyzed.

Overcoming limitation: Using com.sun.source.util.TaskListener

The idea is to handle the type element analysis completion events.

  • Processor.init() method: Register a task listener and handle the type element analysis completion events using the captured annotated type elements.
  • Processor.process() method: Capture the annotated type elements.

Some related references:

Note on implementation approaches

Some third-party dependencies (libraries and frameworks) may be used to implement an annotation processor.

For example, the already mentioned Checker Framework.

Some related references:

Please, note that the Checker Framework processors use @SupportedAnnotationTypes("*").

Draft implementation

Let's consider a draft implementation, which does not use third-party dependencies mentioned in the «Note on implementation approaches» section.

Annotation processor project

Maven project
<properties>
    <auto-service.version>1.0.1</auto-service.version>
</properties>
<plugin>
    <groupId>org.apache.maven.plugins</groupId>
    <artifactId>maven-compiler-plugin</artifactId>
    <configuration>
        <annotationProcessorPaths>
            <path>
                <groupId>com.google.auto.service</groupId>
                <artifactId>auto-service</artifactId>
                <version>${auto-service.version}</version>
            </path>
        </annotationProcessorPaths>
    </configuration>
</plugin>
<dependency>
    <groupId>com.google.auto.service</groupId>
    <artifactId>auto-service-annotations</artifactId>
    <version>${auto-service.version}</version>
</dependency>
AbstractTypeProcessor class: Base class

Let's introduce the base class that has the following abstract method:

public abstract void processType(Trees trees, TypeElement typeElement, TreePath treePath);
import com.sun.source.util.JavacTask;
import com.sun.source.util.TaskEvent;
import com.sun.source.util.TaskListener;
import com.sun.source.util.TreePath;
import com.sun.source.util.Trees;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.element.Name;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.ElementFilter;

// NOTE: It is designed to work only with `@Target(ElementType.TYPE)` annotations!
public abstract class AbstractTypeProcessor extends AbstractProcessor {
    private final AnalyzeTaskListener analyzeTaskListener = new AnalyzeTaskListener(this);
    protected final Set<Name> remainingTypeElementNames = new HashSet<>();
    private Trees trees;

    protected AbstractTypeProcessor() {
    }

    @Override
    public synchronized void init(final ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        trees = Trees.instance(processingEnv);
        JavacTask.instance(processingEnv).addTaskListener(analyzeTaskListener);
    }

    @Override
    public boolean process(final Set<? extends TypeElement> annotations, final RoundEnvironment roundEnv) {
        for (final TypeElement annotation : annotations) {
            final Set<? extends Element> annotatedElements = roundEnv.getElementsAnnotatedWith(annotation);
            final Set<TypeElement> typeElements = ElementFilter.typesIn(annotatedElements);
            final List<Name> typeElementNames = typeElements.stream()
                .map(TypeElement::getQualifiedName)
                .toList();
            remainingTypeElementNames.addAll(typeElementNames);
        }
        System.out.println(
            String.format("Remaining type element names: %s", remainingTypeElementNames)
        );
        return false;
    }

    public abstract void processType(Trees trees, TypeElement typeElement, TreePath treePath);

    private void handleAnalyzedType(final TypeElement typeElement) {
        System.out.println(
            String.format("Handling analyzed type element: %s", typeElement)
        );
        if (!remainingTypeElementNames.remove(typeElement.getQualifiedName())) {
            return;
        }

        final TreePath treePath = trees.getPath(typeElement);
        processType(trees, typeElement, treePath);
    }

    private static final class AnalyzeTaskListener implements TaskListener {
        private final AbstractTypeProcessor processor;

        public AnalyzeTaskListener(final AbstractTypeProcessor processor) {
            this.processor = processor;
        }

        @Override
        public void finished(final TaskEvent e) {
            if (e.getKind() != TaskEvent.Kind.ANALYZE) {
                return;
            }

            processor.handleAnalyzedType(e.getTypeElement());
        }
    }
}
CheckMethodBodies class: Annotation class
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.SOURCE)
public @interface CheckMethodBodies {
}
CheckMethodBodiesProcessor class: Annotation processor
import com.google.auto.service.AutoService;
import com.sun.source.tree.BlockTree;
import com.sun.source.tree.CompilationUnitTree;
import com.sun.source.tree.ExpressionTree;
import com.sun.source.tree.MethodTree;
import com.sun.source.tree.StatementTree;
import com.sun.source.tree.SwitchTree;
import com.sun.source.tree.Tree;
import com.sun.source.util.TreePath;
import com.sun.source.util.TreePathScanner;
import com.sun.source.util.Trees;
import javax.annotation.processing.Processor;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeMirror;

@SupportedAnnotationTypes("org.example.annotation.processor.CheckMethodBodies")
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public final class CheckMethodBodiesProcessor extends AbstractTypeProcessor {
    @Override
    public void processType(final Trees trees, final TypeElement typeElement, final TreePath treePath) {
        final CompilationUnitTree compilationUnitTree = treePath.getCompilationUnit();
        final TestMethodTreePathScanner treePathScanner = new TestMethodTreePathScanner(trees, compilationUnitTree);
        treePathScanner.scan(compilationUnitTree, null);
    }

    private static final class TestMethodTreePathScanner extends TreePathScanner<Void, Void> {
        private final Trees trees;
        private final CompilationUnitTree compilationUnitTree;

        public TestMethodTreePathScanner(
            final Trees trees,
            final CompilationUnitTree compilationUnitTree
        ) {
            this.trees = trees;
            this.compilationUnitTree = compilationUnitTree;
        }

        @Override
        public Void visitMethod(final MethodTree node, final Void unused) {
            System.out.println(
                String.format("Visiting method: %s", node.getName())
            );

            final BlockTree blockTree = node.getBody();
            for (final StatementTree statementTree : blockTree.getStatements()) {
                if (statementTree.getKind() != Tree.Kind.SWITCH) {
                    continue;
                }

                final SwitchTree switchTree = (SwitchTree) statementTree;
                final ExpressionTree switchTreeExpression = switchTree.getExpression();
                System.out.println(
                    String.format("Switch tree expression: %s", switchTreeExpression)
                );

                final TreePath treePath = TreePath.getPath(compilationUnitTree, switchTreeExpression);
                final TypeMirror typeMirror = trees.getTypeMirror(treePath);
                System.out.println(
                    String.format("Tree mirror: %s", typeMirror)
                );
            }
            return null;
        }
    }
}

Test project

Maven project
<plugin>
    <groupId>org.apache.maven.plugins</groupId>
    <artifactId>maven-compiler-plugin</artifactId>
    <configuration>
        <annotationProcessorPaths>
            <path>
                <groupId>org.example</groupId>
                <artifactId>annotation-processor</artifactId>
                <version>1.0.0-SNAPSHOT</version>
            </path>
        </annotationProcessorPaths>
        <showWarnings>true</showWarnings>
    </configuration>
</plugin>

To be able to use the annotation class:

<dependency>
    <groupId>org.example</groupId>
    <artifactId>annotation-processor</artifactId>
    <version>1.0.0-SNAPSHOT</version>
</dependency>
Switcher class: Using annotation
import org.example.annotation.processor.CheckMethodBodies;

@CheckMethodBodies
public final class Switcher {
    public void theMethod() {
        final Integer value = 1;
        switch (value.toString() + "0" + "0") {
            case "100":
                System.out.println("Hundred!");
            default:
                System.out.println("Not hundred!");
        }
    }
}

Testing

Execute the command for the annotation processor project:

mvn clean install

Execute the command for the test project:

mvn clean compile

Observe the output of the annotation processor:

Remaining type element names: [org.example.annotation.processor.test.Switcher]
Remaining type element names: [org.example.annotation.processor.test.Switcher]
Handling analyzed type element: org.example.annotation.processor.test.Switcher
Visiting method: <init>
Visiting method: theMethod
Switch tree expression: (value.toString() + "00")
Tree mirror: java.lang.String

Stand-alone program

It is possible to use javac functionality in a stand-alone program.

It seems that it is necessary to get the tree path and then get the type mirror:

final CompilationUnitTree compilationUnitTree = <…>;
final ExpressionTree switchTreeExpression = <…>;

final TreePath treePath = TreePath.getPath(compilationUnitTree, switchTreeExpression);
final TypeMirror typeMirror = trees.getTypeMirror(treePath);

An excerpt from the documentation: TypeMirror (Java Platform SE 8 ):

public interface TypeMirror extends AnnotatedConstruct

Represents a type in the Java programming language. Types include primitive types, declared types (class and interface types), array types, type variables, and the null type. Also represented are wildcard type arguments, the signature and return types of executables, and pseudo-types corresponding to packages and to the keyword void.

Draft implementation

Input file: Switcher class

public final class Switcher {
    public void theMethod() {
        final Integer value = 1;
        switch (value.toString() + "0" + "0") {
            case "100":
                System.out.println("Hundred!");
            default:
                System.out.println("Not hundred!");
        }
    }
}

Program class

Please, replace the "/path/to/Switcher.java" file path value with the actual file path value.

import com.sun.source.tree.BlockTree;
import com.sun.source.tree.CompilationUnitTree;
import com.sun.source.tree.ExpressionTree;
import com.sun.source.tree.MethodTree;
import com.sun.source.tree.StatementTree;
import com.sun.source.tree.SwitchTree;
import com.sun.source.tree.Tree;
import com.sun.source.util.JavacTask;
import com.sun.source.util.TreePath;
import com.sun.source.util.TreePathScanner;
import com.sun.source.util.Trees;
import java.io.IOException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import javax.lang.model.type.TypeMirror;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileObject;
import javax.tools.SimpleJavaFileObject;
import javax.tools.ToolProvider;

public final class Program {
    public static void main(final String[] args) throws IOException {
        final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
        final JavacTask task = (JavacTask) compiler.getTask(
            null,
            null,
            null,
            null,
            null,
            List.of(new TestFileObject())
        );
        final Iterable<? extends CompilationUnitTree> compilationUnitTrees = task.parse();
        task.analyze();
        final Trees trees = Trees.instance(task);

        for (final CompilationUnitTree compilationUnitTree : compilationUnitTrees) {
            final TestMethodTreePathScanner treePathScanner = new TestMethodTreePathScanner(trees, compilationUnitTree);
            treePathScanner.scan(compilationUnitTree, null);
        }
    }

    private static final class TestFileObject extends SimpleJavaFileObject {
        public TestFileObject() {
            super(URI.create("myfo:/Switcher.java"), JavaFileObject.Kind.SOURCE);
        }

        @Override
        public CharSequence getCharContent(final boolean ignoreEncodingErrors) throws IOException {
            return Files.readString(
                Path.of("/path/to/Switcher.java"),
                StandardCharsets.UTF_8
            );
        }
    }

    private static final class TestMethodTreePathScanner extends TreePathScanner<Void, Void> {
        private final Trees trees;
        private final CompilationUnitTree compilationUnitTree;

        public TestMethodTreePathScanner(
            final Trees trees,
            final CompilationUnitTree compilationUnitTree
        ) {
            this.trees = trees;
            this.compilationUnitTree = compilationUnitTree;
        }

        @Override
        public Void visitMethod(final MethodTree node, final Void unused) {
            final BlockTree blockTree = node.getBody();
            for (final StatementTree statementTree : blockTree.getStatements()) {
                if (statementTree.getKind() != Tree.Kind.SWITCH) {
                    continue;
                }

                final SwitchTree switchTree = (SwitchTree) statementTree;
                final ExpressionTree switchTreeExpression = switchTree.getExpression();
                System.out.println(
                    String.format("Switch tree expression: %s", switchTreeExpression)
                );

                final TreePath treePath = TreePath.getPath(compilationUnitTree, switchTreeExpression);
                final TypeMirror typeMirror = trees.getTypeMirror(treePath);
                System.out.println(
                    String.format("Tree mirror: %s", typeMirror)
                );
            }
            return null;
        }
    }
}

The program output:

Switch tree expression: (value.toString() + "00")
Tree mirror: java.lang.String
Related