Recently I played a bit with compiler plugins for Scala.

My goal is to write a plugin which would generate checks in code for methods annotated with JSR305 annotations. As a first step, I wanted to handle the @Nonnull annotation on parameters. Supposing you have the following code:

def parameterMustBeNull(@Nonnull parameter : Object) = {
 parameter.toString();
}

the compiler plugin should transform it to:

def parameterMustBeNull(@Nonnull parameter : Object) = {
 if (param == null)
  throw new IllegalArgumentException("Parameter 'parameter' should not be null.")
 parameter.toString();
}

That way, the method contract expressed by annotations in the method header doesn’t have to be duplicated by explicit checks in the code later – they are generated automatically.

To achieve that, the plugin must first find the parameters that are annotated with @Nonnull. Then, for each such parameter, it must generate a tree corresponding to the if part. Finally, these generated chunks of code must be added to the beginning of the body of the method. All of this is quite easy to do thanks to the Scala compiler API and the AST that you can inspect and manipulate. But I think it’s best just to look at source code.

In some more detail, as we want to transform the AST, the plugin must use the Transform trait. Then we have to implement the transform method, which looks if the tree passed is a method and if so, looks for the nonnull parameters:

object AnnotationsCheckGenComponent extends PluginComponent with Transform {
 ...
 object AnnotationsCheckGenTransformer extends Transformer {
  override def transform(tree: Tree) = {
   tree match {
    // DefDef is a method definition
    case dd @ DefDef(_, _, _, vparamss, _, _) => {
     // here we add the nonnull checks
    }
    case t => super.transform(t)
  }
}

The vparamss list holds the list of parameters. To search for parameters which are annotated with a nonnull annotation, we can use the very handy Scala list manipulation methods:

vparamss.flatten[ValDef].filter(param => {
 // Checking if the parameter contains a non-null annotation
 param.mods.annotations.exists(annotation => {
  // Checking if the annotation is a non-null annotation
  annotation.constr.find(tree => {
   tree match {
    case Ident(name) => "No[nt][Nn]ull".r.findFirstIn(name.toString).isDefined
    case _ => false
   }
  }).isDefined
 })
})

As you can see above, any nonnull annotation will be detected, not only the ones defined by JSR305. And finally, for each parameter annotated, we must generate the if block containing the check. It’s almost like writing the code directly ;).

If(
 // if: param == ...
 Apply(
  Select(
   Ident(param.name),
   newTermName("$eq$eq")),
  List(Literal(Constant(null)))
 ),
 // then: throw ...
 Throw(
  Apply(
   Select(
    New(Ident(newTypeName("IllegalArgumentException"))),
    newTermName("<init>")),
   List(Literal(Constant("Parameter '%s' should not be null.".format(param.name)))))),
 // else
 EmptyTree)

To use the plugin during compiling, all you need to do is have a jar (you can download a ready one here) that includes the compiled plugin class and a descriptor, and run the compilation with the -Xplugin parameter, for example:

scalac -cp $JSR305_JAR_PATH -Xplugin:annotations-check-gen.jar Example.scala

The only problem I had was in which phase the plugin should be executed. At first I tried executing it after namer, and later, taking the advice I got on the forums (by the way, the forum members are really helpful), after typer. However, then I had to run the typer again on the generated code, and I didn’t yet figure out how to type the generated if in a context of the method, so that the typer has access to the information about available values (here, the value being the parameter checked is used). So for now, instead, the plugin runs after the parser phase (so in fact as the first plugin), and then the result handled by the rest of the phases as if the code was there from the beginning. But if somebody has an idea, on how to run the typer properly, please let me know :).

You can find the entire source code on github. For building and testing I used buildr. I must say that writing the build using that tool was a lot more pleasant than using Maven2, even considering the fact that I don’t know Ruby.

The plugin can be extended to generate checks for other annotations, for example @MatchesPattern, @RegEx etc. Another thing to do is to generate checks also for return values of methods (e.g. if a method is annotated with @Nonnull, this means that the return value shuoldn’t be null). This is a bit more tricky however, as there can be many exit points from the method. Going further, the generation can be extended to more advanced method contracts, expressed for example via typestate annotations.

If you are also interested in writing compiler plugins for Scala you can find a good tutorial here, and a step-by-step guide describing another plugin here.

Adam