Bring our TokenTree model closer to rustc's

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>(tt: &tt::TokenStream<S>) -> CfgExpr {
+        next_cfg_expr(tt.trees()).unwrap_or(CfgExpr::Invalid)
     }
 
     /// Fold the cfg by querying all basic `Atom` and `KeyValue` predicates.
@@ -66,48 +66,52 @@ 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>(mut it: tt::RefTokenTreeCursor<'_, S>) -> Option<CfgExpr> {
     use intern::sym;
 
     let name = match it.next() {
         None => return None,
-        Some(tt::TokenTree::Leaf(tt::Leaf::Ident(ident))) => ident.sym.clone(),
+        Some(tt::TokenTree::Token(tt::Token { kind: tt::TokenKind::Ident(ident, _), .. }, _)) => {
+            ident.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) {
-                Some(tt::TokenTree::Leaf(tt::Leaf::Literal(literal))) => {
-                    it.next();
-                    it.next();
-                    CfgAtom::KeyValue { key: name, value: literal.symbol.clone() }.into()
+    let ret = match it.look_ahead(0) {
+        Some(tt::TokenTree::Token(tt::Token { kind: tt::TokenKind::Eq, .. }, _)) => {
+            match it.look_ahead(1) {
+                Some(tt::TokenTree::Token(
+                    tt::Token { kind: tt::TokenKind::Literal(lit), .. },
+                    _,
+                )) => {
+                    let res = CfgAtom::KeyValue { key: name, value: lit.symbol.clone() }.into();
+                    _ = it.next();
+                    _ = it.next();
+                    res
                 }
                 _ => return Some(CfgExpr::Invalid),
             }
         }
-        Some(tt::TokenTree::Subtree(subtree)) => {
-            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 {
+        Some(tt::TokenTree::Delimited(_, _, _, stream)) => {
+            let mut subs = std::iter::from_fn(|| next_cfg_expr(stream.trees()));
+            let res = match name {
                 s if s == sym::all => CfgExpr::All(subs.collect()),
                 s if s == sym::any => CfgExpr::Any(subs.collect()),
                 s if s == sym::not => {
                     CfgExpr::Not(Box::new(subs.next().unwrap_or(CfgExpr::Invalid)))
                 }
                 _ => CfgExpr::Invalid,
-            }
+            };
+            it.next();
+            res
         }
         _ => CfgAtom::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();
-        }
+    if let Some(tt::TokenTree::Token(tt::Token { kind: tt::TokenKind::Comma, .. }, _)) = it.next() {
+        it.next();
     }
     Some(ret)
 }

crates/hir-expand/src/db.rs

@@ -28,7 +28,7 @@ use crate::{
     MacroDefId, MacroDefKind, MacroFileId,
 };
 /// This is just to ensure the types of smart_macro_arg and macro_arg are the same
-type MacroArgResult = (Arc<tt::Subtree>, SyntaxFixupUndoInfo, Span);
+type MacroArgResult = (Arc<tt::TokenStream>, SyntaxFixupUndoInfo, Span);
 /// Total limit on the number of tokens produced by any macro invocation.
 ///
 /// If an invocation produces more tokens than this limit, it will not be stored in the database and

crates/hir-expand/src/lib.rs

@@ -75,6 +75,7 @@ pub mod tt {
     pub type Punct = ::tt::Punct<Span>;
     pub type Ident = ::tt::Ident<Span>;
     pub type TokenTree = ::tt::TokenTree<Span>;
+    pub type TokenStream = ::tt::TokenStream<Span>;
 }
 
 #[macro_export]

crates/mbe/src/lib.rs

@@ -23,15 +23,14 @@ mod tests;
 
 use span::{Edition, Span, SyntaxContextId};
 use syntax_bridge::to_parser_input;
-use tt::iter::TtIter;
-use tt::DelimSpan;
+use tt::{DelimSpan, RefTokenCursor, Token, TokenCursor, TokenKind, TokenStream};
 
 use std::fmt;
 use std::sync::Arc;
 
 use crate::parser::{MetaTemplate, MetaVarKind, Op};
 
-pub use tt::{Delimiter, DelimiterKind, Punct};
+pub use tt::Delimiter;
 
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum ParseError {
@@ -148,30 +147,32 @@ impl DeclarativeMacro {
 
     /// The old, `macro_rules! m {}` flavor.
     pub fn parse_macro_rules(
-        tt: &tt::Subtree<Span>,
+        stream: &TokenStream<Span>,
         ctx_edition: impl Copy + Fn(SyntaxContextId) -> Edition,
     ) -> DeclarativeMacro {
         // Note: this parsing can be implemented using mbe machinery itself, by
         // matching against `$($lhs:tt => $rhs:tt);*` pattern, but implementing
         // manually seems easier.
-        let mut src = TtIter::new(tt);
+        let mut cursor = RefTokenCursor::new(stream);
         let mut rules = Vec::new();
         let mut err = None;
 
-        while src.len() > 0 {
-            let rule = match Rule::parse(ctx_edition, &mut src) {
+        while let Some((token, _)) = cursor.next() {
+            let rule = match Rule::parse(ctx_edition, &mut cursor) {
                 Ok(it) => it,
                 Err(e) => {
                     err = Some(Box::new(e));
                     break;
                 }
             };
             rules.push(rule);
-            if let Err(()) = src.expect_char(';') {
-                if src.len() > 0 {
+            match cursor.next() {
+                Some((Token { kind: TokenKind::Semi, .. }, _)) => (),
+                Some((Token { span, .. }, _)) => {
                     err = Some(Box::new(ParseError::expected("expected `;`")));
+                    break;
                 }
-                break;
+                None => break,
             }
         }
 
@@ -187,8 +188,8 @@ impl DeclarativeMacro {
 
     /// The new, unstable `macro m {}` flavor.
     pub fn parse_macro2(
-        args: Option<&tt::Subtree<Span>>,
-        body: &tt::Subtree<Span>,
+        args: Option<&tt::TokenStream<Span>>,
+        body: &tt::TokenStream<Span>,
         ctx_edition: impl Copy + Fn(SyntaxContextId) -> Edition,
     ) -> DeclarativeMacro {
         let mut rules = Vec::new();
@@ -210,23 +211,25 @@ impl DeclarativeMacro {
             }
         } else {
             cov_mark::hit!(parse_macro_def_rules);
-            let mut src = TtIter::new(body);
-            while src.len() > 0 {
-                let rule = match Rule::parse(ctx_edition, &mut src) {
+            let mut cursor = RefTokenCursor::new(body);
+            while let Some((token, _)) = cursor.next() {
+                let rule = match Rule::parse(ctx_edition, &mut cursor) {
                     Ok(it) => it,
                     Err(e) => {
                         err = Some(Box::new(e));
                         break;
                     }
                 };
                 rules.push(rule);
-                if let Err(()) = src.expect_any_char(&[';', ',']) {
-                    if src.len() > 0 {
+                match cursor.next() {
+                    Some((Token { kind: TokenKind::Semi | TokenKind::Comma, .. }, _)) => (),
+                    Some((Token { span, .. }, _)) => {
                         err = Some(Box::new(ParseError::expected(
                             "expected `;` or `,` to delimit rules",
                         )));
+                        break;
                     }
-                    break;
+                    None => break,
                 }
             }
         }
@@ -251,19 +254,19 @@ impl DeclarativeMacro {
 
     pub fn expand(
         &self,
-        tt: &tt::Subtree<Span>,
+        tt: &tt::TokenStream<Span>,
         marker: impl Fn(&mut Span) + Copy,
         call_site: Span,
         def_site_edition: Edition,
-    ) -> ExpandResult<(tt::Subtree<Span>, MatchedArmIndex)> {
+    ) -> ExpandResult<(tt::TokenStream<Span>, MatchedArmIndex)> {
         expander::expand_rules(&self.rules, tt, marker, call_site, def_site_edition)
     }
 }
 
 impl Rule {
     fn parse(
         edition: impl Copy + Fn(SyntaxContextId) -> Edition,
-        src: &mut TtIter<'_, Span>,
+        src: &mut RefTokenCursor<'_, Span>,
     ) -> Result<Self, ParseError> {
         let lhs = src.expect_subtree().map_err(|()| ParseError::expected("expected subtree"))?;
         src.expect_char('=').map_err(|()| ParseError::expected("expected `=`"))?;
@@ -360,31 +363,26 @@ impl<T: Default, E> From<Result<T, E>> for ValueResult<T, E> {
 }
 
 pub fn expect_fragment(
-    tt_iter: &mut TtIter<'_, Span>,
+    cursor: &mut RefTokenCursor<'_, Span>,
     entry_point: ::parser::PrefixEntryPoint,
     edition: ::parser::Edition,
     delim_span: DelimSpan<Span>,
 ) -> ExpandResult<Option<tt::TokenTree<Span>>> {
     use ::parser;
-    let buffer = tt::buffer::TokenBuffer::from_tokens(tt_iter.as_slice());
-    let parser_input = to_parser_input(edition, &buffer);
+    let parser_input = to_parser_input(edition, cursor.clone());
     let tree_traversal = entry_point.parse(&parser_input, edition);
-    let mut cursor = buffer.begin();
     let mut error = false;
     for step in tree_traversal.iter() {
         match step {
             parser::Step::Token { kind, mut n_input_tokens } => {
-                if kind == ::parser::SyntaxKind::LIFETIME_IDENT {
-                    n_input_tokens = 2;
-                }
                 for _ in 0..n_input_tokens {
-                    cursor = cursor.bump_subtree();
+                    cursor.next();
                 }
             }
             parser::Step::FloatSplit { .. } => {
                 // FIXME: We need to split the tree properly here, but mutating the token trees
                 // in the buffer is somewhat tricky to pull off.
-                cursor = cursor.bump_subtree();
+                cursor.next();
             }
             parser::Step::Enter { .. } | parser::Step::Exit => (),
             parser::Step::Error { .. } => error = true,
@@ -409,7 +407,7 @@ pub fn expect_fragment(
         curr = curr.bump();
     }
 
-    *tt_iter = TtIter::new_iter(tt_iter.as_slice()[res.len()..].iter());
+    *cursor = TtIter::new_iter(cursor.as_slice()[res.len()..].iter());
     let res = match &*res {
         [] | [_] => res.pop(),
         [first, ..] => Some(tt::TokenTree::Subtree(tt::Subtree {