Separate the public tt API from how they're stored

And don't provide any insight into the actual storage (no references, no "cheat" methods like `flat_tokens()`), as we're going to rewrite it to be more memory efficient.

Author: ChayimFriedman2

crates/cfg/src/cfg_expr.rs

@@ -197,17 +197,18 @@ fn next_cfg_expr(it: &mut tt::iter::TtIter<'_>) -> Option<CfgExpr> {
         Some(_) => return Some(CfgExpr::Invalid),
     };
 
-    let ret = match it.peek() {
+    let mut it_clone = it.clone();
+    let ret = match it_clone.next() {
         Some(TtElement::Leaf(tt::Leaf::Punct(punct)))
             // Don't consume on e.g. `=>`.
             if punct.char == '='
                 && (punct.spacing == tt::Spacing::Alone
-                    || it.remaining().flat_tokens().get(1).is_none_or(|peek2| {
-                        !matches!(peek2, tt::TokenTree::Leaf(tt::Leaf::Punct(_)))
+                    || it_clone.peek().is_none_or(|peek2| {
+                        !matches!(peek2, tt::TtElement::Leaf(tt::Leaf::Punct(_)))
                     })) =>
         {
-            match it.remaining().flat_tokens().get(1) {
-                Some(tt::TokenTree::Leaf(tt::Leaf::Literal(literal))) => {
+            match it_clone.next() {
+                Some(tt::TtElement::Leaf(tt::Leaf::Literal(literal))) => {
                     it.next();
                     it.next();
                     CfgAtom::KeyValue { key: name, value: literal.symbol.clone() }.into()

crates/hir-def/src/nameres/attr_resolution.rs

@@ -114,7 +114,7 @@ pub(super) fn attr_macro_as_call_id(
     let arg = match macro_attr.input.as_deref() {
         Some(AttrInput::TokenTree(tt)) => {
             let mut tt = tt.clone();
-            tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Invisible;
+            tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Invisible);
             Some(tt)
         }
 

crates/hir-def/src/nameres/collector.rs

@@ -2466,8 +2466,8 @@ impl ModCollector<'_, '_> {
                 None => {
                     let explicit_name =
                         attrs.by_key(sym::rustc_builtin_macro).tt_values().next().and_then(|tt| {
-                            match tt.token_trees().flat_tokens().first() {
-                                Some(tt::TokenTree::Leaf(tt::Leaf::Ident(name))) => Some(name),
+                            match tt.token_trees().iter().next() {
+                                Some(tt::TtElement::Leaf(tt::Leaf::Ident(name))) => Some(name),
                                 _ => None,
                             }
                         });

crates/hir-def/src/nameres/proc_macro.rs

@@ -2,10 +2,11 @@
 
 use hir_expand::name::{AsName, Name};
 use intern::sym;
+use itertools::Itertools;
 
 use crate::{
     item_tree::Attrs,
-    tt::{Leaf, TokenTree, TopSubtree, TtElement},
+    tt::{Leaf, TopSubtree, TtElement},
 };
 
 #[derive(Debug, PartialEq, Eq)]
@@ -61,35 +62,35 @@ impl Attrs<'_> {
 
 // This fn is intended for `#[proc_macro_derive(..)]` and `#[rustc_builtin_macro(..)]`, which have
 // the same structure.
-#[rustfmt::skip]
 pub(crate) fn parse_macro_name_and_helper_attrs(tt: &TopSubtree) -> Option<(Name, Box<[Name]>)> {
-    match tt.token_trees().flat_tokens() {
+    if let Some([TtElement::Leaf(Leaf::Ident(trait_name))]) =
+        tt.token_trees().iter().collect_array()
+    {
         // `#[proc_macro_derive(Trait)]`
         // `#[rustc_builtin_macro(Trait)]`
-        [TokenTree::Leaf(Leaf::Ident(trait_name))] => Some((trait_name.as_name(), Box::new([]))),
-
+        Some((trait_name.as_name(), Box::new([])))
+    } else if let Some(
+        [
+            TtElement::Leaf(Leaf::Ident(trait_name)),
+            TtElement::Leaf(Leaf::Punct(comma)),
+            TtElement::Leaf(Leaf::Ident(attributes)),
+            TtElement::Subtree(_, helpers),
+        ],
+    ) = tt.token_trees().iter().collect_array()
+        && comma.char == ','
+        && attributes.sym == sym::attributes
+    {
         // `#[proc_macro_derive(Trait, attributes(helper1, helper2, ...))]`
         // `#[rustc_builtin_macro(Trait, attributes(helper1, helper2, ...))]`
-        [
-            TokenTree::Leaf(Leaf::Ident(trait_name)),
-            TokenTree::Leaf(Leaf::Punct(comma)),
-            TokenTree::Leaf(Leaf::Ident(attributes)),
-            TokenTree::Subtree(_),
-            ..
-        ] if comma.char == ',' && attributes.sym == sym::attributes =>
-        {
-            let helpers = tt::TokenTreesView::new(&tt.token_trees().flat_tokens()[3..]).try_into_subtree()?;
-            let helpers = helpers
-                .iter()
-                .filter_map(|tt| match tt {
-                    TtElement::Leaf(Leaf::Ident(helper)) => Some(helper.as_name()),
-                    _ => None,
-                })
-                .collect::<Box<[_]>>();
-
-            Some((trait_name.as_name(), helpers))
-        }
+        let helpers = helpers
+            .filter_map(|tt| match tt {
+                TtElement::Leaf(Leaf::Ident(helper)) => Some(helper.as_name()),
+                _ => None,
+            })
+            .collect::<Box<[_]>>();
 
-        _ => None,
+        Some((trait_name.as_name(), helpers))
+    } else {
+        None
     }
 }

crates/hir-expand/src/attrs.rs

@@ -36,6 +36,7 @@ use base_db::Crate;
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
 use intern::{Interned, Symbol};
+use itertools::Itertools;
 use mbe::{DelimiterKind, Punct};
 use parser::T;
 use smallvec::SmallVec;
@@ -453,10 +454,10 @@ impl Attr {
     }
 
     /// #[path(ident)]
-    pub fn single_ident_value(&self) -> Option<&tt::Ident> {
+    pub fn single_ident_value(&self) -> Option<tt::Ident> {
         match self.input.as_deref()? {
-            AttrInput::TokenTree(tt) => match tt.token_trees().flat_tokens() {
-                [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] => Some(ident),
+            AttrInput::TokenTree(tt) => match tt.token_trees().iter().collect_array() {
+                Some([tt::TtElement::Leaf(tt::Leaf::Ident(ident))]) => Some(ident),
                 _ => None,
             },
             _ => None,
@@ -492,7 +493,7 @@ fn parse_path_comma_token_tree<'a>(
     args.token_trees()
         .split(|tt| matches!(tt, tt::TtElement::Leaf(tt::Leaf::Punct(Punct { char: ',', .. }))))
         .filter_map(move |tts| {
-            let span = tts.flat_tokens().first()?.first_span();
+            let span = tts.first_span()?;
             Some((ModPath::from_tt(db, tts)?, span, tts))
         })
 }
@@ -611,16 +612,12 @@ impl AttrId {
         else {
             return derive_attr_range;
         };
-        let (Some(first_tt), Some(last_tt)) =
-            (derive_tts.flat_tokens().first(), derive_tts.flat_tokens().last())
+        let (Some(first_span), Some(last_span)) = (derive_tts.first_span(), derive_tts.last_span())
         else {
             return derive_attr_range;
         };
-        let start = first_tt.first_span().range.start();
-        let end = match last_tt {
-            tt::TokenTree::Leaf(it) => it.span().range.end(),
-            tt::TokenTree::Subtree(it) => it.delimiter.close.range.end(),
-        };
+        let start = first_span.range.start();
+        let end = last_span.range.end();
         TextRange::new(start, end)
     }
 }

crates/hir-expand/src/builtin/fn_macro.rs

@@ -210,7 +210,7 @@ fn stringify_expand(
     tt: &tt::TopSubtree,
     span: Span,
 ) -> ExpandResult<tt::TopSubtree> {
-    let pretty = ::tt::pretty(tt.token_trees().flat_tokens());
+    let pretty = ::tt::pretty(tt.token_trees());
 
     let expanded = quote! {span =>
         #pretty
@@ -283,7 +283,7 @@ fn format_args_expand(
 ) -> ExpandResult<tt::TopSubtree> {
     let pound = mk_pound(span);
     let mut tt = tt.clone();
-    tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Parenthesis;
+    tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
     ExpandResult::ok(quote! {span =>
         builtin #pound format_args #tt
     })
@@ -297,14 +297,14 @@ fn format_args_nl_expand(
 ) -> ExpandResult<tt::TopSubtree> {
     let pound = mk_pound(span);
     let mut tt = tt.clone();
-    tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Parenthesis;
-    if let Some(tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal {
-        symbol: text,
-        kind: tt::LitKind::Str,
-        ..
-    }))) = tt.0.get_mut(1)
+    tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
+    let lit = tt.as_token_trees().iter_flat_tokens().nth(1);
+    if let Some(tt::TokenTree::Leaf(tt::Leaf::Literal(
+        mut lit @ tt::Literal { kind: tt::LitKind::Str, .. },
+    ))) = lit
     {
-        *text = Symbol::intern(&format_smolstr!("{}\\n", text.as_str()));
+        lit.symbol = Symbol::intern(&format_smolstr!("{}\\n", lit.symbol.as_str()));
+        tt.set_token(1, lit.into());
     }
     ExpandResult::ok(quote! {span =>
         builtin #pound format_args #tt
@@ -318,7 +318,7 @@ fn asm_expand(
     span: Span,
 ) -> ExpandResult<tt::TopSubtree> {
     let mut tt = tt.clone();
-    tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Parenthesis;
+    tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
     let pound = mk_pound(span);
     let expanded = quote! {span =>
         builtin #pound asm #tt
@@ -333,7 +333,7 @@ fn global_asm_expand(
     span: Span,
 ) -> ExpandResult<tt::TopSubtree> {
     let mut tt = tt.clone();
-    tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Parenthesis;
+    tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
     let pound = mk_pound(span);
     let expanded = quote! {span =>
         builtin #pound global_asm #tt
@@ -348,7 +348,7 @@ fn naked_asm_expand(
     span: Span,
 ) -> ExpandResult<tt::TopSubtree> {
     let mut tt = tt.clone();
-    tt.top_subtree_delimiter_mut().kind = tt::DelimiterKind::Parenthesis;
+    tt.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
     let pound = mk_pound(span);
     let expanded = quote! {span =>
         builtin #pound naked_asm #tt
@@ -478,11 +478,11 @@ fn unreachable_expand(
 
     // Pass the original arguments
     let mut subtree = tt.clone();
-    *subtree.top_subtree_delimiter_mut() = tt::Delimiter {
+    subtree.set_top_subtree_delimiter_kind(tt::DelimiterKind::Parenthesis);
+    subtree.set_top_subtree_delimiter_span(tt::DelimSpan {
         open: call_site_span,
         close: call_site_span,
-        kind: tt::DelimiterKind::Parenthesis,
-    };
+    });
 
     // Expand to a macro call `$crate::panic::panic_{edition}`
     let call = quote!(call_site_span =>#dollar_crate::panic::#mac! #subtree);
@@ -518,16 +518,17 @@ fn compile_error_expand(
     tt: &tt::TopSubtree,
     span: Span,
 ) -> ExpandResult<tt::TopSubtree> {
-    let err = match &*tt.0 {
-        [
-            _,
-            tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal {
-                symbol: text,
-                span: _,
-                kind: tt::LitKind::Str | tt::LitKind::StrRaw(_),
-                suffix: _,
-            })),
-        ] => ExpandError::other(span, Box::from(unescape_symbol(text).as_str())),
+    let err = match tt.iter().collect_array() {
+        Some(
+            [
+                tt::TtElement::Leaf(tt::Leaf::Literal(tt::Literal {
+                    symbol: text,
+                    span: _,
+                    kind: tt::LitKind::Str | tt::LitKind::StrRaw(_),
+                    suffix: _,
+                })),
+            ],
+        ) => ExpandError::other(span, Box::from(unescape_symbol(&text).as_str())),
         _ => ExpandError::other(span, "`compile_error!` argument must be a string"),
     };
 
@@ -556,7 +557,7 @@ fn concat_expand(
         // to ensure the right parsing order, so skip the parentheses here. Ideally we'd
         // implement rustc's model. cc https://github.com/rust-lang/rust-analyzer/pull/10623
         if let TtElement::Subtree(subtree, subtree_iter) = &t
-            && let [tt::TokenTree::Leaf(tt)] = subtree_iter.remaining().flat_tokens()
+            && let Some([tt::TtElement::Leaf(tt)]) = subtree_iter.clone().collect_array()
             && subtree.delimiter.kind == tt::DelimiterKind::Parenthesis
         {
             t = TtElement::Leaf(tt);
@@ -663,7 +664,7 @@ fn concat_bytes_expand(
                 kind,
                 suffix: _,
             })) => {
-                record_span(*span);
+                record_span(span);
                 match kind {
                     tt::LitKind::Byte => {
                         if let Ok(b) = unescape_byte(text.as_str()) {
@@ -679,7 +680,7 @@ fn concat_bytes_expand(
                         bytes.extend(text.as_str().escape_debug());
                     }
                     _ => {
-                        err.get_or_insert(ExpandError::other(*span, "unexpected token"));
+                        err.get_or_insert(ExpandError::other(span, "unexpected token"));
                         break;
                     }
                 }
@@ -733,15 +734,15 @@ fn concat_bytes_expand_subtree(
                 if let Ok(b) = unescape_byte(text.as_str()) {
                     bytes.extend(b.escape_ascii().filter_map(|it| char::from_u32(it as u32)));
                 }
-                record_span(*span);
+                record_span(span);
             }
             TtElement::Leaf(tt::Leaf::Literal(tt::Literal {
                 symbol: text,
                 span,
                 kind: tt::LitKind::Integer,
                 suffix: _,
             })) => {
-                record_span(*span);
+                record_span(span);
                 if let Ok(b) = text.as_str().parse::<u8>() {
                     bytes.extend(b.escape_ascii().filter_map(|it| char::from_u32(it as u32)));
                 }
@@ -796,13 +797,13 @@ fn parse_string(tt: &tt::TopSubtree) -> Result<(Symbol, Span), ExpandError> {
                 span,
                 kind: tt::LitKind::Str,
                 suffix: _,
-            })) => Ok((unescape_symbol(text), *span)),
+            })) => Ok((unescape_symbol(&text), span)),
             TtElement::Leaf(tt::Leaf::Literal(tt::Literal {
                 symbol: text,
                 span,
                 kind: tt::LitKind::StrRaw(_),
                 suffix: _,
-            })) => Ok((text.clone(), *span)),
+            })) => Ok((text.clone(), span)),
             TtElement::Leaf(l) => Err(*l.span()),
             TtElement::Subtree(tt, _) => Err(tt.delimiter.open.cover(tt.delimiter.close)),
         }

crates/hir-expand/src/builtin/quote.rs

@@ -163,7 +163,7 @@ impl ToTokenTree for crate::tt::SubtreeView<'_> {
 
 impl ToTokenTree for crate::tt::TopSubtree {
     fn to_tokens(self, _: Span, builder: &mut TopSubtreeBuilder) {
-        builder.extend_tt_dangerous(self.0);
+        builder.extend_with_tt(self.as_token_trees());
     }
 }
 
@@ -172,10 +172,9 @@ impl ToTokenTree for crate::tt::TtElement<'_> {
         match self {
             crate::tt::TtElement::Leaf(leaf) => builder.push(leaf.clone()),
             crate::tt::TtElement::Subtree(subtree, subtree_iter) => {
-                builder.extend_tt_dangerous(
-                    std::iter::once(crate::tt::TokenTree::Subtree(subtree.clone()))
-                        .chain(subtree_iter.remaining().flat_tokens().iter().cloned()),
-                );
+                builder.open(subtree.delimiter.kind, subtree.delimiter.open);
+                builder.extend_with_tt(subtree_iter.remaining());
+                builder.close(subtree.delimiter.close);
             }
         }
     }