Scala Compiler PluginでAST変換をやってみる

先週、Scalaの公式ページのDivByZeroプラグインを写経して、Scala Compiler Pluginを作成してみました。今度は、いよいよAST変換にトライしてみたいと思います。

動機は、@BeanPropertyアノテーションと@BooleanBeanPropertyアノテーションで、これを毎度毎度フィールドに付与するのってけっこう面倒だと思うのですよねぇ…。クラス単位で付与できないもんかしら？かといって、Scalaのアノテーションを調べても、そんなものはありません。
※@scala.reflect.BeanInfoアノテーションが名前からして近そうでしたが、残念ながら意味が違いました…

というわけで、せっかくAST変換に興味を持ったのでこれをテーマにAST変換を行うCompiler Pluginを作成してみたいと思います。

目標。以下のように、クラス定義に@Beansアノテーションを付与すると…

@Beans
class TestBean {
  val valName: String = "Test"
  var varName: String = "Test"
  val valFlag: Boolean = true
  var varFlag: Boolean = true
}

以下の様に変換されること。

class TestBean {
  @BeanProperty val valName: String = "Test"
  @BeanProperty var varName: String = "Test"
  @BooleanBeanProperty val valFlag: Boolean = true
  @BooleanBeanProperty var varFlag: Boolean = true
}

なお、以下のCompiler Pluginをちょっと参考にしています。近いけど、ちょっと違うのよねぇ。
beans-scalac-plugin
https://github.com/mvv/beans-scalac-plugin

まずは、build.sbt。

name := "beans-plugin"

version := "0.0.1"

scalaVersion := "2.9.0"

organization := "littlewings"

libraryDependencies += "org.scala-lang" % "scala-compiler" % "2.9.0"

scalac-plugin.xml。作成するPluginのクラス名は、「BeansPlugin」とします。

<plugin>
  <name>beans</name>
  <classname>BeansPlugin</classname>
</plugin>

Beans.scala（アノテーション）

class Beans extends Annotation

では、Plugin本体（BeansPlugin.scala）を少しずつ紹介します。

今回のAST変換は、コンパイル時に@BeanProperty、@BooleanBeanPropertyの両アノテーションを付与することで、getter／setterを生成することが目的です。自力ではgetter／setterは生成しません。

ということは、@BeanProperty、@BooleanBeanPropertyをコンパイラが解釈しているフェーズの前に処理に介入する必要があります。これをやっているのは、Scala Compilerのソースを追うと、以下のトレイト（正確には、内部クラス）でした。
scala.tools.nsc.typechecker.Namers

これが実行されているフェーズは「namer」です。「namer」は「parser」のすぐ後なので、実は2番目のフェーズだったりします…。

よって、「parser」フェーズの後で動くようにrunsAfterを宣言します。なお、runsBeforeをオーバーライドして「namer」の前に割り込むことを明記できるらしいですが、今回は書いても書かなくても結果が変わらなかったので一旦コメントアウト。

class BeansPlugin(val global: Global) extends Plugin {
  val name: String = "beans"
  val description: String = "@Beans annotation to convert @BeanProperty or @BooleanBeanProperty"
  val components: List[PluginComponent] = List(Component)

  private object Component extends PluginComponent with Transform {
    val global: BeansPlugin.this.global.type = BeansPlugin.this.global
    val runsAfter: List[String] = List("parser")
    // override val runsBefore: List[String] = List("namer")

    val phaseName: String = BeansPlugin.this.name

    def newTransformer(unit: global.CompilationUnit): global.Transformer = new BeansTransformer

AST変換を行う場合には、PluginComponentを継承する際に、TransformトレイトをMix-inします。そして、newPhaseの代わりにnewTransformerメソッドを記述します。

作成したTransformer…の一部

    class BeansTransformer extends global.Transformer {
      var needImportBeanProperty: Boolean = _
      var needImportBooleanBeanProperty: Boolean = _

