Capturing Tokens by Replaying Them: rustc's Lazy Token Capture

TL;DR: rustc’s parser does not collect tokens into a buffer as it parses. It takes a snapshot of the token cursor and remembers the bump count; the captured range is a half-open interval of bump counts. The tokens are materialised by replaying the cursor, and only if something actually asks for them — which is rare.

The problem

Some AST nodes need their original tokens: anything that will be handed to a proc macro, a derive, or an attribute that reads source text. Most nodes do not.

Collecting tokens eagerly — appending to a Vec as the parser advances — would cost every node in the program. rustc instead does the opposite.

The snapshot

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:17-24

// When collecting tokens, this fully captures the start point. Usually its
// just after outer attributes, but occasionally it's before.
#[derive(Clone, Debug)]
pub(super) struct CollectPos {
    start_token: (Token, Spacing),
    cursor_snapshot: TokenCursor,
    start_pos: u32,
}

and the snapshot operation:

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:94-101

impl<'a> Parser<'a> {
    pub(super) fn collect_pos(&self) -> CollectPos {
        CollectPos {
            start_token: (self.token, self.token_spacing),
            cursor_snapshot: self.token_cursor.clone(),
            start_pos: self.num_bump_calls,
        }
    }

start_pos is a bump counter, not a byte offset or a token index. The captured range is [start_pos, end_pos) in bump-counter space.

Note that TokenCursor::clone() is cheap-ish: the cursor holds a Vec of enclosing delimiters and Rc-shared token trees. The token trees are shared, not copied.

The end position, with the ungluing correction

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:278-294

assert!(
    !(self.break_last_token > 0 && matches!(capture_trailing, Trailing::Yes)),
    "Cannot have break_last_token > 0 and have trailing token"
);
assert!(self.break_last_token <= 2, "cannot break token more than twice");

let end_pos = self.num_bump_calls
    + capture_trailing as u32
    // If we "broke" the last token (e.g. breaking a `>>` token once into `>` + `>`, or
    // breaking a `>>=` token twice into `>` + `>` + `=`), then extend the range of
    // captured tokens to include it, because the parser was not actually bumped past it.
    // (Even if we broke twice, it was still just one token originally, hence the `1`.)
    // When the `LazyAttrTokenStream` gets converted into an `AttrTokenStream`, we will
    // rebreak that final token once or twice.
    + if self.break_last_token == 0 { 0 } else { 1 };

let num_calls = end_pos - collect_pos.start_pos;

The correction: if the last token was unglued (see the previous article), the range must include it, because the parser was never actually bumped past it. Even two breaks of >>= count as one original token.

The lazy stream

rust
// compiler/rustc_ast/src/tokenstream.rs:217-239

enum LazyAttrTokenStreamInner {
    // The token stream has already been produced.
    Direct(AttrTokenStream),

