Writing a Lisp JIT Interpreter with GraalVM Truffle

Nicolas Laurent, a TruffleRuby contributor has written a great introductory tutorial for Truffle, but in case you want a TL;DR:

• Graal is a Java Just-In-Time compiler (written in Java),
• and Truffle is a language implementation framework that lets you build a JIT compiler by just writing an interpreter.

How? One may also think of Truffle as a way to instruct Graal to perform super aggressive inlining (called “partial evaluation or Futamura projections“). For example, (+ 1 2) in a Lispy language may parse into an AST as follows:

Node literal1 = new Node() { int execute() { return 1; } };
Node literal2 = new Node() { int execute() { return 2; } };
Node addNode = new Node() {
        @Child Node lhs = literal1;
        @Child Node rhs = literal2;

        int execute() {
            return lhs.execute() + rhs.execute();
        }
    };

// execute the program
addNode.execute(); // 1 + 2 => 3

Upon compilation, the compiler might attempt to inline, like, everything:

That is, having Graal JIT-compile your AST interpreter is essentially JIT-compiling the interpreted program: we go from an AST to a compiled, constant-folded form.

Let’s have a look at a lengthier example:

class BytecodeInterpreter extends Node {
    @CompilationFinal(dimensions = 1)
    final byte[] instructions = { REF_ARG, ADD_CONST, -100, RETURN };
    // the bytecode is more or less: (lambda (x) (+ x -100))

    @BytecodeInterpreterSwitch
    @ExplodeLoop(kind = LoopExplosionKind.MERGE_EXPLODE)
    int execute(VirtualFrame frame) {
        int pc = 0;
        while (true) {
            byte instruction = instructions[pc++];
            switch (instruction) {
            case REF_ARG:
                frame.setInt(0, (Integer) frame.getArguments()[0]);
                break;
            case ADD_CONST:
                frame.setInt(0, frame.getInt(0) + instructions[pc++]);
                break;
            case RETURN:
                return frame.getInt(0);
            }
        }
    }
}

// execute the program
VirtualFrame frame = createTheFrameSomehow();
new BytecodeInterpreter().execute(frame);

The code above is a barebone bytecode interpreter implemented with Truffle. @CompilationFinal tells Graal to treat instructions as constants when compiling, and @ExplodeLoop tells Graal to unroll loops aggressively. With these, we now have a simplistic bytecode JIT compiler, and the bytecode above partial-evaluates to:

{
    int pc = 0;
    int instruction = REF_ARG; // instructions[pc] -> instructions[0] -> REF_ARG
    pc++;
    switch (instruction) {
    case REF_ARG:
        frame.setInt(0, (Integer) frame.getArguments()[0]);
    }

    instruction = ADD1;        // instructions[pc] -> instructions[1] -> ADD1
    pc++;
    switch (instruction) {
    case ADD1:
        frame.setInt(0, frame.getInt(0) + (-100)); // instructions[2] -> -100
        pc++;
    }

    instruction = RETURN;      // instructions[pc] -> instructions[3] -> RETURN
    switch (instruction) {
    case RETURN:
        return frame.getInt(0);
    }
    // loop unroll stopped by the return statement
}

And the above further simpifies into:

{
    frame.setInt(0, (Integer) frame.getArguments()[0]);
    frame.setInt(0, frame.getInt(0) + (-100));
    return frame.getInt(0);
}

Depending on the stack frame type, Graal might be able to recognize frame.getInt(0) as a local variable and make further optimization:

{
    return (Integer) frame.getArguments()[0] + (-100);
}

Graal offers a tool called Ideal Graph Visualizer (IGV) to visualize compilation graphs. The following graph was produced by Graal running said bytecode:

• Up until the 127 Unbox node, the program reads in a function parameter.
• The 149 + node adds a constant -100 to the unboxed result.
• The 200 BoxNode and 195 Return nodes box the sum and return.

Exactly what we’ve described above.

Partial evaluation alone is not enough for an efficient JIT implementation. For example, we expect (+ 1 2) and (+ 1.0 2.0) to compile to different instructions (maybe ADD and FADD respectively), but for dynamic languages, we don’t always know the argument types at compile time. And this is exactly what speculative compilation (or speculation) can help with.