      override def transform(tree: global.Tree): global.Tree = isClassDef(tree) match {
        case true =>
          val newClassDef = transformClassDef(tree.asInstanceOf[global.ClassDef])
          insertImportIfNeeded(newClassDef)
        case false => super.transform(tree)
      }

Transformerを作成する際には、基本はtransformメソッドをオーバーライドする形で進めていくらしいです。で、引数にASTが渡ってくるので、ここからこれを解釈して変換したりすることになるわけですね。

これがけっこう大変です。目を通すべきAPIは、実はCompiler APIではなくて以下のトレイトおよびそのサブクラスです。
scala.reflect.generic.Trees

あと、実装参考的には以下のトレイトにも目を通しています。
scala.tools.nsc.typechecker.Namers
scala.tools.nsc.typechecker.Typers

構文解析後のTreeは、Treesトレイトのサブクラスとして渡ってくるので、これを分解する時には必ずTreesのサブクラスのAPIを眺めることになります。なお、scala.tools.nsc.GlobalクラスはTreesをMix-inしているので、match式で分解する際にはGlobalクラスのインスタンスが付きまといます。
普通はimport global._するみたいですが、今回は明記したかったのでいちいちglobal.と書くことにしました。

手順としては、まずは渡された構文木がクラス定義かどうかを判定し、@Beansアノテーションが付与されていれば、変換処理に移ります。

      def isClassDef(tree: global.Tree): Boolean = tree match {
        case global.ClassDef(_, _, _, _) => true
        case _ => false
      }

      // クラス定義を変換する
      def transformClassDef(classDef: global.ClassDef): global.ClassDef =
        hasBeansAnnotationClass(classDef) match {
          case true =>
            import classDef._
            val newImpl = transformImpl(impl)
            val removedAnnotations = mods.annotations.filterNot(isBeansAnnotation(_))
            val newMods = global.Modifiers(mods.flags, mods.privateWithin, removedAnnotations, mods.positions)
            global.treeCopy.ClassDef(classDef, newMods, classDef.name, tparams, newImpl)
          case false => classDef
        }

クラスにアノテーションが付与されているかどうかは、こんな感じで確認しました。
※完全修飾名では確認できないっぽいです

      def isBeansAnnotation(tree: global.Tree): Boolean = hasAnnotation(tree, "Beans")

      // 指定された名前のアノテーションがtreeに付与されているか確認する
      def hasAnnotation(tree: global.Tree, targetName: String): Boolean = tree match {
        case global.Apply(global.Select(global.New(global.Ident(name)), _), _) => name.toString == targetName
        case _ => false
      }

で、クラスの中身はTemplateクラスとして表現されており、そのbodyに対して変換処理を行います。

      // Template.implを変換する
      def transformImpl(impl: global.Template): global.Template = {
        val newBody = impl.body.map {
          tree => tree match {
            case vd @ global.ValDef(mods, name, tpt, rhs) if canAppendBooleanBeanProperty(vd) =>
              // @BooleanBeanPropertyアノテーションを付与
              needImportBooleanBeanProperty = true
              val newAnnotations = createAnnotation("BooleanBeanProperty") :: mods.annotations
              val newMods = global.Modifiers(mods.flags, mods.privateWithin, newAnnotations, mods.positions)
              global.treeCopy.ValDef(vd, newMods, name, tpt, rhs)

            case vd @ global.ValDef(mods, name, tpt, rhs) if canAppendBeanProperty(vd) =>
              // @BeanPropertyアノテーションを付与
              needImportBeanProperty = true
              val newAnnotations = createAnnotation("BeanProperty") :: mods.annotations
              val newMods = global.Modifiers(mods.flags, mods.privateWithin, newAnnotations, mods.positions)
              global.treeCopy.ValDef(vd, newMods, name, tpt, rhs)

            case _ if isClassDef(tree) =>
              // インナークラスは再帰的に処理
              transformClassDef(tree.asInstanceOf[global.ClassDef])
            case _ => tree
          }
        }

        global.treeCopy.Template(impl, impl.parents, impl.self, newBody)
      }

すでに@BeanPropertyまたは@BooleanBeanPropertyが付いているかどうかを確認して、付いていればスキップします。あと、publicなフィールドしか相手にしておりません。

      // @BooleanBeanPropertyアノテーションが追加可能か確認する
      def canAppendBooleanBeanProperty(vd: global.ValDef): Boolean = vd.tpt match {
        case global.Ident(name) if vd.mods.isPublic && name.toString == "Boolean" =>
          !vd.mods.annotations.exists(hasAnnotation(_, "BooleanBeanProperty"))
        case _ => false
      }

