Store token trees in contiguous Vec instead of as a tree

I expected this to be faster (due to less allocations and better cache locality), but benchmarked it is not (neither it is slower). Memory usage, however, drops by ~50mb (of `analysis-stats .`). I guess tt construction is just not hot.

This also simplifies using even less memory for token trees by compressing equal span, which I plan to do right after.

Some workflows are more easily expressed with a flat tt, while some are better expressed with a tree. With the right helpers, though (which was mostly a matter of trial and error), even the worst workflows become very easy indeed.

src/tools/rust-analyzer/Cargo.lock

@@ -1389,6 +1389,7 @@ version = "0.0.0"
 dependencies = [
  "proc-macro-api",
  "proc-macro-srv",
+ "tt",
 ]
 
 [[package]]

src/tools/rust-analyzer/crates/cfg/src/cfg_expr.rs

@@ -45,8 +45,8 @@ impl From<CfgAtom> for CfgExpr {
 
 impl CfgExpr {
     #[cfg(feature = "tt")]
-    pub fn parse<S>(tt: &tt::Subtree<S>) -> CfgExpr {
-        next_cfg_expr(&mut tt.token_trees.iter()).unwrap_or(CfgExpr::Invalid)
+    pub fn parse<S: Copy>(tt: &tt::TopSubtree<S>) -> CfgExpr {
+        next_cfg_expr(&mut tt.iter()).unwrap_or(CfgExpr::Invalid)
     }
 
     /// Fold the cfg by querying all basic `Atom` and `KeyValue` predicates.
@@ -66,19 +66,19 @@ impl CfgExpr {
 }
 
 #[cfg(feature = "tt")]
-fn next_cfg_expr<S>(it: &mut std::slice::Iter<'_, tt::TokenTree<S>>) -> Option<CfgExpr> {
+fn next_cfg_expr<S: Copy>(it: &mut tt::iter::TtIter<'_, S>) -> Option<CfgExpr> {
     use intern::sym;
+    use tt::iter::TtElement;
 
     let name = match it.next() {
         None => return None,
-        Some(tt::TokenTree::Leaf(tt::Leaf::Ident(ident))) => ident.sym.clone(),
+        Some(TtElement::Leaf(tt::Leaf::Ident(ident))) => ident.sym.clone(),
         Some(_) => return Some(CfgExpr::Invalid),
     };
 
-    // Peek
-    let ret = match it.as_slice().first() {
-        Some(tt::TokenTree::Leaf(tt::Leaf::Punct(punct))) if punct.char == '=' => {
-            match it.as_slice().get(1) {
+    let ret = match it.peek() {
+        Some(TtElement::Leaf(tt::Leaf::Punct(punct))) if punct.char == '=' => {
+            match it.remaining().flat_tokens().get(1) {
                 Some(tt::TokenTree::Leaf(tt::Leaf::Literal(literal))) => {
                     it.next();
                     it.next();
@@ -87,9 +87,8 @@ fn next_cfg_expr<S>(it: &mut std::slice::Iter<'_, tt::TokenTree<S>>) -> Option<C
                 _ => return Some(CfgExpr::Invalid),
             }
         }
-        Some(tt::TokenTree::Subtree(subtree)) => {
+        Some(TtElement::Subtree(_, mut sub_it)) => {
             it.next();
-            let mut sub_it = subtree.token_trees.iter();
             let mut subs = std::iter::from_fn(|| next_cfg_expr(&mut sub_it));
             match name {
                 s if s == sym::all => CfgExpr::All(subs.collect()),
@@ -104,7 +103,7 @@ fn next_cfg_expr<S>(it: &mut std::slice::Iter<'_, tt::TokenTree<S>>) -> Option<C
     };
 
     // Eat comma separator
-    if let Some(tt::TokenTree::Leaf(tt::Leaf::Punct(punct))) = it.as_slice().first() {
+    if let Some(TtElement::Leaf(tt::Leaf::Punct(punct))) = it.peek() {
         if punct.char == ',' {
             it.next();
         }

src/tools/rust-analyzer/crates/hir-def/src/attr.rs

@@ -1,6 +1,6 @@
 //! A higher level attributes based on TokenTree, with also some shortcuts.
 
-use std::{borrow::Cow, hash::Hash, ops, slice};
+use std::{borrow::Cow, hash::Hash, ops};
 
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
@@ -17,6 +17,7 @@ use syntax::{
     AstPtr,
 };
 use triomphe::Arc;
+use tt::iter::{TtElement, TtIter};
 
 use crate::{
     db::DefDatabase,
@@ -154,15 +155,15 @@ impl Attrs {
 
     pub fn has_doc_hidden(&self) -> bool {
         self.by_key(&sym::doc).tt_values().any(|tt| {
-            tt.delimiter.kind == DelimiterKind::Parenthesis &&
-                matches!(&*tt.token_trees, [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.sym == sym::hidden)
+            tt.top_subtree().delimiter.kind == DelimiterKind::Parenthesis &&
+                matches!(tt.token_trees().flat_tokens(), [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.sym == sym::hidden)
         })
     }
 
     pub fn has_doc_notable_trait(&self) -> bool {
         self.by_key(&sym::doc).tt_values().any(|tt| {
-            tt.delimiter.kind == DelimiterKind::Parenthesis &&
-                matches!(&*tt.token_trees, [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.sym == sym::notable_trait)
+            tt.top_subtree().delimiter.kind == DelimiterKind::Parenthesis &&
+                matches!(tt.token_trees().flat_tokens(), [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.sym == sym::notable_trait)
         })
     }
 
@@ -245,8 +246,8 @@ impl From<DocAtom> for DocExpr {
 }
 
 impl DocExpr {
-    fn parse<S>(tt: &tt::Subtree<S>) -> DocExpr {
-        next_doc_expr(&mut tt.token_trees.iter()).unwrap_or(DocExpr::Invalid)
+    fn parse<S: Copy>(tt: &tt::TopSubtree<S>) -> DocExpr {
+        next_doc_expr(tt.iter()).unwrap_or(DocExpr::Invalid)
     }
 
     pub fn aliases(&self) -> &[Symbol] {
@@ -260,62 +261,47 @@ impl DocExpr {
     }
 }
 
-fn next_doc_expr<S>(it: &mut slice::Iter<'_, tt::TokenTree<S>>) -> Option<DocExpr> {
+fn next_doc_expr<S: Copy>(mut it: TtIter<'_, S>) -> Option<DocExpr> {
     let name = match it.next() {
         None => return None,
-        Some(tt::TokenTree::Leaf(tt::Leaf::Ident(ident))) => ident.sym.clone(),
+        Some(TtElement::Leaf(tt::Leaf::Ident(ident))) => ident.sym.clone(),
         Some(_) => return Some(DocExpr::Invalid),
     };
 
     // Peek
-    let ret = match it.as_slice().first() {
-        Some(tt::TokenTree::Leaf(tt::Leaf::Punct(punct))) if punct.char == '=' => {
-            match it.as_slice().get(1) {
-                Some(tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal {
+    let ret = match it.peek() {
+        Some(TtElement::Leaf(tt::Leaf::Punct(punct))) if punct.char == '=' => {
+            it.next();
+            match it.next() {
+                Some(TtElement::Leaf(tt::Leaf::Literal(tt::Literal {
                     symbol: text,
                     kind: tt::LitKind::Str,
                     ..
-                }))) => {
-                    it.next();
-                    it.next();
-                    DocAtom::KeyValue { key: name, value: text.clone() }.into()
-                }
+                }))) => DocAtom::KeyValue { key: name, value: text.clone() }.into(),
                 _ => return Some(DocExpr::Invalid),
             }
         }
-        Some(tt::TokenTree::Subtree(subtree)) => {
+        Some(TtElement::Subtree(_, subtree_iter)) => {
             it.next();
-            let subs = parse_comma_sep(subtree);
+            let subs = parse_comma_sep(subtree_iter);
             match &name {
                 s if *s == sym::alias => DocExpr::Alias(subs),
                 _ => DocExpr::Invalid,
             }
         }
         _ => DocAtom::Flag(name).into(),
     };
-
-    // Eat comma separator
-    if let Some(tt::TokenTree::Leaf(tt::Leaf::Punct(punct))) = it.as_slice().first() {
-        if punct.char == ',' {
-            it.next();
-        }
-    }
     Some(ret)
 }
 
-fn parse_comma_sep<S>(subtree: &tt::Subtree<S>) -> Vec<Symbol> {
-    subtree
-        .token_trees
-        .iter()
-        .filter_map(|tt| match tt {
-            tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal {
-                kind: tt::LitKind::Str,
-                symbol,
-                ..
-            })) => Some(symbol.clone()),
-            _ => None,
-        })
-        .collect()
+fn parse_comma_sep<S>(iter: TtIter<'_, S>) -> Vec<Symbol> {
+    iter.filter_map(|tt| match tt {
+        TtElement::Leaf(tt::Leaf::Literal(tt::Literal {
+            kind: tt::LitKind::Str, symbol, ..
+        })) => Some(symbol.clone()),
+        _ => None,
+    })
+    .collect()
 }
 
 impl AttrsWithOwner {
@@ -563,7 +549,7 @@ pub struct AttrQuery<'attr> {
 }
 
 impl<'attr> AttrQuery<'attr> {
-    pub fn tt_values(self) -> impl Iterator<Item = &'attr crate::tt::Subtree> {
+    pub fn tt_values(self) -> impl Iterator<Item = &'attr crate::tt::TopSubtree> {
         self.attrs().filter_map(|attr| attr.token_tree_value())
     }
 
@@ -596,12 +582,12 @@ impl<'attr> AttrQuery<'attr> {
     /// ```
     pub fn find_string_value_in_tt(self, key: &'attr Symbol) -> Option<&'attr str> {
         self.tt_values().find_map(|tt| {
-            let name = tt.token_trees.iter()
-                .skip_while(|tt| !matches!(tt, tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident { sym, ..} )) if *sym == *key))
+            let name = tt.iter()
+                .skip_while(|tt| !matches!(tt, TtElement::Leaf(tt::Leaf::Ident(tt::Ident { sym, ..} )) if *sym == *key))
                 .nth(2);
 
             match name {
-                Some(tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal{  symbol: text, kind: tt::LitKind::Str | tt::LitKind::StrRaw(_) , ..}))) => Some(text.as_str()),
+                Some(TtElement::Leaf(tt::Leaf::Literal(tt::Literal{  symbol: text, kind: tt::LitKind::Str | tt::LitKind::StrRaw(_) , ..}))) => Some(text.as_str()),
                 _ => None
             }
         })

src/tools/rust-analyzer/crates/hir-def/src/data.rs

@@ -11,6 +11,7 @@ use la_arena::{Idx, RawIdx};
 use smallvec::SmallVec;
 use syntax::{ast, Parse};
 use triomphe::Arc;
+use tt::iter::TtElement;
 
 use crate::{
     db::DefDatabase,
@@ -156,20 +157,21 @@ impl FunctionData {
     }
 }
 
-fn parse_rustc_legacy_const_generics(tt: &crate::tt::Subtree) -> Box<[u32]> {
+fn parse_rustc_legacy_const_generics(tt: &crate::tt::TopSubtree) -> Box<[u32]> {
     let mut indices = Vec::new();
-    for args in tt.token_trees.chunks(2) {
-        match &args[0] {
-            tt::TokenTree::Leaf(tt::Leaf::Literal(lit)) => match lit.symbol.as_str().parse() {
+    let mut iter = tt.iter();
+    while let (Some(first), second) = (iter.next(), iter.next()) {
+        match first {
+            TtElement::Leaf(tt::Leaf::Literal(lit)) => match lit.symbol.as_str().parse() {
                 Ok(index) => indices.push(index),
                 Err(_) => break,
             },
             _ => break,
         }
 
-        if let Some(comma) = args.get(1) {
+        if let Some(comma) = second {
             match comma {
-                tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) if punct.char == ',' => {}
+                TtElement::Leaf(tt::Leaf::Punct(punct)) if punct.char == ',' => {}
                 _ => break,
             }
         }
@@ -267,8 +269,8 @@ impl TraitData {
             attrs.by_key(&sym::rustc_skip_array_during_method_dispatch).exists();
         let mut skip_boxed_slice_during_method_dispatch = false;
         for tt in attrs.by_key(&sym::rustc_skip_during_method_dispatch).tt_values() {
-            for tt in tt.token_trees.iter() {
-                if let crate::tt::TokenTree::Leaf(tt::Leaf::Ident(ident)) = tt {
+            for tt in tt.iter() {
+                if let tt::iter::TtElement::Leaf(tt::Leaf::Ident(ident)) = tt {
                     skip_array_during_method_dispatch |= ident.sym == sym::array;
                     skip_boxed_slice_during_method_dispatch |= ident.sym == sym::boxed_slice;
                 }
@@ -421,7 +423,7 @@ impl Macro2Data {
             .by_key(&sym::rustc_builtin_macro)
             .tt_values()
             .next()
-            .and_then(|attr| parse_macro_name_and_helper_attrs(&attr.token_trees))
+            .and_then(parse_macro_name_and_helper_attrs)
             .map(|(_, helpers)| helpers);
 
         Arc::new(Macro2Data {

src/tools/rust-analyzer/crates/hir-def/src/data/adt.rs

@@ -10,6 +10,7 @@ use intern::sym;
 use la_arena::Arena;
 use rustc_abi::{Align, Integer, IntegerType, ReprFlags, ReprOptions};
 use triomphe::Arc;
+use tt::iter::TtElement;
 
 use crate::{
     builtin_type::{BuiltinInt, BuiltinUint},
@@ -20,7 +21,7 @@ use crate::{
     },
     lang_item::LangItem,
     nameres::diagnostics::{DefDiagnostic, DefDiagnostics},
-    tt::{Delimiter, DelimiterKind, Leaf, Subtree, TokenTree},
+    tt::{Delimiter, DelimiterKind, Leaf, TopSubtree},
     type_ref::{TypeRefId, TypesMap},
     visibility::RawVisibility,
     EnumId, EnumVariantId, LocalFieldId, LocalModuleId, Lookup, StructId, UnionId, VariantId,
@@ -95,8 +96,8 @@ fn repr_from_value(
     item_tree.attrs(db, krate, of).by_key(&sym::repr).tt_values().find_map(parse_repr_tt)
 }
 
-fn parse_repr_tt(tt: &Subtree) -> Option<ReprOptions> {
-    match tt.delimiter {
+fn parse_repr_tt(tt: &TopSubtree) -> Option<ReprOptions> {
+    match tt.top_subtree().delimiter {
         Delimiter { kind: DelimiterKind::Parenthesis, .. } => {}
         _ => return None,
     }
@@ -106,14 +107,14 @@ fn parse_repr_tt(tt: &Subtree) -> Option<ReprOptions> {
     let mut max_align: Option<Align> = None;
     let mut min_pack: Option<Align> = None;
 
-    let mut tts = tt.token_trees.iter().peekable();
+    let mut tts = tt.iter();
     while let Some(tt) = tts.next() {
-        if let TokenTree::Leaf(Leaf::Ident(ident)) = tt {
+        if let TtElement::Leaf(Leaf::Ident(ident)) = tt {
             flags.insert(match &ident.sym {
                 s if *s == sym::packed => {
-                    let pack = if let Some(TokenTree::Subtree(tt)) = tts.peek() {
+                    let pack = if let Some(TtElement::Subtree(_, mut tt_iter)) = tts.peek() {
                         tts.next();
-                        if let Some(TokenTree::Leaf(Leaf::Literal(lit))) = tt.token_trees.first() {
+                        if let Some(TtElement::Leaf(Leaf::Literal(lit))) = tt_iter.next() {
                             lit.symbol.as_str().parse().unwrap_or_default()
                         } else {
                             0
@@ -127,9 +128,9 @@ fn parse_repr_tt(tt: &Subtree) -> Option<ReprOptions> {
                     ReprFlags::empty()
                 }
                 s if *s == sym::align => {
-                    if let Some(TokenTree::Subtree(tt)) = tts.peek() {
+                    if let Some(TtElement::Subtree(_, mut tt_iter)) = tts.peek() {
                         tts.next();
-                        if let Some(TokenTree::Leaf(Leaf::Literal(lit))) = tt.token_trees.first() {
+                        if let Some(TtElement::Leaf(Leaf::Literal(lit))) = tt_iter.next() {
                             if let Ok(align) = lit.symbol.as_str().parse() {
                                 let align = Align::from_bytes(align).ok();
                                 max_align = max_align.max(align);

src/tools/rust-analyzer/crates/hir-def/src/db.rs

@@ -21,6 +21,7 @@ use crate::{
     item_tree::{AttrOwner, ItemTree, ItemTreeSourceMaps},
     lang_item::{self, LangItem, LangItemTarget, LangItems},
     nameres::{diagnostics::DefDiagnostics, DefMap},
+    tt,
     type_ref::TypesSourceMap,
     visibility::{self, Visibility},
     AttrDefId, BlockId, BlockLoc, ConstBlockId, ConstBlockLoc, ConstId, ConstLoc, DefWithBodyId,
@@ -294,14 +295,14 @@ fn crate_supports_no_std(db: &dyn DefDatabase, crate_id: CrateId) -> bool {
         // This is a `cfg_attr`; check if it could possibly expand to `no_std`.
         // Syntax is: `#[cfg_attr(condition(cfg, style), attr0, attr1, <...>)]`
         let tt = match attr.token_tree_value() {
-            Some(tt) => &tt.token_trees,
+            Some(tt) => tt.token_trees(),
             None => continue,
         };
 
         let segments =
-            tt.split(|tt| matches!(tt, tt::TokenTree::Leaf(tt::Leaf::Punct(p)) if p.char == ','));
+            tt.split(|tt| matches!(tt, tt::TtElement::Leaf(tt::Leaf::Punct(p)) if p.char == ','));
         for output in segments.skip(1) {
-            match output {
+            match output.flat_tokens() {
                 [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.sym == sym::no_std => {
                     return true
                 }

src/tools/rust-analyzer/crates/hir-def/src/macro_expansion_tests/mod.rs

@@ -40,7 +40,7 @@ use crate::{
     resolver::HasResolver,
     src::HasSource,
     test_db::TestDB,
-    tt::Subtree,
+    tt::TopSubtree,
     AdtId, AsMacroCall, Lookup, ModuleDefId,
 };
 
@@ -313,14 +313,14 @@ struct IdentityWhenValidProcMacroExpander;
 impl ProcMacroExpander for IdentityWhenValidProcMacroExpander {
     fn expand(
         &self,
-        subtree: &Subtree,
-        _: Option<&Subtree>,
+        subtree: &TopSubtree,
+        _: Option<&TopSubtree>,
         _: &base_db::Env,
         _: Span,
         _: Span,
         _: Span,
         _: Option<String>,
-    ) -> Result<Subtree, ProcMacroExpansionError> {
+    ) -> Result<TopSubtree, ProcMacroExpansionError> {
         let (parse, _) = syntax_bridge::token_tree_to_syntax_node(
             subtree,
             syntax_bridge::TopEntryPoint::MacroItems,

src/tools/rust-analyzer/crates/hir-def/src/nameres/attr_resolution.rs

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