Basically, without speculation, the compiler might compile (+ x y) into:

if (x instanceof Integer xi && y instanceof Integer yi) {
    return xi + yi;
}
if (x instanceof Float xf && y instanceof Float yf) {
    return xf + yf;
}
// maybe tons of other cases
// - float + int
// - int + float
// - int/float + bigint
// - bigint + int/float/bigint
throw new UnsupportedOperationException();

With speculation, the compiler can somehow guess from previous calls what types x, y are of, and put irrelevant paths in a separate interpreted slow path:

if (!(x instanceof Integer xi) || !(y instanceof Integer yi)) interpretedSlowPath();
return xi + yi;

The interpretedSlowPath usually involves invalidating the current compilation and re-specializing the AST node. Truffle languages usually achieve this by using a mutable @CompilationFinal state field and a special compiler directive transferToInterpreterAndInvalidate():

class SpeculativeAdd extends Node {
    static final int STATE_INT_INT = 1;
    static final int STATE_FLOAT_FLOAT = 2;
    @CompilationFinal
    int state = 0;

    Object execute(Object lhs, Object rhs) {
        if (state == STATE_INT_INT) {
            if (lhs instanceof Integer l && rhs instanceof Integer r) {
                return l + r;
            }
        }
        if (state == STATE_FLOAT_FLOAT) {
            if (lhs instanceof Float l && rhs instanceof Float r) {
                return l + r;
            }
        }
        // other specializations: int + float, float + int, bigints, etc.

        return specialize(lhs, rhs);
    }

    Object specialize(Object lhs, Object rhs) {
        CompilerDirectives.transferToInterpreterAndInvalidate();
        if (lhs instanceof Integer l && rhs instanceof Integer r) {
            state = STATE_INT_INT;
            return l + r;
        }
        if (lhs instanceof Float l && rhs instanceof Float r) {
            state = STATE_FLOAT_FLOAT;
            return l + r;
        }
        // other specializations: int + float, float + int, bigints, etc.

        throw new UnsupportedOperationException();
    }
}

Whenever there is an active state, the compiler treats state as @CompilationFinal and compiles only the active path upon compilation; whenever the state assumptions are invalidated, the specialize method is responsible for invalidating the current compilation and adjusting the state so that the node can be recompiled correctly.

This pattern allows us to realize tons of advanced JIT techniques, including:

• Inline caching: For costly operations, if the inputs are known and vary infrequently, it might be best to cache the results/intermediate products. For example, global variable lookup can really use the technique - instead of consulting a hash table every single time you access a variable, cache the value container.
• Polymorphism: In (+ 1 x), when x can sometimes be an integer and sometimes a floating point number, instead of re-compiling every time the type changes, it is better to just compile these two paths (out of many others) and switch to different paths at runtime.

• Splitting: A highly polymorphic AST graph, however, can impact performance. One way to handle this is making the some parent nodes polymorphic by deep copying the sub-tree.

  

However, it is definitely not fun having to hand-rolling all these ourselves. Instead, Truffle provides a powerful DSL (read: Java annotations) that generates all the boilerplate for us. Using the DSL, The SpeculativeAdd node above is as simple as a few @Specialization:

abstract class MyLanguageExpressionNode extends Node {
    /// An abstract "executeXXX" method to implemented by the DSL
    public abstract Object execute(VirtualFrame frame);
}

@NodeChild(value = "lhs", type = MyLanguageExpressionNode.class)
@NodeChild(value = "rhs", type = MyLanguageExpressionNode.class)
abstract class SpeculativeAdd extends MyLanguageExpressionNode {
    @Specialization
    public static int addInts(int l, int r) { return l + r; }
    @Specialization
    public static float addFloats(float l, float r) { return l + r; }
    @Specialization
    public static float addFloatInt(float l, int r) { return l + r; }
    @Specialization
    public static float addIntFloat(int l, float r) { return l + r; }
}

The code above asks Truffle annotation processor to generate a SpeculativeAddNodeGen class implementing specialization and polymorphism logic.

