ASTを活用してコードの自動修正に挑戦してみよう - Cybozu Inside Out

どうも！アプリケーション基盤チームの@yokotaso です。単純だけど、大量のソースコードの修正が必要な場合、みなさんはどうやって修正していますか？

Junit4からJunit5の移行調査をしていたときに、例外を検証する@Testの expected がJunit5では消えていることがわかりました。社内のコードを調べたところ、修正が必要な箇所が1000箇所くらいということがわかったので、ASTを活用した自動修正ツールを作ってみました。今回は自動修正ツールを使った大量修正の話を紹介します。

ASTとはなにか？

ASTとはプログラミング言語のソースコードの構造をツリーとして表現したもです。

今回はJavaの話なので、javaparserを使って話していきます。

例えば次のようなテストコードがあったとします。

@Test(expected = IllegalArgumentException.class)
public void testNullInput() {
    // setup
    Dog sut = new Dog();
    sut.sit();

    // exercise & verify
    sut.eat(null);
}

このプログラムをASTに分解すると次のような構造になります(簡略化しています)。

f:id:cybozuinsideout:20180904175225p:plain

今回はこのテストコードを次のように変換します。

@Test
public void testNullInput() {
    // setup
    Dog sut = new Dog();
    sut.sit();

    // exercise & verify
    assertThatThrownBy(() -> sut.eat(null)).isInstanceOf(IllegalArgumentException.class);
}

ASTを活用して自動修正ツールを作ってみよう

社内のコード変換につかったコードはjunit-exception-test-replacer においてありますので、もしよかったらご利用ください。

修正したい箇所をVisitorパターンを実装して探す

メソッドに対して@Test(expected=...class) となっているところを探します。

public class ExceptionTestVisitor extends VoidVisitorAdapter<ExceptionTest> {

    private void collectMethodDeclaration(final MethodDeclaration n, final ExceptionTest arg, final Optional<AnnotationExpr> testAnnotation) {
        // 後述
    }

    @Override
    public void visit(final MethodDeclaration n, final ExceptionTest arg) {
        super.visit(n, arg);
        Optional<AnnotationExpr> testAnnotation = n.getAnnotationByClass(Test.class);
        if (!testAnnotation.isPresent()) {
            return;
        }
        // メソッド定義の位置情報を収集する
        collectMethodDeclaration(n, arg, testAnnotation);
    }
}

修正したい場所の位置情報や必要な情報を収集する

ソースコードからASTを構築するとソースコード上の行や列の位置情報を取得することができます。ソースコードの修正対象箇所を収集していきましょう。

このメソッドの中では、次の情報を収集しています。

修正が必要な @Test がプログラムファイルの中のどこに存在しているのか
@Test(expected = ...class) で指定されているクラス名
テストメソッドの中の最後のステートメント

private void collectMethodDeclaration(final MethodDeclaration n, final ExceptionTest arg, final Optional<AnnotationExpr> testAnnotation) {
        testAnnotation.ifPresent((annotationExpr) -> {
            for (Node node : annotationExpr.getChildNodes()) {
                if (!(node instanceof MemberValuePair)) {
                    continue;
                }
                
                MemberValuePair memberValuePair = (MemberValuePair) node;
                if (memberValuePair.getName().asString().equals("expected")) {
                    Range range = annotationExpr.getRange().orElseThrow(IllegalStateException::new);
                    // @Test(expcected=...class) の 例外クラスの情報を取り出す
                    ClassExpr expected = (ClassExpr) memberValuePair.getValue();
                    arg.testAnnotationPositions.put(range, expected);
                    
                    // テストメソッドの最後のステートメントの位置情報を記録する
                    List<Statement> statements = n.getBody().orElseThrow(IllegalStateException::new).getStatements();
                    range = statements.get(statements.size() - 1).getRange().orElseThrow(IllegalAccessError::new);
                    arg.lastStatementPositions.put(range, expected);
                }
            }
        });
}

ASTで解析した情報を使ってパッチを自動生成する

ASTを使ってソースコードを解析した結果、修正が必要なソースコードの位置情報や修正に必要な情報がわかりました。この情報を使ってパッチを自動生成していきましょう。パッチにはJavaのOSSライブラリのDiff Utils を使っています。

ASTのデータ構造を直接書き換える方法もありますが、その場合、空白行やImport文が微妙に変わるなど難しいことが実装してわかったのでパッチを生成する方法をとっています。

`@Test` に修正するパッチを生成

patch = new Patch<>();
Pair<Range, ClassExpr> rangeAndClassExpr = ...;
Range range = rangeAndClassExpr.getOne();
int beginLine = range.begin.line - 1;
int beginColumn = range.begin.column - 1;
String oldSource = sourceCode.get(beginLine);
String indent = oldSource.substring(0, beginColumn);

// @Test(expected = Throwable.class) を次のように置き換えるパッチを生成
// @Test
Chunk<String> oldChunk = new Chunk<>(beginLine, Lists.newArrayList(oldSource));
Chunk<String> newChunk = new Chunk<>(beginLine, Lists.newArrayList(indent + "@Test"));
patch.addDelta(new ChangeDelta<>(oldChunk, newChunk));

例外が発生することを検査するようにコードを修正する

テストメソッドの最後のステートメントで例外が発生する仮定で、テストコードを修正するパッチを生成していきます。テストメソッドの最後のステートメントを assertj の assertThatThrownBy を使って変換します。

assertThatThrownBy(() => /* テストメソッドの最後のステートメント */)
    .isInstanceOf(/* @Test(expecred=) で指定されていたクラス */)

上のようなコードを頭にイメージしながら、パッチを生成していきます。

Pair<Range, ClassExpr> rangeAndClassExpr = ...;
Range range = rangeAndClassExpr.getOne();
int beginLine = range.begin.line - 1;
int beginColumn = range.begin.column - 1;
String indent = sourceCode.get(beginLine).substring(0, beginColumn);

StringBuilder newSource = new StringBuilder(indent).append("assertThatThrownBy(() -> ");
newSource.append(sourceCode.get(beginLine).substring(beginColumn, range.end.column - 1));
newSource.append(").isInstanceOf(").append(getExpectedClass(rangeAndClassExpr.getTwo())).append(");");

Chunk<String> oldChunk = new Chunk<>(beginLine, Collections.singletonList(sourceCode.get(beginLine)));
Chunk<String> newChunk = new Chunk<>(beginLine, Collections.singletonList(newSource.toString()));
patch.addDelta(new ChangeDelta<>(oldChunk, newChunk));