Construct AST nodes in place

A proposal for how the parser builds AST nodes. This sits underneath AstBuilder rather than replacing it - AstBuilder::new/boxed would continue working, and would be reimplemented on top of this.

The below was written by Claude based on a lengthy conversation I had with him, and prototyping. I asked Claude to write it in my voice, which he's only been partially successful in, so it reads a bit odd in places!

The problem

The most common shape in the parser is: parse the children, then allocate the parent around it.

let rhs = self.parse_binary_expression_or_higher(left_precedence);
let span = self.end_span(lhs_span);
Expression::new_binary_expression(span, lhs, op, rhs, self)

The parent's memory is allocated last, so everything going into it has to be held while the child is parsed. lhs and op are known before we call parse_binary_expression_or_higher, but there's nowhere to put them yet, so the compiler keeps them in registers across the call - and lhs is an Expression, which is two registers on its own. Registers that survive a call have to be saved and restored, so we pay:

• Extra push/pop at function entry/exit.
• Register shuffling around the call.
• Parent's fields all written at the end, rather than as they become available.

The overhead is surprisingly large, especially with types with many fields, where a lot of values get shifted to and from the stack before finally being copied into the arena.

This isn't the allocator being slow. It's that we're forced to keep things alive that we could have finished with.

The fix

Allocate the parent's memory first, then write each field as soon as we have it.

A performant API to do this without any unsafety:

// Produces an `ArenaBox<'a, BinaryExpression>`
let bin_expr = BinaryExpression::build(self)
    .span_start(lhs_span)   // written immediately - `lhs_span` is now dead
    .left(lhs)              // written immediately - `lhs` is now dead
    .operator(op)           // written immediately - `op` is now dead
    .right(self.parse_binary_expression_or_higher(left_precedence))
    .span_end(self.prev_token_end)
    .finish();

By the time we call parse_binary_expression_or_higher, everything we already know is in its final home in the arena, and nothing needs to survive the call except the pointer to the node. Fields can be written in whatever order the values arrive.

span is split into span_start and span_end. We know the start before parsing the children and the end only afterwards, so a single .span() would force the start to be held in a register across the call - exactly what we're trying to avoid. Splitting it saved 2 instructions and one callee-saved register in my test, and costs nothing, because end_span builds the Span from two separate registers anyway, so it's two stores either way.

The builder compiles away entirely - it's a pointer, and every method is #[inline(always)].

Why alloc_with doesn't do it

The simpler alternative using existing Allocator APIs would be to use a closure:

self.allocator.alloc_with(|| BinaryExpression {
    left: lhs,
    operator: op,
    right: self.parse_binary_expression_or_higher(left_precedence),
    span: self.end_span(lhs_span),
})

This allocates first, but that isn't enough. A closure produces a whole value, so the struct isn't assembled until the closure body finishes - which is after the child has been parsed. lhs and op still have to survive the parse_binary_expression_or_higher call. In my test it came out between the two: 30 instructions, versus 33 for what we do today and 25 for the builder.

The compiler won't reorder the writes for us either. I tried it with a child that provably can't touch memory, so reordering was obviously safe, and the output was byte-for-byte identical. I don't want to claim anything general about LLVM from one experiment, but no variant I tried moved a write earlier to free up a register. In practice, the order we write fields in the source is the order they happen - so no closure-based API can get there. The writes have to be separate calls.

Design

The builder holds a pointer to uninitialised memory, plus one type parameter per field recording whether that field has been set. finish() is only callable once they're all Yes.

This code would all be codegenned by ast_tools:

/// Shared by every builder.
#[doc(hidden)]
pub struct Yes;
#[doc(hidden)]
pub struct No;

/// Sealed trait that makes it impossible to implement e.g. `impl LeftIsSet for No {}` outside this module.
trait IsYes {}
impl IsYes for Yes {}

/// One of these per field name, so a missing field can name itself in the error.
#[diagnostic::on_unimplemented(message = "`.left()` has not been called on this builder")]
#[doc(hidden)]
pub trait LeftIsSet: IsYes {}
impl LeftIsSet for Yes {}
// ... one trait per field name ...

/// Builder type for `BinaryExpression`.
/// 5 type params `F1` - `F5` represent state of each field.
/// `span` gets 2, because its halves are written at different times.
#[repr(transparent)]
pub struct BinaryExpressionBuilder<'a, F1, F2, F3, F4, F5> {
    node: ArenaBox<'a, MaybeUninit<BinaryExpression<'a>>>,
    marker: PhantomData<(F1, F2, F3, F4, F5)>,
}