The above specializes the AST nodes. But to supply values to the nodes, you will need variables. It turns out that Truffle offers VirtualFrame exactly for this purpose. Storing local variables on a Frame, you can specialize the slots to Object or primitive types so as to remove boxing/unboxing costs. See SLWriteLocalVariableNode.java and SLReadLocalVariableNode.java for how easily one can speculate about unboxed storage of local variables.

Lisp languages offer macros for metaprogramming and dynamic AST generation. For example, an unless form may expand to:

;; macro form
(unless noninteractive
  (message "Hello World!"))
;; expanded form
(if noninteractive
    nil
  (message "Hello World!"))

Emacs Lisp chooses to expand macro at runtime. Using Truffle’s replace method, it is quite straightforward to come up with a MacroNode:

class MacroNode extends MyLanguageExpression {
    final Object macroFunction;
    final Object[] macroArgs;

    @Override
    Object execute(VirtualFrame frame) {
        Object expanded = getFunction(macroFunction).call(macroArgs);
        return replace(objectToAST(expanded)).execute(frame);
    }
}

However, runtime macro evaluation brings some implications impacting the overall interpreter design:

• AST nodes are created lazily.

  Whether a node is a macro node or a function call node affects the AST. In the unless example above, the AST is as follows:

  (macro-node
   (literal unless)
   (literal noninteractive)
   (literal (message "Hello World!"))
   )
  

  All arguments are passed “raw” for macro nodes. If unless were a function, arguments are evaluated before the function call and the AST would be:

  (call-node
   (literal unless)
   (read-variable noninteractive)
   (call-or-macro-node
    (message "Hello World!"))
   )
  

  Now, with runtime macro expansion, we cannot know for centain (at declaration time) whether (something a b) is a call-node or a macro-node. And the only way is to have a call-or-macro-node that dynamically switches to call/macro nodes at runtime.

