Groovy AST Transformation - AOP Style

In my last blog about AST Transformation I followed a sample example and created AssertParamsNotNull local transformation. Further, I wanted to create AST Transformation which is generic enough and performs any check before method is executed. In process, goal is to learn about AST Transformation.

BeforeAdvisor AST Transformation
Use case for BeforeAdvisor AST involves:

• Authorization Checking - Security by checking role from context
• Print Parameter values with which the method is called
• Asserts Parameters are not null
• Check various entry-conditions/Pre-Conditions of the method

To make transform generic enough, idea is to inject a method call to before method of advice for each method annotated with @BeforeAdvisor(MyPreConditionAdvice). Advice method can choose to implement any checks/conditions to be performed before actual methods gets executed.

Few of the subtle characteristics of this type of solution are

• It does not allow changing the method parameters. 
  
• It does not allow to stop execution of method. However, you can throw runtime exception.
• Advice needs no arg constructor and must implement method before

Usage and Client
Starting with class that will use this transform. Following defines script-level method with annotation having advice information.

package com.learn.sts.groovy.ast

@com.learn.sts.groovy.ast.BeforeAdvisor(value= com.learn.sts.groovy.MyAdvice)
def sayHello(name, name2)
{
    println "Hello " + name + name2
}

sayHello("World", "Groovy") 

Sample Advice
The advice that we will try to invoke will be something like this. This advice just prints parameter value that method is being invoked with. But you can implement any of the use cases described above.

package com.learn.sts.groovy

public class MyAdvice
{
    def before(String methodName, List listArg)
    {
        println 'Entering Method ' + methodName + ' with params ' + listArg
    }
}

Annotation
Now lets create the Annotation

package com.learn.sts.groovy.ast;

import java.lang.annotation.Retention
import java.lang.annotation.RetentionPolicy
import java.lang.annotation.Target
import java.lang.annotation.ElementType
import org.codehaus.groovy.transform.GroovyASTTransformationClass
@Retention(RetentionPolicy.SOURCE)
@Target([ElementType.METHOD])
@GroovyASTTransformationClass(["com.learn.sts.groovy.ast.BeforeAdvisorASTTransformation"])
public @interface BeforeAdvisor {
    Class value ();
}

One difference to this annotation is that it declares a method value(). This allows us to get value being passed during annotation declaration on method.

AST Transformation

package com.learn.sts.groovy.ast

//Imports section skipped for brevity

@GroovyASTTransformation(phase=CompilePhase.SEMANTIC_ANALYSIS)
public class BeforeAdvisorASTTransformation implements ASTTransformation 
{
    public void visit(ASTNode[] nodes, SourceUnit source)
    {
        AnnotationNode node = (AnnotationNode) nodes[0];
        final Expression classNode = node.getMember("value")
        
        List<MethodNode> allMethods = source.AST?.classes*.methods.flatten()
        List annotatedMethods = allMethods.findAll{ MethodNode method ->
                 method.getAnnotations(new ClassNode(BeforeAdvisor))
                }
        annotatedMethods.each{MethodNode method -> 
             List existingStatements = method.getCode().getStatements()
            Parameter[] parameters = method.getParameters()
            int i = 0;
            existingStatements.add(i++, initAdviceCall(classNode))
            existingStatements.add(i++, createMethodCall(method, parameters))
        }
      }
  
    public Statement initAdviceCall(classNode)
    {
      return new ExpressionStatement(
              new DeclarationExpression(
                      new VariableExpression("advice"),
                      new Token(Types.ASSIGNMENT_OPERATOR, "=", -1, -1),
                      new ConstructorCallExpression(classNode.getType(), new ArgumentListExpression())
                      )
              )
    }
 
  public Statement createMethodCall(method, Parameter[] parameters){
      List parameterExpressionList = new ArrayList()
      parameters.each{ parameter -> parameterExpressionList.add(new VariableExpression(parameter))}
      return new ExpressionStatement(
        new MethodCallExpression(
            new VariableExpression("advice"),
            "before",
            new ArgumentListExpression(new ConstantExpression(method.getName()),
                    new ListExpression(parameterExpressionList)
            )
        )
      )
  }
}

Couple of things to notice here:

First, creation of the AST tree involves

• Creating instance of Advice Class - This is done through method initAdviceCall
  
• Invocation of the method before with methodName and Parameter List - This is done through method createMethodCall

Creating AST structure is not the easiest task. One approach that has worked for me is the write the sample code that you are trying to generate AST for and then use Groovy AST viewer in eclipse.

Second, inspecting the value on ASTNode and getting Advice class value. First node contains information about the annotation and second node is the annotated node.

Lessons Learned

Compilation
It was required to compile Transformations files first and then compile the classes that use AST Transformation. If I compiled all three of them together the AST Transformation did not kick in. I used Eclipse IDE and it forced me to use Compile Groovy File explicitly each time I made change. "Build Automatically" or "Clean" option did not work.

Script Level Methods vs all Class Methods         
In AssertParamsNotNull AST Transformation in previous blog, I used source.getAST()?.getMethods()  what it did is that it only found top level (Script Level) methods. So AST Transformation did not apply to Class Methods, it only got applied to Script Level Method. For this one I changed to source.ast?.classes*.methods.flatten() (Line 12 in BeforeAdvisorASTTransformation). This also become apparent once you see AST structure for a given class using Groovy AST View in eclipse.

Retrive value from annotations
ASTNode being passed to visit method carries information about the annotation. First element (node[0]) contains information about annotation. You can get the values passed in to annotation during the visit method and use it during the transformation. In case above we get class passed in as value and it gets instantiated and before method gets invoked.

Groovy Compilation
When groovy compiler is invoked, any sourcefile.groovy goes under series of transformations.
From Source --> ANTLR Tokens --> ANTLR AST --> Groovy AST --> Bytecode

Using AST Transformation, we manipulate the way groovy AST gets generated. It allows to insert additonal statements.

There are various CompilerPhase that goes along with this process. I couldn't find much documentation on the process. Best information is found on Jochen Theodorou's blog post

Blogged with the Flock Browser