Zenn
🥰

Scala 3 で型安全にコンパイルタイムマクロを書く.

2021/12/05に公開
Scala
Scala 3
tech

Q: Scala 3 のマクロを勉強したいんだけどなにかおすすめのものってある？
A: (＾ω＾)Scala 3 のマクロを書こう(＾ω＾)

Scala 3 では Scala 2.x では experimental だったマクロが正式にサポートされるようになりました. API も刷新されて型安全に快適にマクロを書いてコンパイル時にさまざまな処理を行うことができます. うれしいですね(＾ω＾)

Scala の主要なライブラリの主要なライブラリの特にコア部分の実装についてはマクロも含めて Scala 3 への移行が終わっていますが、まだまだ コア以外のモジュールに移行作業が残っているライブラリや Scala 3への移行が終わっていないライブラリがあります.

マクロと聞くと難しそうな印象を受けるかもしれませんが、Scala のライブラリのうち、マクロ自体が複雑なロジックを持っているもの(shapeless, refined や magnolia など)は少なく、どちらかというと開発者のDX(developer experience)を改善する目的で書かれた補助的なものが多いので身構えなくても大丈夫です. 今回扱うのも後者に分類されます. また、 Scala 3 のマクロは エディタや IDE のサポートがしっかりしているのでコンパイルの仰せのままにコードを書けばいいので簡単です(ホンマか？🤔)

さて、ということでもし興味があればこの記事や xuwei-k さんの記事、Eugene Yokota さんの記事を参考に Scala の OSS に PR を出してみましょう. Scala 2.x のマクロを Scala 3 に対応させるのは、Scala 3 のマクロを勉強するいい教材です.

https://xuwei-k.hatenablog.com/

https://eed3si9n.com/ja/intro-to-scala-3-macros/

それでは本題に入って Scala 3 のマクロについてみていきましょう.

今回扱うマクロ

Scala の関数型プログラミング界隈で有名な Json ライブラリ、circe の literal マクロを Scala 3 で書きます. circe の設計や Json パーサーについても説明するとかなり長くなってしまうので割愛します.

https://github.com/circe/circe

似た設計の json ライブラリをつくりながら学ぶ拙著を公開しているのでそちらを参考するといいかもしれません. (＆投げ銭は大歓迎ですよ？(＾ω＾))

https://zenn.dev/110416/books/a94a86a2c696d1
literal マクロは次のように使います.

val payload: Json = json"""{ "key1": "value"} """

このマクロは コンパイル時に 文字列をパースして Json 型を返します.

したがって次のような json 文字列を書くとコンパイルエラーになります. うれしいですね(＾ω＾) もうこれなしでは json をかけませんね.

val invalidPayload = json"""{  "key" : { "value"" }""" // compile error .....

世はコンパイラに仕事をさせる時代です(＾ω＾) マクロを使うとだれでも簡単にコンパイラに仕事をさせることができるそうですよ?(＾ω＾)

他の言語でも似たようなことができるはずです. 例えば Rust は普通に書くと無駄にコードが膨れ上がる部分に頻繁にマクロが使われているので、その副次的なメリットとしてコンパイル時に入力文字列をチェックしてくれるものが結構ありますよね.

もちろん、この json 文字列には通常の文字列と同じように変数を代入することができます.

val value1 = 1
val key2 = "hoge"
val interp = json"""{  "key1" : $value1 , $key2   : "value2"  }"""

プログラムの中に json をべた書きするようなユースケースはあまりありませんが、 ちょっとしたデモやスクリプトから API を叩くときにこのマクロは便利です.
( Scala は VS Code や IntelliJ で *.sc や *.worksheet.sc の拡張子のファイルを使ってファイルひとつでスクリプト言語のようにプログラムを書くことができます. もちろんコンパイルチェックも補完も機能します. )

StringContext

Scala では foo"文字列" のかたちで書かれた文字列はStringContext インスタンスの foo メソッドを呼び出すように展開されます.
StringContext は parts パラメタを持ち、もともとの文字列の部分は StringContext の parts の部分に、文字列に ${} で埋め込まれた値は foo メソッドの可変長引数として渡されます.

例えば、json"""{ "key" : ${value} }""" は StringContext(parts = "{ \"key\" : "," }").json(value) に展開されます. このとき StringContext は当然のことながら json メソッドを持っていないので Scala 2.x では暗黙のスコープを探索して StringContext => ???(ただし ??? は json メソッドを持つクラス) の変換を試みます.
implicit のスコープに以下のようなメソッドやクラスが定義されていれば、それをつかって StringContext を変換します.

// class を implicit で定義する場合.
implicit class JsonStringContext(sc: StringContext) extends AnyVal {
    def json(args:Any*):Json = ???
}

