AST Transformations: Creating a Complex AST Transformation

In this (quite long) blog entry we will examine a full implementation of an AST transformation using a closure as an annotation parameter. So be prepared to either be bored out of your mind or alternatively learn a bit about creating AST transformation. Your choice.

The Example Revisited
Until now we have used the AstBuilder to create an AST from a String containing the source code. But the AstBuilder offers two other methods buildFromCode() and buildFromSpec() that are equally interesting and have their respective advantages.

To demonstrate the use of the two methods we change our example a little bit. Until now we have used an annotation that took a String as an argument. We now want to provide a closure instead, which has the advantage that we can provided arbitrary code that results in a boolean expression. First of all we need the changed annotation. We define the type of the provided value (our closure) as Class. This way we can provide Closure objects without colliding with the limits regarding the allowed types of annotation values.

@Retention (RetentionPolicy.SOURCE)
@Target ([ElementType.METHOD])
@GroovyASTTransformationClass ("Requires2Transformation")
public @interface Requires2 {
    Class value() default {true};
}


Our annotated example looks still very much the same, only instead of a string we provide a closure as argument to the annotation. To show that we can access local variables inside the closure scope we use the field a for the comparison.

public class Requires2Example {
    def a = 0

    @Requires2({divisor != a})
    public int divide10By(divisor) {
        10/divisor
    }
}


The transformation should introduce code that checks for the precondition defined as a closure, and if it does not evaluate successfully, throw an exception:

    public int divide10By(divisor) {
        if(! {divisor != a}() )
            throw new Exception('Precondition violated')
        10/divisor
    }


The test looks nearly the same as in the original example.

class Requires2ClassloaderTest extends GroovyTestCase {

    public void testInvokeUnitTest() {
        def file = new File('./Requires2Example.groovy')
        assert file.exists()

        GroovyClassLoader invoker = new GroovyClassLoader()
        def clazz = invoker.parseClass(file)
        def out = clazz.newInstance()

        assert out.divide10By(2) == 5

        def exceptionMessage = shouldFail { out.divide10By(0) }
        assert exceptionMessage =~ /Precondition violated/
    }
}


The AST Transformation
Next is the transformation itself. Here we have more changes. As before we start with defensive programming and check the correct types for the arguments with which the method visit()has been called:

def annotationType = Requires2.class.name

private boolean checkNode(astNodes, annotationType) {
    if (! astNodes) return false
    if (! astNodes[0]) return false
    if (! astNodes[1]) return false
    if (!(astNodes[0] instanceof AnnotationNode)) return false
    if (! astNodes[0].classNode?.name == annotationType) return false
    if (!(astNodes[1] instanceof MethodNode)) return false
    true
}

public void visit(ASTNode[] astNodes, SourceUnit sourceUnit) {

    if (!checkNode(astNodes, annotationType)) {
        addError("Internal Error: wrong arguments", astNodes[0], sourceUnit)
        return
    }
    ...


After the initial check of the arguments we check for the correct type of the expression in the annotation, but now instead of hoping for a ConstantExpression ( the String) we now look for a ClosureExpression. If the type is not correct, we add an error to the error collector using the method addError() (see blog entry).

    MethodNode annotatedMethod = astNodes[1]
    def annotationExpression = astNodes[0].members.value

    if (annotationExpression.class != ClosureExpression) {
         addError("Internal Error: annotation doesn't contain closure", 
                 astNodes[0], sourceUnit)
        return
    }


The remaining part of the visit()-method is nearly the same as before. We call the method createStatements() that returns the if-statement as an AST and add that to the beginning of the AST of the annotated method.

    IfStatement block = createStatements(annotationExpression)

    def methodStatements = annotatedMethod.code.statements
    methodStatements.add(0, block)
}


In the first example we used the buildFromString()-method provided by the AstBuilder to create our AST (see blog entry). This time we want to test the other methods' usefulness to create the AST. Please remember that all AstBuilder-methods return an array of ASTNodes that contains as first element the AST describing the provided fragment, and as second element the whole script class that has been created from it. We normally simply access the first entry for our uses.

AstBuilder - buildFromCode()
This method is very interesting because it can only be called at compile time. If you examine the implementation of the method you only see an IllegalStateException being thrown. The trick here is that it has an AST transformation of its own (class AstBuilderTransformation), that, if it detects a call to this method at compile time, reroutes it to a private class AstBuilderInvocationTrap. This in turn identifies, from the given code, the associated source code, extracts that and calls the AstBuilder method buildFromString(). At first you might think this to be a bit convoluted, but in fact it is brilliant, because this way, the already partway compiled code is reevaluated in the context of the target class.

The following code fragment generates part of the AST that we need for the example, namely throwing the exception if the precondition is not met. Consider the following:

exception = ab.buildFromCode 
                { throw new Exception("Precondition violated in ${this.class}") }


The reference to the class will return the class in which the annotation is found, not our transformation class in which the code has been originally defined. This provides interesting opportunities with the caveat that our beforehand-knowledge of where the annotation might be used is limited at best.

The disadvantage of this approach though is that with buildFromCode() you cannot access the surrounding scope e.g., to add values from local variables to the generated code.