    // From a value of this type we can reconstruct the `TokenStream` seen by
    // the `f` callback passed to a call to `Parser::collect_tokens`, by
    // replaying the getting of the tokens. This saves us producing a
    // `TokenStream` if it is never needed, e.g. a captured `macro_rules!`
    // argument that is never passed to a proc macro. In practice, token stream
    // creation happens rarely compared to calls to `collect_tokens` (see some
    // statistics in #78736) so we are doing as little up-front work as
    // possible.
    //
    // This also makes `Parser` very cheap to clone, since there is no
    // intermediate collection buffer to clone.
    Pending {
        start_token: (Token, Spacing),
        cursor_snapshot: TokenCursor,
        num_calls: u32,
        break_last_token: u32,
        node_replacements: ThinVec<NodeReplacement>,
    },
}

The comment mentions the statistics: token stream creation happens rarely compared to calls to collect_tokens. That is the empirical justification for the whole design.

And it gives a second benefit: the Parser becomes cheap to clone, because there is no intermediate collection buffer. And the parser is cloned — for every macro rule attempt during macro_rules! matching.

The replay

rust
// compiler/rustc_ast/src/tokenstream.rs:252-263

// The token produced by the final call to `{,inlined_}next_and_bump` was not
// actually consumed by the callback. The combination of chaining the
// initial token and using `take` produces the desired result - we
// produce an empty `TokenStream` if no calls were made, and omit the
// final token otherwise.
let mut cursor_snapshot = cursor_snapshot.clone();
let tokens = iter::once(FlatToken::Token(*start_token))
    .chain(iter::repeat_with(|| FlatToken::Token(cursor_snapshot.next_and_bump())))
    .take(*num_calls as usize);

if node_replacements.is_empty() {
    make_attr_token_stream(tokens, *break_last_token)
} else {

The replay is a lazy iterator: once(start_token).chain(repeat_with(next)).take(n). The take is what stops it. And make_attr_token_stream(tokens, break_last_token) re-splits the final glued token using the same table the parser used — so the tokens handed to a proc macro are byte-identical to the source.

Deciding whether to capture: twice

The decision is made in two places, and the second one is sharper.

The early check, before parsing:

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:163-181

let possible_capture_mode = self.capture_cfg;

// We must collect if anything could observe the collected tokens, i.e.
// if any of the following conditions hold.
// - We are force collecting tokens (because force collection requires
// tokens by definition).
let needs_collection = matches!(force_collect, ForceCollect::Yes)
    // - Any of our outer attributes require tokens.
    || needs_tokens(&attrs.attrs)
    // - Our target supports custom inner attributes (custom
    //   inner attribute invocation might require token capturing).
    || R::SUPPORTS_CUSTOM_INNER_ATTRS
    // - We are in "possible capture mode" (which requires tokens if
    // the parsed node has `#[cfg]` or `#[cfg_attr]` attributes).
    || possible_capture_mode;
if !needs_collection {
    return Ok(f(self, attrs.attrs)?.0);
}

and the late check, after parsing, when inner attributes are known:

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:230-266

// When we're not in "definite capture mode", then skip collecting and
// return early if `ret` doesn't support tokens or already has some.
//
// Note that this check is independent of `force_collect`. There's no
// need to collect tokens when we don't support tokens or already have
// tokens.
let definite_capture_mode = self.capture_cfg
    && matches!(self.capture_state.capturing, Capturing::Yes)
    && has_cfg_or_cfg_attr(&ret_attrs);
if !definite_capture_mode && !ret_can_hold_tokens {
    return Ok(ret);
}

The comment in the source is candid about the residual waste:

If we return early here that means the setup, such as cloning the token cursor, was unnecessary. That’s hard to avoid.

What actually triggers capture

rust
// compiler/rustc_parse/src/parser/attr_wrapper.rs:410-427

fn needs_tokens(attrs: &[ast::Attribute]) -> bool {
    // Tokens are needed if...
    attrs.iter().any(|attr| match &attr.kind {
        AttrKind::Normal(normal) => {
            match normal.item.name() {
                // ... a multi-segment attribute is present, e.g. `rustfmt::skip`.
                None => true,
                // ... `cfg_attr` or a single-segment non-builtin attribute is present, e.g.
                // `derive`, `test`, `global_allocator`.
                Some(name) => name == sym::cfg_attr || !rustc_feature::is_builtin_attr_name(name),
            }
        }
        // Synthetic attributes are created only during expansion, and can't re-enter the parser
        // because they have no token form.
        AttrKind::Synthetic(_) => unreachable!(),
        AttrKind::DocComment(..) => false,
    })
}

Builtin attributes (#[doc], #[inline], #[allow]) do not need tokens. #[derive], #[test], #[global_allocator], #[cfg_attr] and any path-qualified attribute do.

The `IntervalSet` for seen attributes

rust
// compiler/rustc_parse/src/parser/mod.rs:274-283

// This state is used by `Parser::collect_tokens`.
#[derive(Clone, Debug)]
struct CaptureState {
    capturing: Capturing,
    parser_replacements: Vec<ParserReplacement>,
    inner_attr_parser_ranges: FxHashMap<AttrId, ParserRange>,
    // `IntervalSet` is good for perf because attrs are mostly added to this
    // set in contiguous ranges.
    seen_attrs: IntervalSet<AttrId>,
}

Attribute IDs are assigned monotonically, so the “already seen” set is mostly a few contiguous ranges — the same interval-set trick the borrow checker uses for regions, applied to a completely different problem.

Where it lives

compiler/rustc_parse/src/parser/attr_wrapper.rs — CollectPos, collect_pos, the two capture decisions.
compiler/rustc_ast/src/tokenstream.rs — LazyAttrTokenStream, the replay, TokenCursor.
compiler/rustc_parse/src/parser/mod.rs — CaptureState, seen_attrs.

Takeaway

“Capture the tokens” sounds like it needs a buffer. It does not — it needs a recipe for reconstructing them. Recording where you were and how far you went is enough, and the recipe is cheaper to store, cheaper to clone, and (crucially) only executed if someone asks. That is the difference between paying O(n) for every AST node and paying O(1).