173. Lowering oxc’s AST to VIR

All the children

In chapter 93, lower.ts turned anything it didn’t model into an opaque instruction taking every variable the subtree mentions. It found them with one small function from estree.ts (simplified):

function children(n: Node): Node[] {
  const out: Node[] = [];
  for (const key of Object.keys(n)) {
    const v = n[key];
    if (v && typeof v.type === 'string')
      out.push(v);
    // (and arrays of nodes)
  }
  return out.sort(byStart);
}

An ESTree node is just an object, so you can ask it for its keys. oxc’s AST (chapter 172) is a set of Rust types: enums with dozens of variants and a struct per node kind. A struct has no Object.keys. Yet the port walks “all the children” in several places: the resolver, hoisting, the mentioned collector. oxc solves this once for everyone.

One method per node kind

Here is the idea in a toy, shaped like oxc’s Visit trait from the oxc_ast_visit crate:

trait Visit: Sized {
    fn visit_expr(&mut self, e: &Expr) {
        walk_expr(self, e);
    }
    fn visit_name(&mut self, _: &str) {}
    fn visit_func(&mut self, b: &Expr) {
        self.visit_expr(b);
    }
}

Every method has a default body. walk_expr is the standard walk-through: a match that calls visit_* on each child, left to right. (: Sized lets default methods pass self to a generic function like walk_expr; oxc’s trait says the same.)

The real trait has more than 250 visit methods, all generated, and each default calls the matching walk_* function, which visits that node’s children^[1]. The port’s README notes that this walk goes in source order, exactly what children() gave with its sort. That matters: the order in which variables are found decides an opaque instruction’s arguments, and so the value numbers.

impl Visit for Mentions {
    fn visit_name(&mut self, n: &str) {
        self.found.push((n.into(), self.inner));
    }
    fn visit_arrow(&mut self, b: &Expr) {
        let saved = self.inner;
        self.inner = true;
        self.visit_expr(b);
        self.inner = saved;
    }
}

Here is the real Mentions from lower.rs, trimmed (walk_arrow_fn stands for oxc’s much longer walk::walk_arrow_function_expression):

impl<'a> Visit<'a> for Mentions {
    fn visit_identifier_reference(
        &mut self,
        it: &IdentifierReference<'a>,
    ) {
        let at = (it.span, self.inner);
        self.found.push(at);
    }
    fn visit_arrow_function_expression(
        &mut self,
        it: &ArrowFunctionExpression<'a>,
    ) {
        let saved = std::mem::replace(
            &mut self.inner,
            true,
        );
        walk_arrow_fn(self, it);
        self.inner = saved;
    }
}

It overrides visit_binding_identifier and visit_function the same way: four methods cover all of JavaScript. std::mem::replace stores true and returns the old value.

enum Expr {
    Num(f64),
    Name(String),
    Add(Box<Expr>, Box<Expr>),
    Func(Box<Expr>),
}

trait Visit: Sized {
    fn visit_expr(&mut self, e: &Expr) {
        walk_expr(self, e);
    }
    fn visit_name(&mut self, _: &str) {}
    fn visit_func(&mut self, b: &Expr) {
        self.visit_expr(b);
    }
}

fn walk_expr<V: Visit>(
    v: &mut V,
    e: &Expr,
) {
    use Expr::*;
    match e {
        Num(_) => {}
        Name(n) => v.visit_name(n),
        Add(a, b) => {
            v.visit_expr(a);
            v.visit_expr(b);
        }
        Func(b) => v.visit_func(b),
    }
}

// Overrides one method.
struct Names(Vec<String>);

impl Visit for Names {
    fn visit_name(&mut self, n: &str) {
        self.0.push(n.into());
    }
}

// Also overrides visit_func, and does
// not walk the body.
struct Outer(Vec<String>);

impl Visit for Outer {
    fn visit_name(&mut self, n: &str) {
        self.0.push(n.into());
    }
    fn visit_func(&mut self, _: &Expr) {
    }
}

fn name(s: &str) -> Box<Expr> {
    Box::new(Expr::Name(s.into()))
}

fn main() {
    // a + (() => b + 1)
    let body = Expr::Add(
        name("b"),
        Box::new(Expr::Num(1.0)),
    );
    let e = Expr::Add(
        name("a"),
        Box::new(Expr::Func(
            Box::new(body),
        )),
    );
    let mut all = Names(vec![]);
    all.visit_expr(&e);
    println!("all:   {:?}", all.0);
    let mut out = Outer(vec![]);
    out.visit_expr(&e);
    println!("outer: {:?}", out.0);
}

