From 287178ba0c63182bb6e302399084eeaad59cdcc8 Mon Sep 17 00:00:00 2001
From: xmakro <xmakro@users.noreply.github.com>
Date: Tue, 8 Sep 2026 07:33:29 -0700
Subject: [PATCH] Count captured token trees without intermediate parser steps

[skip ci]
---
 compiler/rustc_ast/src/tokenstream.rs       | 23 +++++++++
 compiler/rustc_ast/src/tokenstream/tests.rs | 54 +++++++++++++++++++++
 compiler/rustc_parse/src/parser/mod.rs      |  8 +++
 3 files changed, 85 insertions(+)

diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
index df71aad0111cd..538df146c291f 100644
--- a/compiler/rustc_ast/src/tokenstream.rs
+++ b/compiler/rustc_ast/src/tokenstream.rs
@@ -15,6 +15,7 @@ use rustc_data_structures::sync;
 use rustc_macros::{Decodable, Encodable, StableHash, Walkable};
 use rustc_serialize::{Decodable, Encodable};
 use rustc_span::{DUMMY_SP, Span, SpanDecoder, SpanEncoder, Symbol, sym};
+use smallvec::SmallVec;
 use thin_vec::ThinVec;
 
 use crate::ast::AttrStyle;
@@ -1005,6 +1006,28 @@ impl TokenCursor {
         self.curr.bump_to_end()
     }
 
+    pub fn bump_to_end_with_count(&mut self) -> u32 {
+        let mut count = 0;
+        let mut pending = SmallVec::<[std::slice::Iter<'_, TokenTree>; 8]>::new();
+        pending.push(self.curr.stream.0[self.curr.next_idx..].iter());
+        while let Some(iter) = pending.last_mut() {
+            match iter.next() {
+                Some(TokenTree::Token(..)) => count += 1,
+                Some(TokenTree::Delimited(_, _, delim, stream)) => {
+                    if !delim.skip() {
+                        count += 2;
+                    }
+                    pending.push(stream.0.iter());
+                }
+                None => {
+                    pending.pop();
+                }
+            }
+        }
+        self.curr.bump_to_end();
+        count
+    }
+
     /// Note: the outermost stream has depth of 0.
     #[inline]
     pub fn depth(&self) -> usize {
diff --git a/compiler/rustc_ast/src/tokenstream/tests.rs b/compiler/rustc_ast/src/tokenstream/tests.rs
index 6c7e82a97c58e..ea64707f4d228 100644
--- a/compiler/rustc_ast/src/tokenstream/tests.rs
+++ b/compiler/rustc_ast/src/tokenstream/tests.rs
@@ -11,3 +11,57 @@ fn test_token_stream_iter() {
     let iter = ts.iter();
     assert_eq!(iter.size_hint(), (1, Some(1)));
 }
+
+#[test]
+fn test_counted_token_cursor_skip() {
+    use crate::token::{Delimiter, InvisibleOrigin, MetaVarKind};
+    use crate::tokenstream::{DelimSpacing, DelimSpan, Spacing, TokenCursor, TokenTree};
+
+    let leaf = TokenTree::token_alone(TokenKind::ShrEq, DUMMY_SP);
+    let mut stream = TokenStream::new(vec![leaf.clone(); 120]);
+    let delimiters = [
+        Delimiter::Brace,
+        Delimiter::Parenthesis,
+        Delimiter::Bracket,
+        Delimiter::Invisible(InvisibleOrigin::ProcMacro),
+        Delimiter::Invisible(InvisibleOrigin::MetaVar(MetaVarKind::Block)),
+    ];
+    for depth in 0..25 {
+        stream = TokenStream::new(vec![
+            leaf.clone(),
+            TokenTree::Delimited(
+                DelimSpan::dummy(),
+                DelimSpacing::new(Spacing::Alone, Spacing::Alone),
+                delimiters[depth % delimiters.len()],
+                stream,
+            ),
+            leaf.clone(),
+        ]);
+    }
+    let mut cursor = TokenCursor::new(stream);
+    loop {
+        let token = cursor.next_and_bump().0;
+        if token.kind == TokenKind::Eof {
+            break;
+        }
+        if token.kind.open_delim().is_some() {
+            let depth = cursor.depth();
+            let mut stepped = cursor.clone();
+            let mut skipped = cursor.clone();
+            let skipped_count = skipped.bump_to_end_with_count();
+            let skipped_close = skipped.next_and_bump();
+            let mut stepped_count = 0;
+            loop {
+                let next = stepped.next_and_bump();
+                stepped_count += 1;
+                if stepped.depth() < depth {
+                    assert_eq!(next, skipped_close);
+                    break;
+                }
+            }
+            assert_eq!(stepped_count, skipped_count + 1);
+            assert_eq!(stepped.next_and_bump(), skipped.next_and_bump());
+            assert_eq!(stepped.depth(), skipped.depth());
+        }
+    }
+}
diff --git a/compiler/rustc_parse/src/parser/mod.rs b/compiler/rustc_parse/src/parser/mod.rs
index 80c1eeb4ef041..4898a59ad8be1 100644
--- a/compiler/rustc_parse/src/parser/mod.rs
+++ b/compiler/rustc_parse/src/parser/mod.rs
@@ -1406,6 +1406,14 @@ impl<'a> Parser<'a> {
                 self.token_cursor.bump_to_end();
                 self.bump();
                 debug_assert_eq!(self.token_cursor.depth(), target_depth);
+            } else if let TokenTree::Delimited(span, ..) = &tree
+                && !span.close.is_dummy()
+            {
+                // Capturing needs the token count, but no intermediate parser state.
+                // A dummy close span needs the normal loop's fallback-span propagation.
+                self.num_bump_calls += self.token_cursor.bump_to_end_with_count();
+                self.bump();
+                debug_assert_eq!(self.token_cursor.depth(), target_depth);
             } else {
                 loop {
                     // Advance one token at a time, so `TokenCursor::next_and_bump()`