アノテーションプロセッサで AST 変換 - Lombok を参考にして変数の型をコンパイル時に変更

Java のボイラープレートを補完してくれる Lombok の処理内容が興味深かったので、これを真似た簡単なサンプルプログラムを作ってみました。

ソースは http://github.com/fits/try_samples/tree/master/blog/20150117/

はじめに

Lombok はアノテーションプロセッサを使って AST (抽象構文木) の変換を実施しています。

Lombok の使い方

まずは Lombok を使って下記のような Java ソースのコンパイルを試してみます。 val@Value が Lombok の機能を使用している箇所です。

Sample.java
import lombok.val;
import lombok.Value;

public class Sample {
    public static void main(String... args) {
        // lombok.val の使用
        val d = new Data("sample1", 10);
        System.out.println(d);
    }

    // lombok.Value の使用
    @Value
    private static class Data {
        private String name;
        private int value;
    }
}

Service Provider 機能(META-INF/services)を使用するため、javac 時に classpath へ lombok.jar を指定するだけで適用されます。

javac によるコンパイル (Lombok 使用)
> javac -cp lombok.jar Sample.java

CFR を使って Sample.class の内容を確認してみると、lombok.vallombok.Value が消え、代わりに型やメソッドを補完している事が分かります。

Sample.class の内容確認 (CFR 利用)
> java -jar cfr_0_94.jar Sample.class

/*
 * Decompiled with CFR 0_94.
 */
import java.beans.ConstructorProperties;
import java.io.PrintStream;

public class Sample {
    public static /* varargs */ void main(String ... arrstring) {
        Data data = new Data("sample1", 10);
        System.out.println(data);
    }

    private static final class Data {
        private final String name;
        private final int value;

        @ConstructorProperties(value={"name", "value"})
        public Data(String string, int n) {
            this.name = string;
            this.value = n;
        }

        public String getName() {
            return this.name;
        }

        public int getValue() {
            return this.value;
        }

        public boolean equals(Object object) {
            if (object == this) {
                return true;
            }
            if (!(object instanceof Data)) {
                return false;
            }
            Data data = (Data)object;
            String string = this.getName();
            String string2 = data.getName();
            if (string == null ? string2 != null : !string.equals(string2)) {
                return false;
            }
            if (this.getValue() != data.getValue()) {
                return false;
            }
            return true;
        }

        public int hashCode() {
            int n = 1;
            String string = this.getName();
            n = n * 59 + (string == null ? 0 : string.hashCode());
            n = n * 59 + this.getValue();
            return n;
        }

        public String toString() {
            return "Sample.Data(name=" + this.getName() + ", value=" + this.getValue() + ")";
        }
    }

}

Lombok の仕組み

次に、Lombok の仕組みを簡単に説明します。

Lombok はアノテーションプロセッサ内 (lombok.javac.apt.Processor) にて RoundEnvironment を元に AST を取得し変換します。

javac 実行時の処理を大雑把に書くと下記のようになっています。

  • (1) lombok.core.AnnotationProcessor を処理
  • (2) lombok.javac.apt.Processor を処理
  • (3) lombok.javac.JavacTransformer を処理
  • (4) AnnotationVisitor を処理
  • (5) 各種 AST 変換用のハンドラ (lombok.javac.handlers パッケージ内のクラス) を処理

lombok.jar の META-INF/services/javax.annotation.processing.Processor に (1) のクラス名が記載されているため、アノテーションプロセッサの仕組みによって (1) が実行されます。

(5) の各種ハンドラは lombok.javac.HandlerLibrary が以下のファイルから取得し管理します。

  • META-INF/services/lombok.javac.JavacASTVisitor (visitorHandlers)
  • META-INF/services/lombok.javac.JavacAnnotationHandler (annotationHandlers)

JavacASTVisitor インターフェース実装クラスの lombok.javac.handlers.HandleVal (現時点では唯一の JavacASTVisitor 実装ハンドラ) は、lombok.val を型として使っている変数を適切な型に変更するという処理を行います。

HandleVal の処理内容が興味深かったので、今回はこれを真似た簡易的な処理を作ります。

アノテーションプロセッサで AST 変換

ここからは、Lombok における val の処理を真似たアノテーションプロセッサを自作していきます。

(1) AST の取得

