path: root/notes/260625-optimize.md

# AST Optimizations

<!--TOC-->

I might be deluded, but I think it should be possible to make a tree
walking AST interpreter pretty much as efficient as a bytecode VM, if
we're smart enough about the AST representation and a well-selected
set of optimizations performed on it.

Here's what I came up with so far.

## Abandon linked lists (cons cells / pairs)

Using traditional cons cells to represent code forms as linked lists
means you're paying a memory overhead of 50% to store lists.  That's
ridiculous.  Just use arrays.

See my latest [cons cell optimization](260611-fastcons4.html) article
on how I've decided to implement this.  Briefly: Most lists in code
are under 8 elements anyway, so we use the low 3 bits of a pointer to
encode whether the length is 1 to 7, or greater; only if it's greater
does the array being pointed to require a termination marker.

(We could put the length in a 64-bit head if longer than 7, but the
AST walker will be going through the list one way or another, so it
doesn't really matter; having the elements always start at zero is
simpler.)

## Tight packing of short arrays

To further ensure minimal memory footprint of AST nodes, we make sure
that they're allocated tightly in blocks, without any padding between
them.  See `src/zisp/gc/ListPool.zig` for how I've implemented this.
Cache locality and memory density should now be perfect.

## Help the prefetcher (ABANDONED)

Memory prefetching refers to a modern CPU feature in which the CPU
notices that you're accessing memory addresses sequentially, and
decides to prefetch memory beyond what you've requested.

It would be neat if we could order the memory of our AST nodes such
that they occur in the order the interpreter will walk them, but it's
difficult to achieve without significant complications that make the
parser slower in general, having to constantly copy around arrays to
put them in the right order.

I've decided to abandon this idea, because most code jumps around a
lot anyway.  By the time you've accessed 3-4 sequential cache lines,
and the prefetcher started to warm up to the pattern, you suddenly
jump to a completely different point in the AST because a function
call that's not inlined is encountered.

So yeah, scrap that.  Hopefully it makes no difference.

## Quoted data directly in the AST

This may just be a stupid micro-optimization, because data that could
be confused for code to evaluate (lists and identifiers) but is meant
as a constant value can simply be wrapped in `(#QUOTE x)` which is a
pointer to a two-value array which should be very quick to evaluate,
but I had a lot of spare value ranges in my NaN packing strategy so
I've decided to dedicate some of them to "implicitly" quoted values.

A two-value array is a two-value array.  There will be fewer of them
this way, and you'll be saving two 64-bit values per quoted object.

This isn't actually that bad considering Zisp conflates strings and
symbols.  Code with lots of string literals would otherwise end up
with lots of quote forms (automatically emitted by the parser upon
encountering a double-quoted string) if not for this optimization.

## Constant propagation

Standard stuff: If a variable is an immutable reference to a value
known statically, just inline the value when the variable is used.

Example Scheme code:

    ;; before optimization
    (define x 5)
    (define (foo) (+ x 1))

    ;; after
    (define (foo) (+ 5 1))

    ;; in fact, probably this
    (define (foo) 6)

Any operator in the standard library (whose bindings are immutable)
may be marked as "pure" which makes it eligible for evaluation at
optimization time to fold constants.

The same could be done with non-exported functions in user code, so
the following works as well...

    ;; before
    (define x 5)
    (let (foo) (+ x 1))
    (define (bar) (foo))

    ;; after
    (define (bar) 6)

Note: `(let (foo) ...)` is a local function definition in Zisp.  Not
married to this syntax, but I think I like it.

Actually, we may decide *not* to do this optimization for exported
constants, so they can be redefined.  Let's say `define` creates
public bindings and `let` local ones:

    ;; before

    (define x 5)
    (let y 10)
    (let (z) (* y 2))

    (define (foo) (+ x 1))
    (define (bar) (+ y 1))
    (define (qux) (+ (z) 1))

    ;; after

    (define x 5)

    (define (foo) (+ x 1))
    (define (bar) 11)
    (define (qux) 21)

I think I like that.  There could still be a way for other modules to
"statically link" to a module, which would effectively turn public
bindings of that module into local bindings within the importing
module, thus making the constants eligible for folding anyway!

## Lexical closure references by flat index

Now we're getting to the hot stuff.