      // @BeanPropertyアノテーションが追加可能か確認する
      def canAppendBeanProperty(vd: global.ValDef): Boolean = vd.tpt match {
        case global.Ident(name) if vd.mods.isPublic && name.toString != "Boolean" =>
          !vd.mods.annotations.exists(hasAnnotation(_, "BeanProperty"))
        case _ => false
      }

アノテーションを動的に作成する箇所は、こんな感じになっています。

      // 指定された名前のアノテーションを作成する
      def createAnnotation(annotationName: String): global.Tree = global.Apply(
        global.Select(
          global.New(global.Ident(global.newTypeName(annotationName))),
          global.nme.CONSTRUCTOR), Nil)

最後に、後続のフェーズで依存関係のエラーとならないように、import文を追加します。

      // 必要であれば、クラス定義にimportを挿入する
      def insertImportIfNeeded(classDef: global.ClassDef): global.ClassDef = {
        var importSelectors: List[global.ImportSelector] = Nil
        if (needImportBeanProperty) {
          // BooleanBeanPropertyのImportSelectorを作成する
          importSelectors = global.ImportSelector(global.newTermName("BeanProperty"),
                                                  -1,
                                                  global.newTermName("BeanProperty"),
                                                  -1) :: importSelectors
        }

        if (needImportBooleanBeanProperty) {
          // BeanPropertyのImportSelectorを作成する
          importSelectors = global.ImportSelector(global.newTermName("BooleanBeanProperty"),
                                                  -1,
                                                  global.newTermName("BooleanBeanProperty"),
                                                  -1) :: importSelectors
        }

        importSelectors.isEmpty match {
          case true => classDef
          case false =>
            val importDef =
              global.Import(global.Select(global.Ident(global.newTermName("scala")), global.newTermName("reflect")),
                            importSelectors)

            global.treeCopy.ClassDef(classDef,
                                     classDef.mods,
                                     classDef.name,
                                     classDef.tparams,
                                     global.treeCopy.Template(classDef.impl,
                                                              classDef.impl.parents,
                                                              classDef.impl.self,
                                                              importDef :: classDef.impl.body))
        }
      }

なお、元の@Beansアノテーションは除去してあります。

では、これをコンパイルしてTestBean.scalaに対して実行してみましょう。
なお、@BeansアノテーションはCompiler Plugin自身に含めてしまったので、コンパイルにはプラグインの設定とクラスパスの設定の両方が必要です…。

「-Xprint」オプションを指定すると、指定されたフェーズの後のASTを表示することができます。今回の「beans」フェーズを指定すると、「beans」フェーズでの変換結果となるASTが表示されます。

$ scalac -Xprint:beans -Xplugin:beans-plugin/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar -cp beans-plugin/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar TestBean.scala 
[[syntax trees at end of beans]]// Scala source: TestBean.scala
package <empty> {
  class TestBean extends scala.ScalaObject {
    import scala.reflect.{BooleanBeanProperty, BeanProperty};
    def <init>() = {
      super.<init>();
      ()
    };
    @new BeanProperty() val valName: String = "Test";
    @new BeanProperty() var varName: String = "Test";
    @new BooleanBeanProperty() val valFlag: Boolean = true;
    @new BooleanBeanProperty() var varFlag: Boolean = true
  }
}

「parser」フェーズだと、こんな感じ。

$ scalac -Xprint:parser -Xplugin:beans-plun/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar -cp beans-plugin/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar TestBean.scala 
[[syntax trees at end of parser]]// Scala source: TestBean.scala
package <empty> {
  @new Beans() class TestBean extends scala.ScalaObject {
    def <init>() = {
      super.<init>();
      ()
    };
    val valName: String = "Test";
    var varName: String = "Test";
    val valFlag: Boolean = true;
    var varFlag: Boolean = true
  }
}

@Beansアノテーションがなくなり、@BeanPropertyおよび@BooleanBeanPropertyが付与されるようにAST変換されているのがわかります。

もちろん、getter／setterが生成されておりますともさ。

$ javap TestBean
Compiled from "TestBean.scala"
public class TestBean extends java.lang.Object implements scala.ScalaObject{
    public java.lang.String valName();
    public java.lang.String varName();
    public void varName_$eq(java.lang.String);
    public void setVarName(java.lang.String);
    public boolean valFlag();
    public boolean varFlag();
    public void varFlag_$eq(boolean);
    public void setVarFlag(boolean);
    public boolean isVarFlag();
    public boolean isValFlag();
    public java.lang.String getVarName();
    public java.lang.String getValName();
    public TestBean();
}

インナークラスとか入れても、ちゃんと動きますよ。

$ scalac -Xprint:beans -Xplugin:beans-plugin/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar -cp beans-plugin/target/scala-2.9.0.final/beans-plugin_2.9.0-0.0.1.jar TestBean.scala 
[[syntax trees at end of beans]]// Scala source: TestBean.scala
package <empty> {
  import scala.reflect.BeanProperty;
  class TestBean extends scala.ScalaObject {
    import scala.reflect.{BooleanBeanProperty, BeanProperty};
    def <init>() = {
      super.<init>();
      ()
    };
    @new BeanProperty() val valName: String = "Test";
    @new BeanProperty() var varName: String = "Test";
    @new BooleanBeanProperty() val valFlag: Boolean = true;
    @new BooleanBeanProperty() var varFlag: Boolean = true;
    class InnerClass extends scala.ScalaObject {
      def <init>() = {
        super.<init>();
        ()
      };
      @new BeanProperty() var innerVarName: String = "Test"
    }
  }
}

すでにわかっている欠点は、@BeanPropertyおよび@BooleanBeanPropertyアノテーションがimportされているかどうかをチェックしていない（無駄にimportする可能性がある）ことですかねぇ。あと、アノテーションを本当に付与していいかどうかのチェックがかなり甘いのですが、触りとしてはこんなもんじゃないかなぁ？…と思いたいです。

一応、Pluginの完全なコードを貼っておきます。

import scala.tools.nsc.{Global, Phase}
import scala.tools.nsc.plugins.{Plugin, PluginComponent}
import scala.tools.nsc.transform.Transform

class BeansPlugin(val global: Global) extends Plugin {
  val name: String = "beans"
  val description: String = "@Beans annotation to convert @BeanProperty or @BooleanBeanProperty"
  val components: List[PluginComponent] = List(Component)