impl<'a> BinaryExpression<'a> {
    #[inline(always)]
    pub fn build(builder: &impl GetAstBuilder<'a>) -> BinaryExpressionBuilder<'a, No, No, No, No, No> {
        let builder = builder.builder();
        let mut node = ArenaBox::<BinaryExpression>::new_uninit_in(builder.allocator());
        // Set `node_id` automatically
        unsafe { (&raw mut (*node.as_mut_ptr()).node_id).write(Cell::new(builder.node_id())) };
        BinaryExpressionBuilder { node, marker: PhantomData }
    }
}

impl<'a, F1, F2, F3, F4, F5> BinaryExpressionBuilder<'a, F1, F2, F3, F4, F5> {
    #[inline(always)]
    pub fn left(mut self, left: Expression<'a>) -> BinaryExpressionBuilder<'a, F1, F2, Yes, F4, F5> {
        unsafe { (&raw mut (*self.node.as_mut_ptr()).left).write(left) };
        BinaryExpressionBuilder { node: self.node, marker: PhantomData }
    }

    // ... one method per field ...

    // `span` gets 3 methods - both halves at once, or either half on its own
    #[inline(always)]
    pub fn span_start(mut self, start: u32) -> BinaryExpressionBuilder<'a, Yes, F2, F3, F4, F5> {
        unsafe { (&raw mut (*self.node.as_mut_ptr()).span.start).write(start) };
        BinaryExpressionBuilder { node: self.node, marker: PhantomData }
    }

    #[inline(always)]
    pub fn span_end(mut self, end: u32) -> BinaryExpressionBuilder<'a, F1, Yes, F3, F4, F5> {
        unsafe { (&raw mut (*self.node.as_mut_ptr()).span.end).write(end) };
        BinaryExpressionBuilder { node: self.node, marker: PhantomData }
    }

    #[inline(always)]
    pub fn span(mut self, span: Span) -> BinaryExpressionBuilder<'a, Yes, Yes, F3, F4, F5> {
        unsafe { (&raw mut (*self.node.as_mut_ptr()).span).write(span) };
        BinaryExpressionBuilder { node: self.node, marker: PhantomData }
    }

    #[inline(always)] // No-op at runtime
    pub fn finish(self) -> ArenaBox<'a, BinaryExpression<'a>>
    where
        F1: SpanStartIsSet,
        F2: SpanEndIsSet,
        F3: LeftIsSet,
        F4: OperatorIsSet,
        F5: RightIsSet,
    {
        unsafe { ArenaBox::assume_init(self.node) }
    }
}

Two additions to ArenaBox that this needs:

impl<'alloc, T> ArenaBox<'alloc, T> {
    #[inline(always)]
    pub fn new_uninit_in<A: GetAllocator>(allocator: &A) -> ArenaBox<'alloc, MaybeUninit<T>> {
        // This IS needed here. `ArenaBox::from_non_null` also contains this assertion,
        // but it's passed a `MaybeUninit<T>` pointer, and `MaybeUninit` is always non-drop, so it always passes.
        // This assertion checks that `T` itself is non-drop.
        const { Self::ASSERT_T_IS_NOT_DROP };

        let allocator = allocator.allocator();
        let ptr = allocator.alloc_layout(Layout::new::<T>()).cast::<MaybeUninit<T>>();
        unsafe { ArenaBox::from_non_null(ptr) }
    }

    /// # SAFETY
    /// The `T` must be fully initialised.
    #[inline(always)]
    pub unsafe fn assume_init(boxed: ArenaBox<'alloc, MaybeUninit<T>>) -> ArenaBox<'alloc, T> {
        unsafe { ArenaBox::from_non_null(ArenaBox::into_non_null(self).cast::<T>()) }
    }
}

Notes