まずは AST を取得して出力するだけのアノテーションプロセッサを作ります。

Compiler Tree API を使用するので、Gradle でビルドする場合は JDK の tools.jar を dependencies へ設定しておきます。

build.gradle
apply plugin: 'java'

dependencies {
    compile files("${System.properties['java.home']}/../lib/tools.jar")
}

次に、Service Provider 設定ファイルを用意しておきます。

このファイルを用意しておけば、JAR ファイルを classpath へ指定するだけで sample.SampleProcessor1 が実行されます。

src/main/resources/META-INF/services/javax.annotation.processing.Processor
sample.SampleProcessor1

AbstractProcessor を extends してアノテーションプロセッサを作成します。

@SupportedAnnotationTypes("*") アノテーションを付与する事で、アノテーションの使用有無に関わらずコンパイル対象の全ソースを RoundEnvironment から取得できるようになります。

Lombok のソースでは JCCompilationUnit へキャストして使っていますが、CompilationUnitTree が AST に該当します。

CompilationUnitTree を取得するため Trees を使っています。

src/main/sample/SampleProcessor1.java
package sample;

import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedSourceVersion;
import javax.annotation.processing.SupportedAnnotationTypes;

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

import com.sun.source.tree.CompilationUnitTree;

import com.sun.source.util.Trees;
import com.sun.source.util.TreePath;

@SupportedSourceVersion(SourceVersion.RELEASE_8)
@SupportedAnnotationTypes("*") //コンパイル対象の全ソースを対象とする
public class SampleProcessor1 extends AbstractProcessor {
    private Trees trees;

    @Override
    public void init(ProcessingEnvironment procEnv) {
        trees = Trees.instance(procEnv);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        // コンパイル対象の全ソースを処理
        roundEnv.getRootElements().stream().map(this::toUnit).forEach(u -> {
            System.out.println("----- CompilationUnitTree -----");
            // AST の内容を出力
            System.out.println(u);
        });

        return false;
    }

    // AST の取得
    private CompilationUnitTree toUnit(Element el) {
        TreePath path = trees.getPath(el);
        return path.getCompilationUnit();
    }
}

今回は Java 8 の API を使ったので @SupportedSourceVersion(SourceVersion.RELEASE_8) を付けています。

ビルド

> gradle build

:compileJava
:processResources UP-TO-DATE
:classes
:jar
:assemble
:compileTestJava UP-TO-DATE
:processTestResources UP-TO-DATE
:testClasses UP-TO-DATE
:test UP-TO-DATE
:check UP-TO-DATE
:build

BUILD SUCCESSFUL

動作確認 (アノテーションプロセッサの適用)

ビルド結果の build/libs/anp-sample1.jar ファイルを以下の Java ソース (A1.java と A2.java) のコンパイルに使ってみます。

A1.java
public class A1 {
    public void sample() {
        int i = 10;
    }
}

下記の var は後で使うので今のところは気にしないで下さい。

A2.java
public class A2 {
    // 下記の var はインターフェースやクラスで定義しても同じ
    @interface var {}

    public var a = 10;
    public String b = "bbb";

    public void sample() {
        var msg = "test data";
        System.out.println(msg);
    }
}

Lombok と同様に javac 時に classpath へ anp-sample1.jar を指定するだけです。

実行結果 (javac 実行)
> javac -cp anp-sample1.jar *.java

----- CompilationUnitTree -----

public class A1 {

    public A1() {
        super();
    }

    public void sample() {
        int i = 10;
    }
}
----- CompilationUnitTree -----

public class A2 {

    public A2() {
        super();
    }

    @interface var {
    }
    public var a = 10;
    public String b = "bbb";

    public void sample() {
        var msg = "test data";
        System.out.println(msg);
    }
}
A2.java:5: エラー: 不適合な型: intをvarに変換できません:
        public var a = 10;
                       ^
A2.java:9: エラー: 不適合な型: Stringをvarに変換できません:
                var msg = "test data";
                          ^
エラー2個

CompilationUnitTree オブジェクトを println するとソースが出力されました。

ここで重要なのは、var を使った箇所のエラーが SampleProcessor1 処理の後に出力されている点です。 (var は int や String では無いので当然のエラーです)

つまり、型チェックはアノテーションプロセッサの後に実施される事が分かります。

