Macros in Scala: moderately complex example

I recently began to discover the power of newly-added macros in Scala. While they are still labeled as "experimental", they already pack quite a punch.

Most tutorials on the web usually use some simple examples - like assert macro, or debugging helper.

In this post, I'll try to make a bit deeper dive into the subject, in hope that I will be able to explore darker corners of macro engineering.

The use-case for today's macro is actually something I needed in one of my real projects - I had complex hierarchy of case classes, and I wanted a quick way to visualize it (without needing to read mile-long text files). So I created a macro, that takes case class, and creates a frame with JTree containing the hierarchy of case class members.

Example use-case for the macro looks as follows:

case class Author(name: String, birthYear: Int)
case class Book(title: String, authors: List[Author])
val hobbit = Book("Hobbit: There And Back Again", List(Author("J.R.R.Tolkien", 1892)))
showTree(hobbit)

There result is quite simple, but useful nevertheless:

Not bad for one macro call, huh?

SBT setup

Since macros and main code should be compiled in separate compilation units, we'll need two separate projects in our sbt setup (note that it's not possible to achieve that using plain build.sbt files, so we'll need to dive into full project configuration). Here's the code:

import sbt._
import Keys._

object build extends Build {
  val sv = "2.10.2"

  lazy val root = Project("main", file(".")) settings (
    scalaVersion := sv
  ) dependsOn macros

  lazy val macros = Project("macros", file("macros")) settings (
    scalaVersion := sv,
    libraryDependencies ++= Seq(
      "org.scala-lang" % "scala-compiler" % sv,
      "org.scala-lang" % "scala-swing" % sv
    )
  )
}

Here, we create macro project, that requires compiler and swing libraries, and our main project, that only depends on our macro code. Note that packaging of "main" project wouldn't include macro code or compiler code - you'll need to jump through several more hoops if you want to distribute your macro alongside your application/library (more information in sbt docs).

Scaffolding

Now, when the build system is configured, we can get to actual macro writing. The stub implementation for our macro will look like this (in file macros/ShowTree.scala):

import scala.language.experimental.macros
import scala.reflect.macros.Context
object ShowTree {
  def showTree[T](obj: T) = macro _showTree[T]
  def _showTree[T](c: Context)
                  (obj: c.Expr[T])
                  (implicit wtt: c.WeakTypeTag[T]): c.Expr[Unit] = ???
}

Here, showTree is macro "use-case", and _showTree is actual macro implementation.