AstBuilder - buildFromSpec()
The third of the methods that AstBuilder offers to create ASTs provides a DSL (domain-specific language) to define your AST. There is no formal specification for the DSL, but the keywords are defined as method names in the class AstSpecificationCompiler (this is the delegate for the closure specifying the AST, and thus can implement the keywords).
Here is the example that we will use in our method createStatements() to create the AST representing the if-statement.

List<ASTNode> res = ab.buildFromSpec {
    ifStatement {
        booleanExpression {
            not {
                methodCall {
                    expression.add(closureAST)
                    constant "call"
                    argumentList {
                    }
                }
            }
        }
        block {
            expression.add(exception)
        }
        empty()
    }
}


Now examine the builder statements and compare them to the AST in the GroovyConsole (Script->Inspect AST)  using the example code of what the transformation should produce and you see the similarities. Two times I use the code expression.add(<local variable>) to insert ASTs held in the respective local variables. The AST in the variable exception for instance can be created using the call to buildFromCode() shown above. This way we can inject arbitrary sub-ASTs into the one we are just building using the DSL.

AstBuilder - Comparing the Provided Methods
We now have worked with each of the three methods that are provided by AstBuilder, and each of these methods has respective advantages and disadvantages. buildFromString() allows you to construct your code right before creating the AST using simple operations on the String representing the code, but this comes at the disadvantage that the compiler does not check the validity of the provided code beforehand. Thus you might introduce some subtle errors that can be difficult to trace. buildFromCode() allows you to facilitate your development environment and some parts of the compiler to ensure that you are provided with syntactically correct code, but has the disadvantage that you cannot access local variables or splice in some sub-AST. buildFromSpec() provides the functionality to do this, but is a little bit cumbersome. Thus, each of the methods has their merits and in every nontrivial case you will combine them.

The Source Code of The Closure
Now we have everything we need to create the transformation. The only thing that is less than perfect is the exception message we create using buildFromCode(). It is much more elegant to use buildFromString() as we did in our first example and use the source code of the closure expression provided in the annotation. Something like the following:

    def statements = "throw new Exception('Precondition violated: $source')"
    AstBuilder ab = new AstBuilder()
    BlockStatement exception
    exception = ab.buildFromString(CompilePhase.SEMANTIC_ANALYSIS, false, statements)[0]

Problem: We do not have the source of the closure. What we have is the sourceUnit, which contains the full source code that is associated with the compilation and we have information in the AST about the lines and columns in which the source code is found. This means, that, while not readily accessible, we can extract the source code from the source unit. The following code shows how this is done. It is derived from Hamlet D'Arcy's code in the class AstBuilderTransformation, where he has to solve the exact same problem.

private String convertToSource(ASTNode node, SourceUnit sourceUnit) {
    if (node == null) {
        addError("Error in convertToSource: node is null", node, sourceUnit)
        return ""
    }

    FileReaderSource sourceFileReader = sourceUnit.getSource()
    int first = node.getLineNumber()
    int last  = node.getLastLineNumber()
    StringBuilder result = new StringBuilder();

    for (int line in first..last) {
        String content = sourceFileReader.getLine(line, null);
        if (content == null) {
            addError("Error accessing source for closure.", node, sourceUnit);
            content = ""
        }
        if (line == last) content = content[0 .. node.getLastColumnNumber()-1]
        if (line == first) content = content[node.getColumnNumber()-1 .. -1]

        result.append(content).append('\n');
    }
    result.toString().trim()
}


The Final Solution
Here is the full implementation of the method createStatements(). We use both buildFromString() and buildFromSpec() to create our final AST and use the method convertToSource() to access the source code associated with the closure provided in the annotation. And this is it.

IfStatement createStatements(ClosureExpression closureAST, SourceUnit sourceUnit) {
    String source = convertToSource(closureAST, sourceUnit)
    def statements = "throw new Exception('Precondition violated: $source')"
    AstBuilder ab = new AstBuilder()
    BlockStatement exception 
    exception = ab.buildFromString(CompilePhase.SEMANTIC_ANALYSIS, false, statements)[0]

    List<ASTNode> res = ab.buildFromSpec {
        ifStatement {
            booleanExpression {
                not {
                    methodCall {
                        expression.add(closureAST)
                        constant "call"
                        argumentList {
                        }
                    }
                }
            }
            block {
                expression.add(exception)
            }
            empty()
        }
    }
    IfStatement is = res[0]
    return is
}


End 
Now we have covered all you need to easily create even complex AST transformations. There are still some dark and murky corners in the compiler and the AST generation, but using the tools that we have been provided with is better than a flashlight in the dark. So, have fun playing with the compiler.

Sourcecode

Requires2.groovy	The annotation for our example using a closure
Requires2Example.groovy	The example class using the annotation
Requires2ClassLoaderTest.groovy	The test using a Groovy Classloader
Requires2Transformation.groovy	The implementation of the AST transformation

Interesting Classes

org.codehaus.groovy.ast.builder.AstBuilderTransformation	The global AST transformation that identifies buildFromCode() calls and creates the AST.
org.codehaus.groovy.ast.builder.AstSpecificationCompiler	The helper class implementing the DSL for the AstBuilder