Groovy AST Transformations by Example: Adding Methods to Classes

What can you do with a Groovy AST Transformation? A difficult question, considering the answer is "almost anything". One frequent use of AST Transformations is to write Groovy changes into the JVM .class files so that standard Java tools will know about them. For instance, the bytecode changes produced by @Delegate are visible and usable from plain old Java classes. This example is going to walk you through how to add synthesized methods to a class at compile time using a Groovy AST Transformation and the AST Builder. It makes an interesting test case since it explores some of the edge cases and common pitfalls. This should be fun...

(And for those of you who want to skip all this and just see the code, check out the latest Groovy example in SVN and just run the Ant script).

svn co http://svn.codehaus.org/groovy/trunk/groovy/groovy-core/src/examples/astbuilder

To start, let's better define our example. Consider the class with a public API and also a standard main() method. Frequently, the main() method simply instantiates the class and calls a single method. Wouldn't it be nice to annotate a method so that it is treated like a main method? What if this were possible:

class MainExample {

   @Main
   public void greet() {
       println "Hello from the greet() method!"
   }   
}

At compile time, this annotation would synthesize a proper main method and run the same code as greet(). (We'll worry about calling an instance method from a static method later). Transforming this annotation in the AST means that both Java and Groovy would see the main() method. Let's do it; there are three things you need:

1. An @Main marker interface
2. An ASTTransformation implementation make the AST changes
3. A test harness (trust me, you want this)
   

The @Main marker interface is the easiest piece. The marker interface is used to a standard Java 5 annotation to an ASTTransformation implementation:

@Retention (RetentionPolicy.SOURCE)
@Target ([ElementType.METHOD])
@GroovyASTTransformationClass (["MainTransformation"])
public @interface Main {  }

This is just a standard Java 5 annotation. The retention policy is SOURCE so that other objects don't see the annotation at runtime when using reflection (there is no need). The ElementType is METHOD because we only want to apply the annotation to methods. And the GroovyASTTransformationClass is "MainTransformation". This value is the fully qualified classname of the ASTTransformation implementation that the Groovy compiler is going to call when an @Main annotation is found. In practice, you will need to specify a package... I left one off here for simplicity.

So this annotation will wire @Main into a class called MainTransformation. I suppose we'd better go make this class then...

The ASTTransformation interface defines one method:

void visit(ASTNode[], SourceUnit)

The best way to understand the contents of the two parameters is to write several tests and inspect them in a debugger, and then use the AST Browser to analyze the syntax tree. In general, the ASTNode array contains what you want to manipulate and the SourceUnit provides services like error and warning reporting. Here is the skeleton of our MainTransformation that is going to add a synthesized main() method to the class:

@GroovyASTTransformation(phase = CompilePhase.INSTRUCTION_SELECTION)
public class MainTransformation implements ASTTransformation {
  
   void visit(ASTNode[] astNodes, SourceUnit sourceUnit) {

       if (!astNodes) return
       if (!astNodes[0]) return
       if (!astNodes[1]) return
       if (!(astNodes[0] instanceof AnnotationNode)) return
       if (astNodes[0].classNode?.name != Main.class.getName()) return
       if (!(astNodes[1] instanceof MethodNode)) return

       MethodNode annotatedMethod = astNodes[1]
       ClassNode declaringClass = astNodes[1].declaringClass

       // makeMainMethod synthesizes a method!
       MethodNode mainMethod = makeMainMethod(annotatedMethod)
       declaringClass.addMethod(mainMethod)
   }
   ...

Again, there are a few interesting things to notice here...

• The compiler phase chosen was INSTRUCTION_SELECTION, which was a completely arbitrary decision. Please, if you want to know which phase to target then ask on the groovy-user mailing list (and asking will give us better hints as to what we need to document on the wiki).
• The body of the transformations contains many guard clauses as a safeguard against NullPointerExceptions and unexpected input. The AST of classes changes from version to version. Defensive programming and thorough test cases highly recommended.
• The MethodNode that was annotated comes through in the ASTNode array, and from there you can navigate most anywhere within the syntax tree.
• The makeMainMethod is going to copy the body of the annotated method and then we will add that method back onto the declaring class.
  

The last thing to do is define the makeMainMethod that actually creates the main():

MethodNode makeMainMethod(MethodNode source) {
   def ast = new AstBuilder().buildFromSpec {
       method('main', ACC_PUBLIC | ACC_STATIC, Void.TYPE) {
           parameters {
               parameter 'args': String[].class
           }
           exceptions {}
           block { }
       }
   }
   MethodNode newMainMethod = ast[0]
   newMainMethod.code = source.code
   newMainMethod
}

Again, a few bits and pieces to make special note of:

• The AstBuilder.buildFromSpec is a DSL convenience over the concrete constructors of ASTNode and its subclasses. Thus, the API of buildFromSpec mirrors the API of the ASTNode classes. If you want to understand the arguments and parameters types, then look at the ASTNode Javadoc. In this regard, the buildFromSpec is more about making the code easier to read rather than easier to write.
• The return type of the method is Void.TYPE and not Void.class. Void.class is a real class with real usages, Void.TYPE is the symbol for "little vee void".
• The modifiers are ACC_PUBLIC | ACC_STATIC. Notice the use of | and not &. Let me editorialize briefly: I don't understood bitmasks, I never understood them, and people who design APIs around parameters typed as "int modifers" will be first against the wall when the revolution comes. If null was the billion dollar mistake then bitmasks at least deserve to be in the tens of millions range. Anyway...
• The codeblock of the new main method is just a reference to the codeblock of the method annotated with @Main. This will cause duplicate bytecode to be written out into the .class file (bad) but does make for a small, clear example (good). In production, you probably want to invoke the @Main method instead of copy the body.
  

Now, as you can see from the javap output, the compiled MainExample.groovy contains a real, live Java main class:

hdarcy:$ groovyc MainExample.groovy
hdarcy:$ javap MainExample
Compiled from "MainExample.groovy"
public class MainExample extends java.lang.Object implements groovy.lang.GroovyObject{
   ...
   public void greet();
   ...
   public static void main(java.lang.String[]);
   ...
}

Phew! That was a lot of information. The last thing to add is a simple test case using Groovy's TranformTestHelper class. A debugger helps AST Transformation development immensely, and this integration test is one way to hook one up to your transformation. I tested this in IDEA but it should work fine in Eclipse/Netbeans as well:

class MainIntegrationTest extends GroovyTestCase {

    public void testInvokeUnitTest() {

       def file = new File('./MainExample.groovy')
       assert file.exists()

       def invoker = new TranformTestHelper(new MainTransformation(), CompilePhase.CANONICALIZATION)

       def clazz = invoker.parse(file)
       def tester = clazz.newInstance()
       tester.main(null)       // main method added with AST transform
   }
}

This test case takes the example class as a simple File and then compiles it using the MainTransformation we defined. This is a great way to invoke both local and global transformations from unit tests. I'm sure there is an easier way to invoke the static main method without instantiating the class, but I'm tired and don't care all that much. It's just an example!

And that's it... we created a marker interface to trigger an annotation, created an ASTTransformation to add a method onto a class, and wrote a test harness. There are just two loose ends to tie up:

How can a static main method invoke a non static method like Greet? We'll, it can't really. For this example we just copied the method body into the main method. In practice, you'd want to make sure only static methods were annotated.

How are inner classes handled? Short answer: not the same way. You have a test case, so just write a test method and find out!

All the code was checked into the Groovy examples folder along with an Ant script. There's no better time to get the source:

svn co http://svn.codehaus.org/groovy/trunk/groovy/groovy-core/src/examples/astbuilder

cd examples/astbuilder and away you go.