Macro implementation takes several argument lists. First, context - it contains a lot of useful information about macro call site, and almost all types inside the macro are dependent on this instance. Second argument list contains expressions for arguments that were actually passed into macro at call-site - in our case, expression that computes the actual instance we will be displaying. The third argument list is used to get TypeTag for the type of instance being displayed (that's actually our only source of information about the instance, and thus is most important part of the macro).

The general idea behind our macro is to take TypeTag for passed case class, and generate some code that will take the case class and show the frame with hierarchy tree.

Let's expand our stub:

def _showTree[T](c: Context)
                (obj: c.Expr[T])
                (implicit wtt: c.WeakTypeTag[T]): c.Expr[Unit] = {
  import c.universe._
  val root = nodes(c)(obj)(wtt.tpe)
  reify {
    UIManager.setLookAndFeel("com.sun.java.swing.plaf.nimbus.NimbusLookAndFeel")
    val frame = new MainFrame {
      contents = new ScrollPane {
        contents = Component.wrap(new JTree(root.splice))
      }
    }
    frame.visible = true
  }
}

One of the most useful helpers in our macro arsenal is "reify" function. It takes a piece of valid scala code, and returns an AST for it - and we can even insert our own Expr's inside it with the help of .splice call (as seen in "root.splice", line 10).

Note call to "nodes" function on 3d line. Since we want to handle case class hirerarchy recursively, we'll need to extract some of our code into recursive function, that will take type tag, find it's type members, and call itself with new type tags. Here's the signature for that function:

def nodes(c: Context)
         (obj: c.Expr[Any])
         (tpe: c.universe.Type): c.Expr[DefaultMutableTreeNode] = ???

Note how it returns an expression for DefaultMutableTreeNode.

Meat of the macro

Let's now proceed to the implementation of "nodes" function.

We are going to handle several possibilities - case classes, sealed traits that have case classes as their implementing classes, lists of values, and all other values.

So inside the "nodes" function, we first check if incoming type is a case class, and then create fresh tree node for it:

if (tpe.normalize.typeSymbol.asClass.isCaseClass) {
  val node = newTermName(c.fresh("node$"))
  val nodeVal = ValDef(Modifiers(), node, TypeTree(), reify {
    new DefaultMutableTreeNode(c.literal(tpe.toString.split("\\.").last).splice)
  }.tree)

This snippet generates code like follows:

val node$42 = new DefaultMutableTreeNode("Book")

Then, we'll need to enumerate case class fields and get their accessor functions. For this, the following helper is useful:

def caseClassAccessors(c: Context)(tpe: c.universe.Type): List[c.universe.MethodSymbol] = {
  import c.universe._
  tpe.declarations.collect {
    case acc: MethodSymbol if acc.isCaseAccessor => acc
  }.toList
}

Given the list of method symbols, we now can generate the code that walks the list of fields and outputs code for creating tree nodes from them. When we encounter a field that is typed as a list, we'll create a node with each list element as a separate child. If, on the other hand, the field is typed as a case class, we'll call "nodes" recursively to get a tree node for that case class. For all other values, we'll simply add "name : value" plain-text node.

val nodeExpr = c.Expr[DefaultMutableTreeNode](Ident(node))
val children = caseClassAccessors(c)(tpe).map { acc =>
  val fieldName = c.literal(acc.name.toString)
  val fieldValue = c.Expr(Select(obj.tree.duplicate, acc.name))
  if (acc.returnType <:< typeOf[Iterable[Any]]) {
    val TypeRef(_, _, typeParam :: Nil) = acc.returnType
    val nodesCall = nodes(c)(c.Expr(Ident(newTermName("v"))))(typeParam)
    reify {
      val listNode = new DefaultMutableTreeNode(fieldName.splice)
      fieldValue.asInstanceOf[c.Expr[Iterable[Any]]].splice.foreach { v =>
        listNode.add(nodesCall.splice)
      }
      nodeExpr.splice.add(listNode)
    }
  } else if (acc.returnType.typeSymbol.isClass &&
             acc.returnType.typeSymbol.asClass.isCaseClass ||
             acc.returnType.typeSymbol.asClass.isTrait) {
    val subNodes = nodes(c)(fieldValue)(acc.returnType)
    reify {
      nodeExpr.splice.add(subNodes.splice)
    }
  } else {
    // handle all other types
    reify {
      nodeExpr.splice.add(new DefaultMutableTreeNode(
        s"${fieldName.splice} : ${fieldValue.splice.toString}"))
    }
  }
}
val addNodes = c.Expr[Unit](Block(children.map(_.tree), Literal(Constant(()))))
reify {
  c.Expr[Unit](nodeVal).splice
  addNodes.splice
  nodeExpr.splice
}

If we encounter a trait, we first check if it is sealed (and abort if it isn't, to avoid runtime errors), then output a match statement, that will determine the exact class type at runtime, and call appropriate code.

} else if (tpe.typeSymbol.asClass.isTrait) {
  if (!tpe.typeSymbol.asClass.isSealed)
    c.abort(c.enclosingPosition, s"Can't handle non-sealed trait ${tpe}")
  else {
    c.Expr[DefaultMutableTreeNode](Match(
      obj.tree.duplicate,
      tpe.normalize.typeSymbol.asClass.knownDirectSubclasses.map(_.asType.toType).toList.map { tp =>
        val name = newTermName(c.fresh("cl$"))
        CaseDef(
          Bind(name, Typed(Ident(nme.WILDCARD), TypeTree(tp))),
          EmptyTree,
          nodes(c)(c.Expr[Any](Ident(name)))(tp).tree)
      }
    ))
  }

Example of generated code (if we assume we have trait Fruit and subclasses Apple and Banana):

obj match {
  case cl$24 @ (_ : Apple) => ...
  case cl$25 @ (_ : Banana) => ...
}

Sadly, current macro implementation contains quite a lot of bugs and corner cases - thus I wasn't able to simply use Typed(Ident(name), TypeTree(tp)) instead of bind (compiler complains that he can't find cl$N values). But that will change in near future, hopefully.

And if incoming type is not a case class or a trait, we simply output a node with string representation of an instance (and a type, for debugging purposes):

} else {
  reify {
    new DefaultMutableTreeNode(
      s"${obj.splice.toString} :: ${c.literal(tpe.toString).splice}")
  }
}

Closing notes

Complete code is available here. Happy hacking!