Using Gradle to Build & Apply AST Transformations

by Matt Cholick

Recently, I wanted to both build and apply local ast transformations in a Gradle project. While I could find several examples of how to write transformations, I couldn't find a complete example showing the full build process. A transformation has to be compiled separately and then put on the classpath, so its source can't simply sit in the rest of the Groovy source tree. This is the detail that tripped me up for a while.

I initially setup a separate GroovyCompile task to process the annotation before the rest of the source (stemming from a helpful suggestion from Peter Niederwieser on the Gradle forums). While this worked, a much simpler solution for getting transformations to apply is to setup a multi-project build. The main project depends on a sub-project with the ast transformation source files. Here's a minimal example's directory structure

ast/build.gradle
ast build file
ast/src/main/groovy/com/cholick/ast/Marker.groovy
marker interface
ast/src/main/groovy/com/cholick/ast/Transform.groovy
ast transformation
build.gradle
main build file
settings.gradle
project hierarchy configuration
src/main/groovy/com/cholick/main/Main.groovy
source to transform

For the full working source (with simple tests and no * imports), clone https://github.com/cholick/gradle_ast_example

The root build.gradle file contains a dependency on the ast project:

dependencies {
    ...
    compile(project(':ast'))
}

The root settings.gradle defines the ast sub-project:

include 'ast'

The base project also has src/main/groovy/com/cholick/main/Main.groovy, with the source file to transform. In this example, the ast transformation I've written puts a method named 'added' onto the class.

package com.cholick.main

import com.cholick.ast.Marker

@Marker
class Main {
    static void main(String[] args) {
        new Main().run()
    }
    def run() {
        println 'Running main'
        assert this.class.declaredMethods.find { it.name == 'added' }
        added()
    }
}

In the ast sub-project, ast/src/main/groovy/com/cholick/ast/Marker.groovy defines an interface to mark classes for the ast transformation:

package com.cholick.ast

import org.codehaus.groovy.transform.GroovyASTTransformationClass

import java.lang.annotation.*

@Retention(RetentionPolicy.SOURCE)
@Target([ElementType.TYPE])
@GroovyASTTransformationClass(['com.cholick.ast.Transform'])
public @interface Marker {}

Finally, the ast transformation class processes source classes and adds a method:

package com.cholick.ast

import org.codehaus.groovy.ast.*
import org.codehaus.groovy.ast.builder.AstBuilder
import org.codehaus.groovy.control.*
import org.codehaus.groovy.transform.*

@GroovyASTTransformation(phase = CompilePhase.INSTRUCTION_SELECTION)
class Transform 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 != Marker.class.name) return

        ClassNode annotatedClass = (ClassNode) astNodes[1]
        MethodNode newMethod = makeMethod(annotatedClass)
        annotatedClass.addMethod(newMethod)
    }
    MethodNode makeMethod(ClassNode source) {
        def ast = new AstBuilder().buildFromString(CompilePhase.INSTRUCTION_SELECTION, false,
                "def added() { println 'Added' }"
        )
        return (MethodNode) ast[1].methods.find { it.name == 'added' }
    }
}

Thanks Hamlet D'Arcy for a great ast transformation example and Peter Niederwieser for answering my question on the forums.