Zisp won't allow capturing mutable variables by reference in closures;
you'll need to copy the variable's current value into a non-mutable
variable if you want to capture it.  (There's precedent for this in
other languages, like Java; nothing newly invented here.)

You can still simply capture a reference to a mutable heap object,
like a one-element array aka "box" object type, so you're not really
losing any power.  IMO, this actually makes code cleaner anyway, as
capturing mutable state in a closure becomes clearly demarcated via
explicit heap allocation.

Given that captures are immutable, they can be passed by value rather
than by reference.  This is conceptually implemented by the following
algorithm, though we'll optimize it greatly:

1. Whenever a `lambda` is evaluated, walk through it to locate the
   lexically captured variables, giving each a number.

2. Put their values in a `captures` array and associate it with the
   lambda object.

3. Replace the lexical variable references in the lambda body with
   indexing operations into said array.

Although well-written code probably shouldn't be calling `lambda` all
the time, it would still be bad if we need to walk the entire body
every time we encounter one.

The walk only needs to happen once; we can then create a "template"
array that contains references to the *current* scope, and the next
time the lambda is encountered, we copy the array and evaluate the
contained local references, then attach that new array to the new
closure object.

Here's some Scheme-like pseudo-code to demonstrate:

    ;; Raw source

    (define (higher-order x y)
      (lambda (z)
        (+ x y z)))

    (define foo (higher-order 1 2))

    ;; Optimized; using #<foo> to represent "internal" stuff

    (define (higher-order x y)
      (#<make_lambda> #<copy_lexicals [x y]> (z)
        (+ #<lexical:0> #<lexical:1> z)))

    (define foo (higher-order 1 2))
    ;foo = #<lambda lexicals:[1 2] (z)
              (+ #<lexical:0> #<lexical:1> z)>

This "fattens" closures a bit but I think it's worth it.  I'm pretty
sure that various Scheme implementations do this when they can prove
that a captured variable is immutable.  Another cache locality win,
and no 2D environment traversal at execution time.

If you look carefully, you'll notice that we inserted a node into the
AST: `#<copy_lexicals [x y]>`.  This is to be avoided, as I want the
AST optimizations to happen purely in-place to preserve the tightly
packed nature of the list-arrays making up the source code, and to
avoid unnecessary allocations and code bloat.

Fortunately, `#<make_lambda>` could actually be represented with just
a few bits...  In the NaN packing strategy, there is a special type of
pointer which basically means "pointer to optimized code expression."
The first element of the destination, then, doesn't need to be treated
as an arbitrary NaN-packed value, but can instead be a sort of tightly
packed instruction with a payload.  Say, lower 8 bits for the opcode,
and higher 56 bits as payload.  Now `#<make_lambda>` can itself have
the pointer to `[x y]` in its payload, implicitly meaning those are
the lexicals to be copied.

This still means there's an extra allocation happening which isn't
*strictly* necessary as we could walk the body every single time the
`lambda` is executed (i.e., at closure creation time), but let's not
be unreasonable.  So long as I keep the original shape of the AST,
only needing small auxiliary allocations, I consider it a success.

To summarize: References to captured lexicals are now literally just
`lexicals[i]` where `lexicals` could, I suppose, be a thread-local
global variable, given that only one function is active at a time, or
it could be an implicitly passed argument; we'll see.

Meanwhile, creating a lambda is still reasonably cheap, requiring only
the copying of a (usually very small) array and iterating through it
to replace "local variable reference" objects with their values.

Those "local variable reference" objects can also just be flat array
indexes, thanks to other optimizations...

## Eager stack frame allocation

This may just be a micro-optimization; I've not thought about it that
much yet.  But basically, the idea is to walk the body of a lambda,
and note how many local variables it uses, so you can allocate its
entire stack frame at once, then turn local variable references into
flat array indexes.

We can combine this with the previous strategy: The same walk through
the body can count local declarations, numbering them and turning them
into literally just `locals[i]` (regardless of whether it's assignment
or reference).  The count can be stored alongside the array of lexical
captures, changing our closure shape on the heap to:

    #<lambda meta:{ locals_count:N lexicals:[...] } (args) (body)>

The `locals` array could be a thread local or implicit argument like
the `lexicals` array, or it could actually just be based on the stack
pointer.  (Referring to the Zisp stack, not the native stack.)

Oh, and function parameters are of course just the first few entries
in this (conceptual) locals array (which may just be a stack pointer)
and get populated at the call-site with the passed arguments.

## Lexical and local reference representation

Above, I've said that lexical and local variable references literally
become `lexicals[i]` and `locals[i]`, but that was a bit of a lie.
This is an interpreter after all.

What I meant is:

The AST will contain nodes that are NaN-packed values (like almost
everything in the AST) with a tag saying e.g. "lexical reference,"
holding a payload value of up to 48 bits, which is the index.

So, the interpreter's code would contain something like this:

    const tag: u16 = getTag(nan_value);
    const payload: u48 = getPayload(nan_value);
    return switch (tag) {
      ...
      LOCAL_REF => locals[payload],
      LEXICAL_REF => lexicals[payload],
      ...
    };

Realistically, you don't need 48 bits for a local or lexical variable
index.  Even 16 bits would be more than enough.  So, there's a bit of
waste here compared to what we could do with bytecode.

## VM opcode in expression head

I've already explained this above as part of the lexical references
optimization:

Once a list (array) has been analyzed and deemed to be a valid form,
the pointer to it can have its tag switched to mean "optimized code
expression pointer" and then the first element doesn't need to be
treated as an arbitrary NaN-packed value anymore, and can instead
denote various operations (8 bits) with a payload (56 bits).

Now in our AST we can essentially have VM opcodes standing for things
like IF, WHILE, LAMBDA (as explained above), CALL, and so on, with a
payload that can fit a 48-bit pointer with 8 bits to spare.

Given that various operations like IF don't actually need a payload,
this is likely to end up being one of the remaining sources of memory
waste in our AST, along with the generally pointer-rich structure as
it is still a tree with a uniform 64-bit NaN-packed value/node format.

## Wrapping up

I think this is all the ideas I had floating in my mind so far; might
add more later if I can think of any.

There's no question that bytecode would be more compact than this AST
despite all the optimizations.  Wasting 56 bits on an unused opcode
payload here, 32 bits on an overly large locals index there, and the
generally pointer-rich structure of the tree...  These will add up,
and bytecode could fix all that.  But the beauty of this strategy is
that you literally just work directly on the data structure returned
by the parser, doing only in-place mutations and a small number of
auxiliary allocations.  Even after optimizations, the AST should be
almost trivial to serialize in a pretty-printed format to get either
some nice insights into what the optimizer did, or, if allocation of
debug metadata was enabled, it may even be possible to turn it back
into original source form.

So, the low latency and the REPL/debug experience will be impeccable.
Not to mention the incredible simplicity of it all: We don't need a
complex bytecode design, a bytecode VM, a bytecode debugger, etc.;
only a few "VM opcodes" in the first position of some arrays which
simply represent AST nodes.

And if you need peak performance, there's always native compilation.

In my opinion, JIT is unnecessary, because one can just use GCC or
LLVM for AOT-quality codegen at runtime these days.  (They call it
JIT, like in libgccjit, but it's really AOT as a library.)  Either
way, you can skip the bytecode.  There's actually a little research
indicating that JIT from AST is as good as, if not better than, JIT
from bytecode:

[AST vs. Bytecode: Interpreters in the Age of Meta-Compilation](https://stefan-marr.de/downloads/oopsla23-larose-et-al-ast-vs-bytecode-interpreters-in-the-age-of-meta-compilation.pdf)

That's using RPython and GraalVM, but I think the same should apply
when using LLVM or GCC.  Any optimizations specific to the Lispish
nature of a language should be possible to perform on the original
AST, or something close to it; after that, send it to LLVM/GCC and
have them perform the brunt of complex, slow, AOT-tier optimizations.

Transforming to bytecode first may even act as obfuscation and end up
worsening the work of GCC/LLVM on the code.  Don't quote me on that,
though.

In Zisp, I want to support what I've dubbed "manual JIT" meaning the
programmer is in control over what should be native-compiled and when.

A performance-critical application could be shipped as source code,
but compile itself at startup, for instance.  I've written more about
this idea here:

[The interpreter and the compiler](260522-interpreter.html)

And when you want the REPL experience, or quick source-only scripts
for which you don't even want to deal with annoying on-disk bytecode
cache files, then the in-place optimizing AST interpreter should do a
fine job.