• Fields can be set in any order.
• span is special-cased into span_start / span_end / span, and takes two state parameters rather than one. .span() sets both at once for the cases where the whole span is known together. Every other field gets one method and one parameter.
• All the unsafe is in generated code, call sites are safe. You can't call finish() early and you can't forget a field.
• If parsing fails part way, the half-built node is abandoned. AST types aren't Drop, so that's some wasted arena bytes and nothing else.
• Defaulted fields don't appear at call sites. BinaryExpression::new fills node_id from the builder; Function::new also defaults scope_id, pure and pife. build writes those immediately and starts their state at Yes, so call sites set exactly the fields they pass to new today.
• Most results get wrapped straight after (Expression::BinaryExpression(node)), which is unchanged - finish() returns the same ArenaBox that boxed does today.
• Marker types rather than const bool parameters, because only types can carry a custom diagnostic. Forgetting a field gives error: `.operator()` has not been called on this builder.
• finish is in the same impl block as the setters, with where bounds, rather than a separate impl BinaryExpressionBuilder<'a, Yes, Yes, Yes, Yes, Yes>. That's what makes the per-field diagnostic possible.
• PhantomData is needed because the parameters are otherwise unused. It's zero-sized, costs nothing.
• The not-Drop assertion needs care in both new methods. ArenaBox's invariant is enforced by every constructor referencing ASSERT_T_IS_NOT_DROP, but on alloc_uninit the natural form is vacuous - needs_drop::<MaybeUninit<T>>() is always false, so it would accept a T: Drop. Hence asserting on T explicitly. Equally, assume_init has to go via from_non_null (which asserts) rather than constructing the struct directly, or alloc_uninit::<T>() + assume_init() becomes a silent route to putting a Drop type in the arena.

Out-parameters for large inline fields

Most parse_* methods return something small: Expression, Statement, TSType and BindingPattern are all 16 bytes, returned in two registers. Cheap, and not worth changing.

Some AST nodes are stored inline in their parent rather than boxed, and get returned by value:

• Function is 96 bytes and holds id: Option<BindingIdentifier<'a>> inline
• BindingIdentifier is 32 bytes, and parse_binding_identifier() returns it by value

On x86-64 SysV, 32 bytes is too big to return in registers, so it comes back via a hidden pointer to a temporary on stack, which is then copied into the Function in the arena. The builder can hand the callee an out-parameter (a type wrapping &mut MaybeUninit<BindingIdentifier>) pointing straight at the field, so it writes there once and the copy disappears:

func_builder.id_with(|slot| self.parse_binding_identifier_into(slot))

parse_binding_identifier_into would receive the slot type and convert it to a ZST "Done" token. To ensure the "Done" token represents "this slot is filled" not "some slot is filled", the slot and token types would hold an unforgeable branded lifetime that ties them together.

In my test this saved 2 instructions in the callee and 2 in the caller, and removed a 48-byte stack frame plus a 40-byte copy.

Only do this where the value is too large to return in registers. For a 16-byte return it's a net loss - the callee has to keep the out-pointer alive across its own internal calls, costing it a register and a push/pop pair, while the caller saves one instruction. I measured small types as 2 instructions worse in the callee. Other candidates worth checking: FormalParameters (48 bytes), FormalParameter (72), VariableDeclarator (56), CallExpression (64), Class (144).

Nodes end up closer to traversal order

This may matter more than the instruction counts.

Today we allocate a parent after all its children, so nodes land in memory in post-order: a parent sits next to its last child, while its first child - the next thing any tree walk visits - is a whole subtree away. Allocating the parent first puts nodes in roughly pre-order, with a parent immediately next to its first child.

Roughly, not exactly: the parser backtracks and re-parses in places, so abandoned allocations get interleaved. But the bulk of the tree would be stored close to the order it's read in, and everything downstream - semantic, linter, transformer, codegen - walks top-down. That's a cache argument that applies to every consumer of the AST, not just to parsing.

This is reasoning, not a measurement, and it's the main thing I'd want benchmarked.

Should we also bump upwards?

The arena bumps downwards, so allocating parents first puts them at higher addresses than their children - pre-order, but descending. Bumping upwards would make it ascending, which suits hardware prefetching better.

I measured upward bumping as costing nothing once we're on Arena<8> (an arena with 8-byte minimum alignment, so alignment rounding disappears - blocked on moving strings out, since they're the only sub-8-byte allocations). But it needs care: one reasonable way to write the fast path ties with downward, another costs an instruction per allocation.

Worth a separate follow-up rather than bundling it in. Descending pre-order may well be fine.

Next steps

Proposed next steps:

1. Prototype builders for a handful of hot node types by hand, before touching codegen.
2. Benchmark the parser, and separately benchmark the linter and codegen over an AST built this way, to test the ordering claim.
3. Decide separately whether out-parameters for large inline fields are worth the extra API surface.

If it holds up, generating the builders is mechanical - we generate ast_builder.rs already. Migrating to the new APIs could be performed by a combination of ast-grep codemod + LLM (codemod written by LLM too), same as previous AstBuilder migration.

Start by reading the existing AstBuilder and ArenaBox APIs, then inspect ast_tools and the parser's parse_* entry points described in the proposal. Determine the scope and design before implementation; done would mean an agreed approach for in-place AST construction and the related large inline-field handling.