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

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

An example of why this could be useful, is we can asuch 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 Boilerplating

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/*.scalacp 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.

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.

--

--

--

Adrian is a software engineer, passionate about elegant code, photography and writing. Connect @ nenuadrian.com

Love podcasts or audiobooks? Learn on the go with our new app.

Recommended from Medium

Get a FREE Cooking School Layout Pack for Divi

19coders

Xevious — Positions and Waypoints

Introduction To Kubernetes

Logo of Kubernetes

Builder Design Pattern

What are Residential Proxies?

Programming A Markov-Chain Based Reddit Bot

Learning Algorithms — Heap | Part 02

Veeam Backup & Replication: A VMCE Preparation Guide

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
Adrian M. Nenu 😺

Adrian M. Nenu 😺

Adrian is a software engineer, passionate about elegant code, photography and writing. Connect @ nenuadrian.com

More from Medium

KivaKit — Type Conversion

Existential Types and Variance in Scala

Project Reactor — Part 1

Java Sealed Class