Truffle言語で関数呼び出しを実装する

Truffleを使って言語を実装してみたのだけど、やはり関数呼び出しもやりたい。
Truffleでの言語実装を最小手数ではじめる - きしだのはてな
簡易Truffle言語に変数を実装する - きしだのはてな

とりあえず関数

今回、関数は埋め込みで呼び出しだけ行ってみます。ということで、とりあえずランダムを返す関数を実装。

public abstract class RandNode extends MathNode {
    static Random rand = new Random();

    @Specialization
    public long rnd() {
        return rand.nextInt(10);
    }
}


ちなみに、native-imageするときstaticイニシャライザはコンパイル時に実行されるので、なにも考えずにネイティブイメージを作ると乱数が固定されて毎回同じ値を返します。
--delay-class-initialization-to-runtime=RandNodeをつける必要があります。
Understanding Class Initialization in GraalVM Native Image Generation

CallTargetを保持する

関数としてはCallTargetが欲しいので、CallTargetを保持するクラスを作ります。

class MathFunction {
    private final RootCallTarget callTarget;

    public MathFunction(RootCallTarget callTarget) {
        this.callTarget = callTarget;
    }    
}


CallTargetを作るときにRootNodeが必要になるのだけど、前回作ったものは変数定義を入れてしまったので、改めてRootNodeを作っておきます。

class FuncRootNode extends RootNode {
    MathNode function;

    public FuncRootNode(TruffleLanguage<?> language, FrameDescriptor frameDescriptor,
            MathNode function) {
        super(language, frameDescriptor);
        this.function = function;
    }

    @Override
    public Object execute(VirtualFrame frame) {
        return function.executeGeneric(frame);
    }
}


そしたら、こんな感じでパースのときにCallTargetを作って保持しておきます。

static MathFunction RAND_FUNC;
MathFunction createBuiltin(FrameDescriptor frame) {
    if (RAND_FUNC == null) {
        RandNode rand = RandNodeGen.create();
        RAND_FUNC = new MathFunction(Truffle.getRuntime().createCallTarget(
                new FunctionNodes.FuncRootNode(this, frame, rand)));
    }
    return RAND_FUNC;
}

実際には関数名->関数オブジェクトのMapに保持することになると思います。

呼び出し

呼び出しは、CallTargetのcallメソッドを呼び出すと行えます。

class MathInvokeNode extends MathNode {
    MathFunction function;

    public MathInvokeNode(MathFunction function) {
        this.function = function;
    }

    @Override
    Object executeGeneric(VirtualFrame frame) {
        return function.callTarget.call();
    }
}

今回はつかってないけど引数があるときはcallの引数で渡します。受け取り側では、VirtualFrameのgetArguments()で。


パースのときに数値以外は変数として扱うようにしたのだけど、randと書くと乱数関数を呼び出すことにします。

    MathNode parseNode(FrameDescriptor frame, String value) {
        try {
            return LongNode.of(value);
        } catch (NumberFormatException ex) {
            if ("rand".equals(value)) {
                return new MathInvokeNode(createBuiltin(frame));


これでこんな感じで動くようになります。

String exp = "12+34+56+aa+rand";
MathCommand.main(new String[]{exp});

呼び出しの最適化

さて、メソッド呼び出しを効率化しましょう。
IndirectCallNode/DirectCallNodeを使います。ここではDirectCallNodeを。

static class InvokeNode extends MathNode {
    DirectCallNode callNode;

    public InvokeNode(MathFunction function) {
        callNode = Truffle.getRuntime()
                .createDirectCallNode(function.callTarget);
    }

    @Override
    Object executeGeneric(VirtualFrame frame) {
        return callNode.call(new Object[0]);
    }
}

これでいいんだけど、実際に言語を作るとき、これだと再帰関数ではこの時点ではcallTargetが準備できてなかったりしてうまういかないことがあります。
なので遅延処理することになります。

static class InvokeNode extends MathNode {
    MathFunction function;
    DirectCallNode callNode;

    public InvokeNode(MathFunction function) {
        this.function = function;
    }

    @Override
    Object executeGeneric(VirtualFrame frame) {
        if (callNode == null) {
            callNode = Truffle.getRuntime().createDirectCallNode(function.callTarget);
        }
        
        return callNode.call(new Object[0]);
    }
}

けど、これ動くには動くけど遅くなるしGraalVMで動かすと文句いわれる。
ということで、@CompilationFinalをつけます。そして値変更時にCompilerDirectives.transferToInterpreterAndInvalidate()でGraalコンパイラに教えてあげる。

static class MathInvokeNode extends MathNode {
    MathFunction function;
    @CompilationFinal
    DirectCallNode callNode;

    public MathInvokeNode(MathFunction function) {
        this.function = function;
    }

    @Override
    Object executeGeneric(VirtualFrame frame) {
        if (callNode == null) {
            CompilerDirectives.transferToInterpreterAndInvalidate();
            callNode = Truffle.getRuntime().createDirectCallNode(function.callTarget);
        }
        
        return callNode.call(new Object[0]);
    }
}


今回のコードだと処理速度に影響はないけど、実際の言語実装したときに10%~30%くらい速くなった。