Scala Compiler Plugin — Annotation based method AST rewriting, wrapping and substitution to inject behaviour dynamically

As a relative newcomer to the Scala Compiler plugin realm, I’ve found the entry curve quite steep when it comes to achieving this type of behavior. 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 the execution of extra functionality.

An example of why this could be useful is we can such as sending blocks of code for remote execution on a distributed platform, without having to write the code everywhere you need it.

This could of course be achieved by leveraging function call, this is just a, depending on the use-case, more elegant approach.

Compiler Plugin Boilerplate

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>

Injecting code around our original code

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 (the q”””).

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 which the annotation was added. This permits us to inject arbitrary code that will be executed at runtime.

Building/compiling the project

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/*.scala
cp src/scalac-plugin.xml out
(cd out; jar cf ../compiler-plugin-release.jar .)

Results

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.

nenuadrian/scala-compiler-plugin-method-rewrite

GitHub code

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

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