• Static analysis is doomed.

  Truffle frees us from doing register allocation, but we still need to allocate a frame slot for every local variable ourselves:

  (lambda (x)          ; x -> frame slot #0
    (let ((y (* x 2))  ; y -> frame slot #1
          (z (* x x))) ; z -> frame slot #2
      (+ x y z)))
  

  For some languages, slot allocation can be done at parse time, but not with runtime macros:

  (pcase-let ((`(,_ ,_ ,uid ,gid)
               (file-attributes user-init-file)))
    (list uid gid))
  

  pcase-let is a macro that introduces new let forms, declaring new variables. Since macros expand at runtime, it means we also assign slot numbers to those variables at runtime.

  This is fine: we can always introduce more states in our AST and keeping a slot count is not a big deal. The only problem is that Truffle requires knowing the frame size (total slot counts) at declaration time, which I don’t think is possible with Emacs Lisp. So, we need another trick to handle this.

In Lisps, control flows, quotes and statement blocks are typically implemented as “special forms”:

(if condition
    (progn
      (statement1)
      (statement2))
  (else-statement1)
  (else-statement2))

In some lisps, these special forms are more-or-less identified by their corresponding symbols, in that you cannot define your own my-if by aliasing my-if to if. Well, this isn’t the case with Emacs:

(defalias 'my-if #'if) ; now you can use 'my-if as a special form
(my-if (special-form-p 'my-if)
       "special"
       "not special")

special

And it is (or was?) a feature that actually saw some usages:

;; org-compat.el:
(define-obsolete-function-alias 'org-without-partial-completion 'progn "9.6")

;; byte-run.el:
;; This was undocumented and unused for decades.
(defalias 'inline 'progn
  "Like `progn', but when compiled inline top-level function calls in body.
You don't need this.  (See bytecomp.el commentary for more details.)

\(fn BODY...)")

The solution is quite simple: since we already have runtime macro-expansion, it is quite straightforward to just also put our special form logic there.

Now we have talked about macros and special forms. Let’s go on to the last list form: function calls.

In most Lisps, there will be some “primitive” functions that have to be provided by the runtime, like arithmetic operations ((+ 1 x) and (- 1 y)) or list/cons-related functions ((car cons) and (cdr cons)). Semantically, +, -, car, cdr are all functions: one can use funcall to indirect-call them, or pass them to mapcar as a mapping function. However, we really want to some special processing for them: function calls are too costly for common constructs like these:

Instead, we want to inline the addition instructions into the AST graph so that Truffle can specialize each node, hopefully turning each into a simple integer addition operation:

For subroutines/user-defined functions that are not intended to be inlined, you can still provide basic optimizations for them by using direct method dispatch (and caching) as much as possible, as is covered by some other tutorials out there (Writing a Language in Truffle. Part 3: Making my Language (Much) Faster and Graal Truffle tutorial part 9 – performance benchmarking).

Internally, GNU Emacs Lisp keeps local variables in a list (Vinternal_interpreter_environment) and it implements variable capturing by keeping the very list in the closure object.

Truffle languages, as is shown above, usually use VirtualFrame objects for local variables, and they have very good reasons to do so: Truffle is aware of VirtualFrame and can inline the costs away. (Still remember the BytecodeInterpreter example above?)

To add to this, Emacs makes quite heavy use of the environment list in its oclosures (closures with associated metadata) as of Emacs 30:

;; an advice oclosure object (with up to 6 slots)
#[
  ;; [0]: arguments
  (&rest r)
  ;; [1]: lisp code
  (
   (if t nil car cdr how props) ; prevent environment trimming by cconv
   (apply car cdr r)
   )
  ;; [2]: environment
  ((car . test-advice) (cdr . test-main) (how . :around) (props))
  ;; [3]: nil
  nil
  ;; [4]: docstring
  advice
  ]

An advice oclosure object has four fields: car, cdr, how, props, and we can clearly see these fields are stored right in the environment slot of the closure.

I don’t have a solution for this. Currently, whenever oclosure requests it, my implementation produces a list of captured variables from the parent scope (stored as a MaterializedFrame), but it never propagates the changes backwards. Let’s hope this works. Otherwise, I don’t know, we might have to ditch MaterializedFrame/VirtualFrame altogether.

Multiple closures created in a function might only partially shares the same MaterializedFrame. Consider the following program:

;; -*- lexical-binding: t -*-
(let (closures)
  (dolist (i '(1 2 3))
    (push (lambda () i)
          closures))
  (mapcar #'funcall closures))

The three closures created in the loop return 3, 2, 1 respectively, meaning that each closure captures a different i. However, for variable outside the loop:

;; -*- lexical-binding: t -*-
(let (j closures)
  (dolist (temp '(1 2 3))
    (setq j temp)
    (push (lambda () j)
          closures))
  (mapcar #'funcall closures))

Now they all return 3, with j shared between all three closures.

Since we already split a stack frame into a series of chained frames (3.1), we don’t need to do anything special: (j . 1) and (i . X) is stored on different Truffle frames, and we allocate a new frame every iteration. By doing so, the closures naturally get to share j while keeping their own i:

For languages with for loops with per-iteration scopes, this can get more complex though:

const functions = [];
for (let i = 0; i < 5; i++) {
  functions.push(() => i);
}
return functions.map(f => f());

See how each JS closure above gets a different i while the let i = 0 declaration is technically outside the loop body? (GraalJs handles this by copying frames every iteration (when needed).)

GraalJS does a bit more to prune the captured frames. Since Truffle frames hold all local variables, it can prevent unused values from garbage collection when there are closures keeping references to the frame chain. And GraalJs does static analysis to decide what frame slots are safe to clear, keeping only the ones used by closures.

Actually, Emacs offers this too, after bootstrapping, Emacs calls cconv-make-interpreted-closure when creating closures and it then prunes unused conses off the environment list. But, as always, it assumes the captured environment to be a list, which I am not sure how to deal with.

No! That’s enough rambling for today. The above has covered most things I’ve tried and done in the little AST interpreter and should get you an interpreter that can potentially run on par with Java. However, since Emacs nowadays also runs bytecode and native-comp products (compiled with libgccjit), there will probably be another post dedicated to writing bytecode interpreters (specifically, a hand-rolled Emacs Lisp bytecode interpreter and a nativecomp LIMPLE interpreter (to be) implemented with Truffle bytecode DSL).

By the way, you can find the interpreter here: https://codeberg.org/gudzpoz/Juicemacs/ (it requires Java 23 and probably some patience to getting it running though ;-)