Scala Compiler Plugin — Annotation based method AST rewriting / wrapping / substitution

As a relative new-comer to the Scala Compiler plugin realm, I’ve found the entry curve quite steep. This is a quick introduction to creating a plugin, traversing the AST and wrapping the body-block of an annotated method into an inner one so as to allow execution of extra functionality (such as remote execution).

If you’re familiar with creating the boilerplate for a Scala Compiler plugin, you might want to skip ahead.

We’ll need a scalac-plugin.xml which gets packed in our plugin’s .jar

<plugin> 
<name>CompilerPlugin</name>
<classname>com.nenuadrian.CompilerPlugin</classname>
</plugin>

The meat of the plugin is in CompilerPlugin.scala where we’ll be using a TypingTransformer to traverse the compiler generated Abstract Syntax Tree [AST] and manipulate our method’s body to our heart’s content.

import scala.tools.nsc.Global
import scala.tools.nsc.Phase
import scala.tools.nsc.plugins._
import scala.tools.nsc.transform._

The above is the skeleton for creating a plugin able to parse the AST; we’ll enhance our match statement to change the block of annotated methods, using quasiquotes.

override def transform(tree: Tree) = tree match {          
case dd: DefDef if (dd.mods.annotations.size > 0) =>
println(dd)
val ddd = treeCopy.DefDef(dd, dd.mods, dd.name,
dd.tparams, dd.vparamss, dd.tpt, Block(
q"""println("Inside - before")""",
DefDef(Modifiers(), TermName("runMethod"), List(),
List(), TypeTree(), dd.rhs)
,
q"val r = runMethod",
q"""println("Inside - after")""",
q"r"
))
println(ddd)
ddd
case _ => super.transform(tree)
}

The bolded code shows the creation point of an inner method with the body of the original one.

To build this plugin you can use something like the following (which is well documented enough in other articles)

mkdir out 

Given that our test class is

class wrapThisMethod extends StaticAnnotation {}
class TestClass {
@wrapThisMethod def myMethod: Int = {
4
}
}

When we compile it using our compiler plugin

scalac -classpath compiler-plugin-release.jar -feature -language:postfixOps -Xplugin:compiler-plugin-release.jar -d out src/*.scala

If we take a quick peek at the output we’ll see something like this before the transformation occurs for our annotated method (which is dd)

@new wrapThisMethod() def myMethod: Int = 4

And the compiler will also tell us what it has come up with after our transformation

@new wrapThisMethod() def myMethod: Int = {
println("Inside - before");
def runMethod = 4;
val r = runMethod;
println("Inside - after");
r
}

conclusion

We’ve created an inner method (runMethod) with the body of our initial method. This inner method gets called towards the end and its result gets pushed out of the parent, while we can enhance this execution in whatever way we want (from simply measuring its runtime to sending it to a remote grid for execution).

This is certainly not the most brilliant solution and it’s definitely hacky, only setting us on the right path. However, with the reduced specific to this type of intent number of examples available out there, it has taken me personally long enough to figure this out that I found it worthwhile to record it for my and your future use.

gitHub code

Feel free to PR and I’ll approve and update the article with better approaches ;)

The Scala compiler plugins opens multiple universes of possibilities to us and I recommend taking advantage of this ability to provide extra tooling to developers in our environment

Full-stack software engineer — “The computer programmer is a creator of universes” ~ Joseph Weizenbaum