  private object Component extends PluginComponent with Transform {
    val global: BeansPlugin.this.global.type = BeansPlugin.this.global
    val runsAfter: List[String] = List("parser")
    // override val runsBefore: List[String] = List("namer")

    val phaseName: String = BeansPlugin.this.name

    def newTransformer(unit: global.CompilationUnit): global.Transformer = new BeansTransformer

    class BeansTransformer extends global.Transformer {
      var needImportBeanProperty: Boolean = _
      var needImportBooleanBeanProperty: Boolean = _

      override def transform(tree: global.Tree): global.Tree = isClassDef(tree) match {
        case true =>
          val newClassDef = transformClassDef(tree.asInstanceOf[global.ClassDef])
          insertImportIfNeeded(newClassDef)
        case false => super.transform(tree)
      }

      def isClassDef(tree: global.Tree): Boolean = tree match {
        case global.ClassDef(_, _, _, _) => true
        case _ => false
      }

      // クラス定義を変換する
      def transformClassDef(classDef: global.ClassDef): global.ClassDef =
        hasBeansAnnotationClass(classDef) match {
          case true =>
            import classDef._
            val newImpl = transformImpl(impl)
            val removedAnnotations = mods.annotations.filterNot(isBeansAnnotation(_))
            val newMods = global.Modifiers(mods.flags, mods.privateWithin, removedAnnotations, mods.positions)
            global.treeCopy.ClassDef(classDef, newMods, classDef.name, tparams, newImpl)
          case false => classDef
        }