// implicit def を与える場合.
class JsonStringContext(sc: StringContext) extends AnyVal {
  def json(args:Any*):Json = ???
}
implicit def toJsonStringContext(sc:StringContext):JsonStringContext = ???

Scala 2.x のマクロのエントリーポイント

Scala 2.x のマクロでは StringContext を暗黙の変換で JsonStringContext に変換し、json メソッドをエントリーポイントとしてマクロを呼び出して String => Json に変換しています.

マクロでは Expr 型の AST を操作してコンパイル時に処理を行います. scala 2.x 系では macro キーワードを使って次のようにエントリーポイントを定義します. 左辺の引数 args は AST に持ち上げられたうえで右辺で処理されます.

implicit class JsonStringContext(sc:StringContext) {
  def json(args:Any*) : Json = macro JsonLiteralMacros.jsonStringContext
}

マクロの具体的な内容については後程見てみましょう.

Scala 3.x のマクロのエントリーポイント

Scala 3.x では extension 機能が追加されたのであるクラスに新しいメソッドをはやすために暗黙の変換を使う必要はありません. うれしいですね(＾ω＾)

マクロでは Expr[T] 型の AST を操作してコンパイル時に処理を行います.

jsonImpl は sc: Expr[StringContext] と args: Expr[Seq[Any]] を引数に取ります.
プログラムから AST を得るには ' を、AST を評価するには $ を使います. なおトップレベルの $ はマクロのエントリーポイントでしか使えません.

extension (inline sc:StringContext) {
  inline def json(args:Any*) : Json = ${ JsonLiteralMacros.jsonImpl('sc,'args) }
}

マクロをざっと読む

scala 2.x の literal マクロは package.scala,JsonLiteralMacro.scala, LiteralInstanceMacro.scala の3つのファイルから成ります. package は残りのファイルで書いたコードをまとめて、公開するためのファイルなのでマクロの実体はほぼ JsonLiteralMacro.scala と LiteralInstanceMacro.scala にあります.

