Using Macro Annotations to Reduce Boilerplate

/ / Engineering

Motivation

Often when modelling certain types of problems Scala (e.g., using third-party data), I’ll find myself needing a bunch of similar but slightly different case classes under a common trait. While this isn’t a big deal with one or two classes, when you have a large number it becomes tedious to do things like change, add, or remove fields that all classes have in common:

case class A(id: Long, name: String, description: Option[String], content: String, created: DateTime)
case class B(id: Long, name: String, description: Option[String], content: String, created: DateTime, xs: List[Long])
case class C(id: Long, name: String, description: Option[String], content: String, created: DateTime, ys: List[String])

You can of course add the common fields to the trait, but these are no longer passed as constructor parameters so constructing instances becomes ugly:

trait Base {
    val id: Long
    val name: String
    val description: Option[String]
    val content: String
    val created: DateTime
}

case class A extends Base
case class B(xs: List[Long]) extends Base

val a = new A {
    override val id = 1L
    override val name = "a1"
    override val description = None
    override val content = "some content"
    override val datetime = new DateTime()
}

Extending an abstract class makes construction nice again, but now you have to type everything an additional time:

abstract class Base(id: Long, name: String, description: Option[String], content: String, created: DateTime)
case class A(id: Long, name: String, description: Option[String], content: String, created: DateTime) extends Base(id: Long, name: String, description: Option[String], content: String, created: DateTime)
case class B(id: Long, name: String, description: Option[String], content: String, created: DateTime, xs: List[Long]) extends Base(id: Long, name: String, description: Option[String], content: String, created: DateTime)

Wouldn’t it be nice to extend a trait and have undeclared vals become contructor parameters without have to explicitly write them out in every implementation? With a little help from macro annotations you can do just that:

trait Base {
    val id: Long
    val name: String
    val description: Option[String]
    val content: String
    val created: DateTime
}

@Extends[Base]
case class A

@Extends[Base]
case class B(xs: List[Long])

How does that work? The short version is that the Extends macro annotation finds all undeclared vals in the trait Base, injects them as constructor params for A and B, and adds Base to the list of traits they extend at compile time. But let’s dive into the code.

As an aside, I’d like to give a hearty thanks to this post from Martin Raison at Kifi andthis one from Imran Rashid for helping me figure all this out.

Implementation

To start off, we have to import a few things:

import scala.annotation.{StaticAnnotation, compileTimeOnly}
import scala.language.experimental.macros
import scala.reflect.macros.whitebox.Context

Next, we define the annotation class, which requires the macro paradise compiler plugin. (If you haven’t used it before, just addaddCompilerPlugin("org.scalamacros" % "paradise" % "2.1.0" cross CrossVersion.full) to your build.sbt.) Note that this class doesn’t really do anything; the actually implementation is defined in its companion object.

@compileTimeOnly("use macro paradise")
class Extends[T](defaults: Any*) extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro Extends.impl
}

Now let’s get to the meat of the implementation. As you can see above, we’ll be defining a method called impl in the Extends object that acts as a macro, meaning that it receives the context (a whitebox one in this case) and returns a modified expression. There are a handful of other methods referenced by or defined within our impl so the flow of the code might get lost here, but you can find the full implementation in this gist.

To start off, we import the universe from our context and then start inspecting the trait we’re mixing in. First, we extract the trait we’re extending by looking at the annotated portion of the tree and using a quasiquote to unlift its components, and then create a “companion” Type instance.

object Extends {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._

    val targetTrait: Tree = c.prefix.tree match {
      case q"new $name[$typ]" => typ
    }

    val tpe: Type = c.typecheck(q"0.asInstanceOf[$targetTrait]").tpe

We next get a list of all undefined vals in our target trait. Note that this is exactly how you’d do this from a TypeTag with reflection at runtime, but since this is a macro it’s run at compile time.

val values: List[TermSymbol] = tpe.members.collect { case d if d.isTerm => d.asTerm }.filter(t => t.isStable && t.isMethod && t.isAbstract).toList.reverse

Next up, a few helper methods. First of all, one to deconstruct our target case class, again using quasiquotes:

def extractClassNameAndFields(classDecl: ClassDef): (TypeName, Seq[Tree], Seq[Tree], Seq[Tree]) = {
    try {
        val q"case class $className(..$fields) extends ..$bases { ..$body }" = classDecl
        (className, fields, bases, body)
    } catch {
        case _: MatchError => c.abort(c.enclosingPosition, "Annotation is only supported on case classes")
    }
}

Then one to add our trait’s values to the existing constructor arguments in the target case class. I find it slightly nicer to have the common fields at the beginning of the argument list.

def modifiedFields(fields: List[ValDef]): List[ValDef] = {
    val newParams = values.map { v =>
        ValDef(Modifiers(Flag.CASEACCESSOR), v.name, TypeTree(v.typeSignature.typeConstructor), EmptyTree)
    }
    newParams ++ fields
}

And finally a method to return the modified case class with its new parameter list and our trait mixed in.

def modifiedClass(classDecl: ClassDef): Expr[_] = {
    val (className, fields, bases, body) = extractClassNameAndFields(classDecl)
    val ctor = modifiedFields(fields.asInstanceOf[List[ValDef]])
    c.Expr(q"case class $className(..$ctor) extends $targetTrait with ..$bases { ..$body }")
}

Now just have to apply this method to the annotated classes and return them.

annottees.map(_.tree) match {
    case (classDecl: ClassDef) :: Nil => modifiedClass(classDecl)
    case _ => c.abort(c.enclosingPosition, "Invalid annottee")
}

And that’s the entire implementation. Now, it’s worth noting a few drawbacks and limitations. First, even if you package up this macro as a standalone entity, projects using it directly must include the macro paradise plugin (though projects using those projects do not). Next, as you can see in the code above this only works on case classes. I don’t think there’s any reason it couldn’t work on regular classes with a few changes, but I haven’t needed it. You can only extend one trait using this technique, though you could easily work around this by creating a trait that mixes together all the traits you want to extend. I haven’t tried using this in Eclipse, but IntelliJ shows a lot of spurious errors when using classes defined with this annotation (though only within the project in which they’re defined). Finally, this was developed for a specific use case and has only been tested under those conditions, so I can’t guarantee that this won’t break down in other situations.

All that said, I’ve found this a handy solution for eliminating boilerplate without losing type safety. Once again, the the full implementation is in this gist; feel free to leave any comments or suggestions there.

TOP