Two passes, two files each

As in TypeScript, the plugin makes two passes. resolve.rs pairs every identifier with its binding (chapter 42), then lower.rs builds VIR. The resolver implements Visit and overrides the node kinds that resolve.ts handles in its own switch: functions, blocks, loops, switch, try (it counts its depth, for assigned_in_try), catch, classes, declarations, identifiers, assignments, imports and exports.

The bigger change is how things are named. resolve.ts keys its maps by node objects: refs: Map<Node, JsBinding>, and a binding’s owner is the function node itself. JavaScript compares objects by identity, so that is free. The port keys everything by plain values instead. An owner is FnKey::Program or FnKey::Func(Span) (a function’s span is unique), refs is a HashMap<Span, JsBindingId>, and bindings live in a Vec, named by a typed index (chapter 170). No table holds a reference into the AST, so none needs the AST’s lifetime.

Lowering, function by function

lower.rs uses the same rule as lower.ts: a variable stays an SSA value only if it is a simple kind, used in no other function, never assigned in a try, and not initialized with an array.

In TypeScript, lowerFunction holds its helpers (read, write, expr…) as closures sharing the function builder fb. Rust closures can’t share mutable state that freely, so the port has one Lowerer struct, with the per-function state in one field, f. A nested function swaps in its own with std::mem::replace(&mut self.f, state), lowers its body, and puts the outer state back.

Real output, printed by a ten-line program of mine (not in the repo) that calls JsPlugin.lower and print_func:

let n = 0;
const ps = [];
function add(v) {
  ps.push({ x: v });
  n++;
}
add(2);
console.log(this, n);

func add(%0):
  b0:
    %0 = param
    %1 = load ps
    %2 = record {x} %0
    %3 = push %1, %2
    %4 = load n
    %5 = unop + %4
    %6 = const 1
    %7 = binop + %5, %6
    store n, %7
    return

v stays an SSA value (%0). ps is array-initialized and n is used inside add, so both are bindings. In <module>, this becomes %20 = opaque(ThisExpression): ast.rs maps oxc’s variant back to the ESTree type name, so explain reads as before.

No AST after lowering

The TS engine keeps ESTree nodes in SourceRef.node, and the emitter asks them questions, such as (through a parent map) “is this a[i] the object of a member expression?”. oxc allocates its AST in the parse’s arena, so keeping it would put a lifetime on the whole Module. Instead node is a NodeRef, an index into a side table of facts (trimmed):

pub enum JsNode {
    Plain,
    Member {
        object: Span,
        property: Span,
        simple_property: bool,
        // ...
        direct: bool,
    },
    // Write, Object, Call, ForOf ...
}

direct answers the parent question: lowering sets a flag just before it lowers a member expression’s object. Chapter 174 uses these facts.

Your turn

enum Expr {
    Num(f64),
    Name(String),
    Add(Box<Expr>, Box<Expr>),
    Call(Box<Expr>, Vec<Expr>),
    Func(Box<Expr>),
}

trait Visit: Sized {
    fn visit_expr(&mut self, e: &Expr) {
        walk_expr(self, e);
    }
    fn visit_name(&mut self, _: &str) {}
    fn visit_func(&mut self, b: &Expr) {
        self.visit_expr(b);
    }
}

fn walk_expr<V: Visit>(
    v: &mut V,
    e: &Expr,
) {
    use Expr::*;
    match e {
        Num(_) => {}
        Name(n) => v.visit_name(n),
        Add(a, b) => {
            v.visit_expr(a);
            v.visit_expr(b);
        }
        Call(f, args) => {
            v.visit_expr(f);
            for a in args {
                v.visit_expr(a);
            }
        }
        Func(b) => v.visit_func(b),
    }
}

// Each name, and whether it is inside
// a nested function.
#[derive(Default)]
struct Mentions {
    found: Vec<(String, bool)>,
    inner: bool,
}

impl Visit for Mentions {
    fn visit_name(&mut self, n: &str) {
        todo!()
    }
    fn visit_func(&mut self, b: &Expr) {
        todo!()
    }
}

What you learned

• The visitor pattern: a trait with one method per node kind, whose defaults walk the children. You override only what you need. Overriding replaces the walk, so call it yourself if you want the children.
• oxc’s Visit walks in source order, like children(), so value numbers match the reference.
• resolve.rs and lower.rs port their TS files case by case. Node identity becomes spans and typed ids (FnKey, HashMap<Span, _>). Closures sharing fb become one Lowerer with a swappable FnState.
• Instructions point into a side table of JsNode facts, not into a live AST.