それらしいところを探すために StringContext を探してみると次のようなコードが見つかります.

  final def jsonStringContext(args: c.Expr[Any]*): Tree = c.prefix.tree match {
    case Apply(_, Apply(_, parts) :: Nil) =>
      val stringParts = parts.map {
        case Literal(Constant(part: String)) => part
        case _ =>
          c.abort(
            c.enclosingPosition,
            "A StringContext part for the json interpolator is not a string"
          )
      }
      // ...
      val replacements: Seq[Replacement] = args.map(argument => Replacement(stringParts, argument.tree))
      
      // ...
      val jsonString = stringParts.zip(replacements.map(_.placeholder)).foldLeft("") {
          case (acc, (part, placeholder)) =>
            val qm = "\""

            s"$acc$part$qm$placeholder$qm"
        } + stringParts.last

        parse(jsonString, replacements) match {
          case Right(tree) => tree

これをざっと読んでいくと次のような処理をしていることがわかります.

  1. Expr[Any]* からStringContext のparts とjson メソッドの可変長引数を取り出す
  2. 変数が代入されているところをReplacement クラスのインスタンスと対応させる. Replacement はユニークな id を持つ.
  3. 元の文字列の変数が代入されている箇所を Replacement の id に置き換える.
  4. json 文字列をパースする. パーサーには jawn を利用している.

Json のパースにはパフォーマンスを考慮してか jawn という パーサーを使っているようです. jawn は 型[T] に対して型クラス Facade[T] を与えることでアドホックにパーサーの動作を返ることができます.

parse メソッドの定義元を見てみると次のようになっています. 動的にクラスを呼び出しているようですね.

  private[this] def jawnFContextClass = Class.forName("org.typelevel.jawn.FContext")
  private[this] def jawnParserClass = Class.forName("org.typelevel.jawn.Parser$")
  private[this] def jawnParser = jawnParserClass.getField("MODULE$").get(jawnParserClass)
  private[this] def jawnFacadeClass = Class.forName("org.typelevel.jawn.Facade")
  private[this] def parseMethod = jawnParserClass.getMethod("parseUnsafe", classOf[String], jawnFacadeClass)

さらに見ていくと次のような安全ではなさそうな処理が出てきます. ウッ😖 おそらく初見では何をしているのか理解に苦しみますね... これは jawn のパーサーを使うために Facade というクラスを定義する必要があるのですが、それをコンパイル時に与えるためにこのような定義になっているのだと思います. むじゅかしいぽよ～😖

protected[this] val invokeWithArgs: (String, Array[Class[_]], Array[Object]) => Object = {
      // format: off
      case ("jnull", _, _)         => q"_root_.io.circe.Json.Null"
      case ("jfalse", _, _)        => q"_root_.io.circe.Json.False"
      case ("jtrue", _, _)         => q"_root_.io.circe.Json.True"
      case ("singleContext", _, _) => new SingleContextHandler(replacements).asProxy(jawnFContextClass)
      case ("arrayContext", _, _)  => new ArrayContextHandler(replacements).asProxy(jawnFContextClass)
      case ("objectContext", _, _) => new ObjectContextHandler(replacements).asProxy(jawnFContextClass)
      case ("jstring", Array(cls), Array(arg: CharSequence)) if cls == classOf[CharSequence]       => toJsonString(arg)
      case ("jstring", Array(cls, _), Array(arg: CharSequence, _)) if cls == classOf[CharSequence] => toJsonString(arg)
      // format: on
      case ("jnum", Array(clsS, clsDecIndex, clsExpIndex), Array(s: CharSequence, decIndex, expIndex))
          if clsS == classOf[CharSequence] && clsDecIndex == classOf[Int] && clsExpIndex == classOf[Int] =>

Json への Replacement の埋め込みは次のコードで定義されているようです. 文字列に変数が代入されている部分では、与えられた変数の型に応じて適切な Encoder を呼び出して T => Json,(jsonkey は T => String) への変換をしています. この処理は jawn のパーサーの文字列 -> Json の処理の途中で呼び出されています.

 protected[this] final def toJsonKey(s: String): Tree =
      replacements.find(_.placeholder == s).fold(q"$s")(_.asKey)

    protected[this] final def toJsonString(s: CharSequence): Tree =
      replacements.find(_.placeholder == s).fold(q"_root_.io.circe.Json.fromString(${s.toString})")(_.asJson)

Scala 3 のマクロを書く

さて、これらを Scala 3 に対応させていきます.

jsonImpl は次のような定義になります. Expr を操作するには 暗黙のパラメタ(using節)で Quotes を与える必要があります.

def jsonImpl(sc:Expr[StringContext],args:Expr[Seq[Any]])(using q:Quotes)

これによって scala 2.x で Apply(_, Apply(_, parts) :: Nil) のように AST のケースクラスのパターンマッチであらわしていた処理は ' を使って書くことができます. 下のコードを見てください. また、 Expr(AST) に対して様々な メソッドが生えています. これらを使うことでより安全に AST を操作することができます. 例えば、 valueOrAbort は Expr[T] から T の値を得ます. ここでは Expr[Seq[String]] から Seq[String] を得ています.

プリミティブ型に対して FromExpr 型クラスが事前に定義されているのでこのようにして値を得ることができます.

(余談ですが,Int, String, Float などは Singleton 型として扱うことでより詳しい情報をコンパイル時に得ることができます. うれしいですね(＾ω＾) 例えば、次のようにコンパイル時チェック可能な合同式を定義したりすることができます.)

def jsonImpl(sc:Expr[StringContext],args:Expr[Seq[Any]])(using q:Quotes) :Expr[Json]= {
    import q.reflect.*
    val stringParts = sc match {
      case '{StringContext($parts:_*)} => parts.valueOrAbort
    }

    val replacements = args match {
      case Varargs(argExprs) =>
        argExprs.map(Replacement(stringParts,_))
      case other => report.error("Invalid arguments for json literal.");Nil
    }

    val jsonString = stringParts.zip(replacements.map(_.placeholder)).foldLeft("") {
      case (acc,(part,placeholder)) =>
        val qm = "\""
        s"$acc$part$qm$placeholder$qm"   
    } + stringParts.last
    // ...
     Parser.parseFromString[Expr[Json]](jsonString) match {
      case Success(jsonExpr)=> jsonExpr
      case Failure(e) => 
        report.error(e.toString)
        '{null}
    }

Scala 2.x で q"_root_.io.circe.Json.fromString(${s.toString})")(_.asJson) というふうにscala.reflect.api.Quasiquotes.Quasiquote.api.apply[scala.Any] を使っていた処理は次のように書くことができます. コード量は少し増えましたが補完がきいたりコンパイルエラーがでたり、など型のメリットを享受することができます.

def asJson(using q:Quotes):Expr[Json] = {
    import q.reflect.*
    argument match {
    case '{$arg: t} => {
      arg.asTerm.tpe.widen.asType match {
              case '[t] => 
                Expr.summon[Encoder[t]] match {
                  case Some(encoder) => '{ $encoder.apply($arg.asInstanceOf[t]) }
                  case None => report.error(s"could not find implicit Encoder for ${Type.show[t]}", arg);'{null}
            }
          }
    }
  }

パターンマッチによる変数と型の抽出

argument:Expr[Any] に対して case '{$arg: t} => { でパターンマッチすることで、型を変数 t として抽出することができます. 型は json メソッドの可変長引数の型が Seq[Any]なので Any ですが、文字列に変数を代入するときにコンパイラは変数 arg の型情報を持っているはずなので arg.asTerm.tpe.widen.asType で型情報を取り出せます.

※ arg.asTerm.tpe.widen.asType は Scala 2.x の c.typecheck(argument).tpe に対応します. runtime の情報にアクセスしようとするほど型安全性を壊す可能性のある API を使う必要があるので注意が必要ですね.

パターンマッチによる型の処理

arg.asTerm.tpe,arg.asTerm.tpe.widen は TypeRepr 型です. (直感的,雑な説明をするとプログラムが foo と Expr(foo) となるように型は T と TypeRepr[T] になる. ). TypeRepr からは asType メソッドを呼び出すことで型を取り出すことができます.

Scala 3 では 型に対してもパターンマッチすることができるので t から具体的な型を取り出して implicit のスコープにある Encoder[t] を呼び出します. Scala 3 では tuple が HList(複数の型の値を型の情報を保ったまま保持できるListのようなデータ構造) に相当します. コンパイル時に型を操作するときに tuple をしばしば使います.

(Encoder は circe が与える型クラス. ある型 T から Json へ変換する機能を持つ. implicit のスコープに Encoder[T] を与えることで T に対してアドホックに asJson を呼び出せるようになる. )

さて、StringContext から得た parts と json メソッドの引数を ユニークな id で置き換えた jsonString をパースするには Facade クラスを与える必要があります.

     Parser.parseFromString[Expr[Json]](jsonString) match {
      case Success(jsonExpr)=> jsonExpr
      case Failure(e) => 
        report.error(e.toString)
        '{null}
    }

Scala 2.x ではクラスを動的に呼び出したり次のような複雑な処理を書いていました.

      case ("arrayContext", _, _)  => new ArrayContextHandler(replacements).asProxy(jawnFContextClass)
      case ("objectContext", _, _) => new ObjectContextHandler(replacements).asProxy(jawnFContextClass)
      case ("jstring", Array(cls), Array(arg: CharSequence)) if cls == classOf[CharSequence]       => toJsonString(arg)

Scala 3 では inline given と Expr.summon[T] を使うことで同様の処理を安全に記述できます. 型クラスの定義と呼び出しに関しては、マクロを使うからと言って特別なことをする必要はほとんどなく、普段と同じようにコードを書くことができます. うれしいですね(＾ω＾)

    inline given Facade[Expr[Json]] with {

      private def toJsonKey(s:String):Expr[String] =  replacements.find(_.placeholder == s )
      .fold(Expr(s.toString))(_.asKey)
      private def toJsonString(s:String):Expr[Json] = replacements.find(_.placeholder == s )
      .fold{ val strExpr = Expr(s.toString);'{ Json.fromString($strExpr)  }   }(_.asJson)
        // ...
        def isObj: Boolean = false
  
        def add(s: CharSequence): Unit = {
          val strExpr = toJsonString(s.toString)
          values = '{ $strExpr :: $values }
        }
        def add(v: Expr[Json]): Unit =  values = '{  $v :: ${values} }
        def finish(): Expr[Json] = '{Json.arr($values.reverse:_*)}


 Expr.summon[Encoder[t]] match {
                  case Some(encoder) => '{ $encoder.apply($arg.asInstanceOf[t]) }
                  case None => report.error(s"could not find implicit Encoder for ${Type.show[t]}", arg);'{null}

※ inline given は Scala 2.x の implicit def, implicit object による型クラスの定義 ,summon は Scala 2.x の implicitly[T] に対応する

さて、マクロが書けたら適切なディレクトリにコードを配置します. sbt は Scala のビルドバージョンごとのディレクトリ構成に対応しているので、 Scala 2.x に依存するコードを src/main/scala-2 以下に、 Scala 3.x に依存するコードを src/main/scala-3 以下に配置して、細かい設定をごにょごにょすればやるべきことは終わりです. テストがパスするか確認します.

余談ですが、sbt は Mac/Windows/Linux × JVM/js/native × compile/benchmark/test × .../2.11/2.12/2.13/3/... × 各種モジュールなどなど地獄のような組み合わせに対応出来る数少ないビルドツールですね. sbt には足を向けて寝れませんね.
(sbt はどちらの方角にあるのだろうか...)

さあ、これであなたも Scala 3 のマクロが書けますね. Scala ３ 対応が終わっていないライブラリはまだありますよ？ マクロを書いてさっそく Scala の OSS にコントリビュートしましょう(＾ω＾)

手元でマクロの動作を確認できる最小構成のレポジトリはこちら. スターしてね❤

https://github.com/i10416/literalmacro

ちなみに筆者が出した circe への PR はこちら.

https://github.com/circe/circe/pull/1874

Discussion