であれば、アノテーションプロセッサ内で適切な型に変えてしまえばコンパイルを通す事が可能という事になり、Lombok の val は実際にこの仕組みを利用しています。

(2) AST 内の var 型を Object へ変更

次に、(1) でエラーが発生した A2.java 内の varjava.lang.Object へ変更するように実装します。

CompilationUnitTree (AST) は Visitor パターンで処理できるようになっています。

今回は var を使っている変数定義の部分を処理させたいだけなので、TreeScannervisitVariable メソッドをオーバーライドした内部クラス (VarVisitor) を用意しました。

JCVariableDeclvartype フィールドに変数の型が設定されており、これを変更すれば変数の型が変わります。

vartype が var という名前の型 (所属パッケージに関係なく) であれば java.lang.Object へ変更するようにしてみました。

src/main/sample/SampleProcessor2.java
package sample;
・・・
import com.sun.source.tree.VariableTree;
・・・
import com.sun.tools.javac.model.JavacElements;

import com.sun.tools.javac.processing.JavacProcessingEnvironment;

import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.JCTree.JCExpression;

@SupportedSourceVersion(SourceVersion.RELEASE_8)
@SupportedAnnotationTypes("*")
public class SampleProcessor2 extends AbstractProcessor {
    private Trees trees;
    private TreeMaker maker;
    private JavacElements elements;

    @Override
    public void init(ProcessingEnvironment procEnv) {
        trees = Trees.instance(procEnv);

        JavacProcessingEnvironment env = (JavacProcessingEnvironment)procEnv;

        maker = TreeMaker.instance(env.getContext());
        elements = JavacElements.instance(env.getContext());
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        // コンパイル対象の全ソースを処理
        roundEnv.getRootElements().stream().map(this::toUnit).forEach(this::processUnit);

        return false;
    }

    private CompilationUnitTree toUnit(Element el) {
        TreePath path = trees.getPath(el);
        return path.getCompilationUnit();
    }

    private void processUnit(CompilationUnitTree cu) {
        // Visitor パターンで AST を処理
        cu.accept(new VarVisitor(), null);
    }

    private class VarVisitor extends TreeScanner<Void, Void> {
        // 変数定義の処理
        @Override
        public Void visitVariable(VariableTree node, Void p) {
            System.out.println("visitVariable: " + node);

            if (node instanceof JCVariableDecl) {
                JCVariableDecl vd = (JCVariableDecl)node;

                if ("var".equals(vd.vartype.toString())) {
                    JCExpression ex = maker.Ident(elements.getName("java"));
                    ex = maker.Select(ex, elements.getName("lang"));
                    ex = maker.Select(ex, elements.getName("Object"));
                    // 型を java.lang.Object へ変更
                    vd.vartype = ex;
                }
            }
            return null;
        }
    }
}

java.lang.Object の JCExpression を構築している箇所はもっと上手い方法がありそうですが、とりあえず Lombok の処理を真似ました。

src/main/resources/META-INF/services/javax.annotation.processing.Processor
sample.SampleProcessor2

ビルド

> gradle build

・・・

動作確認 (アノテーションプロセッサの適用)

(1) の時と同様に A1.java と A2.javaコンパイルに使ってみます。

実行結果 (javac 実行)
> javac -cp anp-sample2.jar *.java

visitVariable: int i = 10
visitVariable: public var a = 10
visitVariable: public String b = "bbb"
visitVariable: var msg = "test data"

今回はコンパイルエラーが発生せず、A1.java と A2.java 内の変数定義を visitVariable で処理しています。

CFR で A2.class の内容を確認してみると、var を使っていた変数の型が、一応 java.lang.Object に変わっています。 ("test data" の直接の型は String となっていますが)

A2.class の内容確認 (CFR 利用)
> java -jar cfr_0_94.jar A2.class

/*
 * Decompiled with CFR 0_94.
 */
import java.io.PrintStream;

public class A2 {
    public Object a = 10;
    public String b = "bbb";

    public void sample() {
        String string = "test data";
        System.out.println((Object)string);
    }

    static @interface var {
    }

}

最後に、この仕組みを流用すれば Java で型に別名を付ける機能 (ソース内だけの型エイリアス・型シノニムのようなもの) を実現できるかもしれません。