Playing with Groovy AST Transformations

In the current project that I’m working on right now, the base backend programing language is Groovy. I didn’t have the chance to work with it before, so it was a great opportunity to learn it and find its secrets. The good news is that I had experience with Java so the syntax was pretty much similar, also Groovy runs on top of the JVM, which means that is compiled into JVM bytecode. Besides that it is great that you can add Groovy code to existing Java applications cause the generated bytecode is totally compatible with the one generated by Java. You can use existing powerful Java libraries out of the box, even not specify data types so at the end it behaves like a dynamic language, it also performs type inference. One of the things that amazed me was the ability to modify a classes at compilation time, before bytecode generation.

AST Transformations

AST means Abstract Syntax Tree which is composed of nodes that correspond to Groovy language constructs. The tree structure lends itself to process using visitor design pattern. An AST Transformation is a compiler hook that Groovy provides into the compilation process, allows the manipulation of the AST during compilation prior to bytecode generation. There are two types of AST transformations, local and global. Local are more common, are annotation driven and indicates the AST transformation to be applied. AST is walked and AST transformation applied to nodes that are annotated with transformation annotation. Groovy comes with many local AST transformation like @ToString, @EqualsAndHashCode, etc. Global are less common, are applied to every source unit in compilation. Uses jar file service provider mechanism to identify global AST transformations.

Implementing AST Transformation

One of the requirements of the current project was to be able to internationalize some text fields of the domain classes, allowing multi language support. When I read the requirement, instantly came into my mind the idea of implementing it using a local transformation, marking those fields to be i18n with an annotation. As an example, will present a domain class with some fields that will become multi language, here is the entity.

class Card {

    String type
    String code
    @I18N
    String title
    @I18N
    String text
    
}

As you can see, I marked the title and text fields as internationalizable, using a custom @I18N annotation that is defined below.

/**
 * Annotation used to mark a class field as multi language.
 * It will aware the I18N AST Transformer to add an extra field to support multi language.
 * The new field will be an instance of {@link com.example.I18NMap}
 */
@Retention(RetentionPolicy.SOURCE)
@Target(ElementType.FIELD)
@GroovyASTTransformationClass("com.example.I18NASTTransformation")
public @interface I18N {

    /**
     * The name of the field to be added, otherwise the name will be the same
     * as the annotated field name prepending an i letter. i.e: iFirstName
     * @return
     */
    String fieldName() default ''
}

@GroovyASTTransformationClass is responsible of telling the compiler what ASTTransformation should be applied when finding @I18N annotation. The fieldName allows the developer to specify the name that the multi language field will have, next is the corresponding transformation class that will perform the work of adding the new internationalizable fields. It implements org.codehaus.groovy.transform.ASTTransformation

/**
 * AST Transforamtion class in charge of adding a new multi language field
 * on the fields that were marked with {@link com.example.I18N}
 */
@GroovyASTTransformation(phase = CompilePhase.CANONICALIZATION)
class I18NASTTransformation implements ASTTransformation {

    @Override
    void visit(ASTNode[] astNodes, SourceUnit sourceUnit) {
        if (safeToAddProperty(astNodes)) {
            def annotationNode = (AnnotationNode)astNodes[0]
            def fieldNode = (FieldNode)astNodes[1]
            addProperty(annotationNode, fieldNode)
        }
        else {
            throw new CompilationFailedException("@I18N annotation failed")
        }
    }

    private void addProperty(AnnotationNode annotationNode, FieldNode fieldNode) {
        def classNode = fieldNode.owner
        def propertyName = getFieldName(annotationNode, fieldNode)
        if(!containsField(classNode, propertyName)) {
            def field = new FieldNode(propertyName, FieldNode.ACC_PRIVATE, new ClassNode(LanguageMap), new ClassNode(classNode.class), null)
            classNode.addProperty(new PropertyNode(field, PropertyNode.ACC_PUBLIC, null, null))
        }
    }

    private String getFieldName(AnnotationNode annotationNode, FieldNode fieldNode) {
        def fieldName = annotationNode.getMember('fieldName')?.value
        return fieldName ?: "i${capitalizeFirst(fieldNode.name)}".toString()
    }

    private Boolean safeToAddProperty(ASTNode[] astNodes) {
        def invalid = (!astNodes || astNodes.length != 2 || !(astNodes[1] instanceof FieldNode) || ((FieldNode)astNodes[1]).type.name != String.name ||
            !(astNodes[0] instanceof AnnotationNode) || ((AnnotationNode)astNodes[0]).classNode.name != I18N.name)
        return !invalid
    }

    private boolean containsField(ClassNode classNode, String propertyName) {
        classNode.fields.find { fieldNode ->
            fieldNode.name == propertyName
        } != null
    }
    
    private static String capitalizeFirst(str) {
        char[] array = str.toCharArray()
        array[0] = Character.toUpperCase(array[0])
        return new String(array)
    }
}

The new fields are instances of I18NMap class, where the key is the iso code and the value is the text in its native language.

/**
 * Class which purpose is to store multi language strings
 * key: iso6391Code
 * value: text in native language
 */
class I18NMap extends HashMap<String,String> {

    private static final String SEPARATOR = ":"

    public I18NMap() {
        super()
    }

    public I18NMap(List values) {
        init(values)
    }

    public I18NMap(String[] values) {
        init(values)
    }

    private init(values) {
        values.each { value ->
            def lang = decodeValue(value)
            if(lang.size() == 2) this.put(lang[0], lang[1])
        }
    }
    
    private String[] decodeValue(String value) {
        if (value == "") return []
        def pair = value.split(SEPARATOR, 2)
        if (pair.size() != 2) {
            throw new RuntimeException("invalid language map entry: $value")
        }
        return pair
    }
}

Hope with this little example you find AST Transformations something useful, to take into account in some situations.