      // 必要であれば、クラス定義にimportを挿入する
      def insertImportIfNeeded(classDef: global.ClassDef): global.ClassDef = {
        var importSelectors: List[global.ImportSelector] = Nil
        if (needImportBeanProperty) {
          // BooleanBeanPropertyのImportSelectorを作成する
          importSelectors = global.ImportSelector(global.newTermName("BeanProperty"),
                                                  -1,
                                                  global.newTermName("BeanProperty"),
                                                  -1) :: importSelectors
        }

        if (needImportBooleanBeanProperty) {
          // BeanPropertyのImportSelectorを作成する
          importSelectors = global.ImportSelector(global.newTermName("BooleanBeanProperty"),
                                                  -1,
                                                  global.newTermName("BooleanBeanProperty"),
                                                  -1) :: importSelectors
        }

        importSelectors.isEmpty match {
          case true => classDef
          case false =>
            val importDef =
              global.Import(global.Select(global.Ident(global.newTermName("scala")), global.newTermName("reflect")),
                            importSelectors)

            global.treeCopy.ClassDef(classDef,
                                     classDef.mods,
                                     classDef.name,
                                     classDef.tparams,
                                     global.treeCopy.Template(classDef.impl,
                                                              classDef.impl.parents,
                                                              classDef.impl.self,
                                                              importDef :: classDef.impl.body))
        }
      }


      // Template.implを変換する
      def transformImpl(impl: global.Template): global.Template = {
        val newBody = impl.body.map {
          tree => tree match {
            case vd @ global.ValDef(mods, name, tpt, rhs) if canAppendBooleanBeanProperty(vd) =>
              // @BooleanBeanPropertyアノテーションを付与
              needImportBooleanBeanProperty = true
              val newAnnotations = createAnnotation("BooleanBeanProperty") :: mods.annotations
              val newMods = global.Modifiers(mods.flags, mods.privateWithin, newAnnotations, mods.positions)
              global.treeCopy.ValDef(vd, newMods, name, tpt, rhs)

            case vd @ global.ValDef(mods, name, tpt, rhs) if canAppendBeanProperty(vd) =>
              // @BeanPropertyアノテーションを付与
              needImportBeanProperty = true
              val newAnnotations = createAnnotation("BeanProperty") :: mods.annotations
              val newMods = global.Modifiers(mods.flags, mods.privateWithin, newAnnotations, mods.positions)
              global.treeCopy.ValDef(vd, newMods, name, tpt, rhs)

            case _ if isClassDef(tree) =>
              // インナークラスは再帰的に処理
              transformClassDef(tree.asInstanceOf[global.ClassDef])
            case _ => tree
          }
        }

        global.treeCopy.Template(impl, impl.parents, impl.self, newBody)
      }

      // 指定された名前のアノテーションを作成する
      def createAnnotation(annotationName: String): global.Tree = global.Apply(
        global.Select(
          global.New(global.Ident(global.newTypeName(annotationName))),
          global.nme.CONSTRUCTOR), Nil)

      // @BooleanBeanPropertyアノテーションが追加可能か確認する
      def canAppendBooleanBeanProperty(vd: global.ValDef): Boolean = vd.tpt match {
        case global.Ident(name) if vd.mods.isPublic && name.toString == "Boolean" =>
          !vd.mods.annotations.exists(hasAnnotation(_, "BooleanBeanProperty"))
        case _ => false
      }

      // @BeanPropertyアノテーションが追加可能か確認する
      def canAppendBeanProperty(vd: global.ValDef): Boolean = vd.tpt match {
        case global.Ident(name) if vd.mods.isPublic && name.toString != "Boolean" =>
          !vd.mods.annotations.exists(hasAnnotation(_, "BeanProperty"))
        case _ => false
      }

      // @Beansアノテーションがクラスに付与されているか確認する
      def hasBeansAnnotationClass(classDef: global.ClassDef): Boolean =
        classDef.mods.annotations.exists(isBeansAnnotation)

      def isBeansAnnotation(tree: global.Tree): Boolean = hasAnnotation(tree, "Beans")

      // 指定された名前のアノテーションがtreeに付与されているか確認する
      def hasAnnotation(tree: global.Tree, targetName: String): Boolean = tree match {
        case global.Apply(global.Select(global.New(global.Ident(name)), _), _) => name.toString == targetName
        case _ => false
      }
    }
  }
}

いやぁ、けっこう大変でしたわ。