Memoization with Groovy Custom AST Tranformation

Main.groovy

package akihiro4chawon

class Main {
    @Memoize
    def fib(a) {
        //a <= 1 ? a : fib(a - 2) + fib(a - 1)
        if (a <= 1) return a
        else return fib(a - 2) + fib(a - 1)
    }

    public static void main(String[] args) {
        println(new Main().fib(40))
    }
}

Memoize.groovy

package akihiro4chawon

import org.codehaus.groovy.transform.GroovyASTTransformationClass
import java.lang.annotation.*

@Retention (RetentionPolicy.SOURCE)
@Target ([ElementType.METHOD])
@GroovyASTTransformationClass ('akihiro4chawon.MemoizeTransformation')
public @interface Memoize { }

MemoizeTransformation.groovy

package akihiro4chawon

import org.codehaus.groovy.ast.*
import org.codehaus.groovy.ast.expr.*
import org.codehaus.groovy.ast.stmt.*
import org.codehaus.groovy.ast.builder.AstBuilder
import org.codehaus.groovy.control.*
import org.codehaus.groovy.control.messages.Message
import org.codehaus.groovy.transform.*
import static org.objectweb.asm.Opcodes.*

@GroovyASTTransformation(phase = CompilePhase.INSTRUCTION_SELECTION)
public class MemoizeTransformation implements ASTTransformation {
    def traceOut(s) {javax.swing.JOptionPane.showMessageDialog(null, s)}

    void visit(ASTNode[] astNodes, SourceUnit sourceUnit) {
        // defensive programming
        if (!astNodes || !astNodes[0] || !astNodes[1]) return
        if (!(astNodes[0] instanceof AnnotationNode)) return
        if (!(astNodes[1] instanceof MethodNode)) return

        // validate AnnotationNode
        MethodNode annotatedMethod = astNodes[1]
        if (annotatedMethod.parameters.length == 0) {
            def errMsg = 'cannot @Memoize a method with no parameters'
            sourceUnit.errorCollector.addError(
                    Message.create("$errMsg: $annotatedMethod.name", sourceUnit))
            return
        }
        if (annotatedMethod.parameters.length >= 2) {
            def errMsg = 'cannot @Memoize a method with multiple parameters'
            sourceUnit.errorCollector.addError(
                    Message.create("$errMsg: $annotatedMethod.name", sourceUnit))
            return
        }
        if (annotatedMethod.returnType.name == "void") {
            def errMsg = 'cannot @Memoize a void method'
            sourceUnit.errorCollector.addError(
                    Message.create("$errMsg: $annotatedMethod.name", sourceUnit))
            return
        }

        ClassNode declaringClass = annotatedMethod.declaringClass
        makeMethodMemoized(declaringClass, annotatedMethod)

        traceOut("Annotation OK!");
    }

    private void aliasMethod(String alias, MethodNode org) {
        org.declaringClass.addMethod(
            new MethodNode(alias, org.modifiers, org.returnType, org.parameters, org.exceptions, org.code))
    }

    private void makeMethodMemoized(ClassNode classNode, MethodNode methodNode) {
traceOut("Entering makeMethodMemoized:\n\t" + classNode.dump() + "\n\t" + methodNode.dump())

        // add field of memo map
        def memoFieldName = "memo${methodNode.name}"
        classNode.addField("$memoFieldName", ACC_PRIVATE, new ClassNode(Map.class),
            new ConstructorCallExpression(new ClassNode(HashMap.class), new ArgumentListExpression()))

        // 本体を保存する
        def methodName = methodNode.name
        def methodBackupName = "${methodName}WithoutMemo"
        aliasMethod(methodBackupName, methodNode)
traceOut("\tmethod chaining completed\n")

        Parameter[] params = methodNode.getParameters()
        String parameterName = params[0].getName()

traceOut("\tabout to make memoizing code\n")
        // メモ化処理をするラッパーコード
        def memoizedValueName = 'memoizedValue'
        def computedValueName = 'computedValue'
        def memoizingCode = new AstBuilder().buildFromSpec { block {
            expression {
                declaration {
                    variable memoizedValueName
                    token "="
                    methodCall {
                        variable memoFieldName
                        constant 'get'
                        argumentList {
                            variable parameterName
                        }
                    }
                }
            }
            ifStatement { // if memoized then {return memoized_value} else {compute -> store -> return}
                booleanExpression { // condition clause
                    variable memoizedValueName
                }
                returnStatement {   // if-body
                    variable memoizedValueName
                }
                block {             // else-body
                    expression {
                        declaration {
                            variable computedValueName
                            token '='
                            methodCall {
                                variable "this"
                                constant methodBackupName
                                argumentList {
                                    variable parameterName
                                }
                            }
                        }
                    }
                    expression {
                        methodCall {
                            variable memoFieldName
                            constant 'put'
                            argumentList {
                                variable parameterName
                                variable computedValueName
                            }
                        }
                    }
                    returnStatement {
                        variable computedValueName
                    }
                }
            }
        }}
traceOut("\tmake memoizing code finished\n")
        methodNode.setCode(memoizingCode.head())
    }
}