Javassist 3.1とjavaagentでmemorization

プログラミング

( 派生クラスでのメモ化は[id:bellbind:20050910:p2] )

Javassist 3.1はJava5のannotationが取り出せるようになっていました。

これとJava5のツールインタフェースであるjavaagentを組み合わせればmemorize機能は実現できました。

以下コード。

アノテーションMemorize.java

import java.lang.annotation.*;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Memorize {}

HashMapのキーに引数を突っ込むためのラッパーArgs.java

import java.util.*;

public class Args {
    private Object[] args;
    public Args(Object[] args) {
        this.args = args;
    }
    
    public int hashCode() {
        return Arrays.deepHashCode(this.args);
    }
    
    public boolean equals(Object o) {
        if (o instanceof Args) {
            return Arrays.deepEquals(this.args, ((Args) o).args);
        }
        return false;
    }
}

AgentクラスMemorizeAgent.java

import java.lang.instrument.*;
import java.security.*;
import java.io.*;
import javassist.*;

public class MemorizeAgent implements ClassFileTransformer {
    private Instrumentation inst;
    private ClassPool pool;
    private CtClass memoClass;
    private CtField.Initializer memoInitializer;
    
    public MemorizeAgent(Instrumentation inst) throws Exception {
        this.inst = inst;
        this.pool = new ClassPool();
        this.pool.appendSystemPath();
        this.memoClass = this.pool.get("java.util.HashMap");
        this.memoInitializer = CtField.Initializer.byExpr("new java.util.HashMap()");
    }
    
    public static void premain(String agentArgs, Instrumentation inst) throws Exception {
        inst.addTransformer(new MemorizeAgent(inst));
    }
    
    public byte[] transform(ClassLoader loader, String className,
                              Class<?> classBeingRedefined,
                              ProtectionDomain protectionDomain,
                              byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            return this.getMemorizedClass(classfileBuffer);
        } catch (Exception ex) {
            throw new IllegalClassFormatException(ex.getMessage());
        }
    }
    
    private byte[] getMemorizedClass(byte[] classfileBuffer) throws Exception {
        ByteArrayInputStream istream = new ByteArrayInputStream(classfileBuffer);
        CtClass newClass = this.pool.makeClass(istream);
        int index = 0;
        for (CtMethod method: newClass.getMethods()) {
            if (!this.isMemorizedMethod(method)) continue;
            
            // add memo field
            String fieldName = "$_memo_map_$" + index;
            CtField memoField = new CtField(memoClass, fieldName, newClass);
            newClass.addField(memoField, memoInitializer);
            
            // escape original method
            CtMethod bodyMethod = new CtMethod(method, newClass, null);
            String bodyMethodName = "$_memo_method_$" + index;
            bodyMethod.setName(bodyMethodName);
            newClass.addMethod(bodyMethod);
            
            // replace memorized body
            String code = "{" +
              "Args args = new Args($args);" +
              "Object result = " +  fieldName + ".get(args); " +
              "if (result != null) return ($r) result;" +
              "result = ($w) " + bodyMethodName + "($$);" +
              fieldName + ".put(args, result);" +
              "return ($r) result;" +
              "}";
            method.setBody(code);
            
            index++;
        }
        return newClass.toBytecode();
    }
    
    private boolean isMemorizedMethod(CtMethod method) throws Exception {
        for (Object annotation: method.getAnnotations()) {
            if (annotation instanceof Memorize) return true;
        }
        return false;
    }
}

置き換えるボディのコードは試行錯誤で理解できました：


javaagent用マニフェストファイルmemorizeagent.mf

Premain-Class: MemorizeAgent


そしてフィボナッチクラスなどMain.java

public class Main {
    public static void main(String[] args) {
        Fib fib = new Fib();
        System.out.println(fib.fib(80));
    }
}

class Fib {
    @Memorize
    public long fib(int n) {
        System.out.println("when n=" + n);
        if (n < 2) return 1;
        return fib(n - 1) + fib(n - 2);
    }
}

これは普通にメモ化せずに実行すると多分終わらない。

ビルド＆実行

$ jar cvfm memorizeagent.jar memorizeagent.mf
$ javac -classpath javassist.jar;. *.java
$ java -classpath javassist.jar;. -javaagent:memorizeagent.jar Main
when n=80
when n=79
when n=78
...
when n=1
when n=0
37889062373143906

javaagentで渡すjarはマニフェストファイルが入ってるだけでよい。

AspectJでやるならこれ以上短くかけないとダメでしょうね。

Python2.4のdecorator版

これが一番シンプルにみえる。decoratorのコールは@のコード位置ごとなのかな。するとオーバーライドしたら消える？