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

Adrian M. Nenu
Oct 24, 2017 · 4 min read
Image for post
Image for post

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._
class CompilerPlugin(override val global: Global)
extends Plugin {
override val name = "compiler-plugin"
override val description = "Compiler plugin"
override val components =
List(new CompilerPluginComponent(global))
}
class CompilerPluginComponent(val global: Global)
extends PluginComponent with TypingTransformers {
import global._
override val phaseName = "compiler-plugin-phase"
override val runsAfter = List("parser")
override def newPhase(prev: Phase) =
new StdPhase(prev) {
override def apply(unit: CompilationUnit) {
unit.body = new MyTypingTransformer(unit).transform(unit.body)
}
}
class MyTypingTransformer(unit: CompilationUnit)
extends TypingTransformer(unit) {
override def transform(tree: Tree) = tree match {
case _ => super.transform(tree)
}
}
def newTransformer(unit: CompilationUnit) =
new MyTypingTransformer(unit)
}

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.

Image for post
Image for post

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

mkdir out fsc -language:postfixOps -feature -d out src/*.scalacp src/scalac-plugin.xml out(cd out; jar cf ../compiler-plugin-release.jar .)

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.

Image for post
Image for post

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

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch

Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore

Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store