mbe: split Op::Leaf and handle multi-character puncts

Author: lowr

crates/hir-def/src/macro_expansion_tests/mbe.rs

@@ -1630,3 +1630,48 @@ const _: i32 = -0--1--2;
     "#]],
     );
 }
+
+#[test]
+fn test_punct_without_space() {
+    // Puncts are "glued" greedily.
+    check(
+        r#"
+macro_rules! foo {
+    (: : :) => { "1 1 1" };
+    (: ::) => { "1 2" };
+    (:: :) => { "2 1" };
+
+    (: : : :) => { "1 1 1 1" };
+    (:: : :) => { "2 1 1" };
+    (: :: :) => { "1 2 1" };
+    (: : ::) => { "1 1 2" };
+    (:: ::) => { "2 2" };
+}
+
+fn test() {
+    foo!(:::);
+    foo!(: :::);
+    foo!(::::);
+}
+"#,
+        expect![[r#"
+macro_rules! foo {
+    (: : :) => { "1 1 1" };
+    (: ::) => { "1 2" };
+    (:: :) => { "2 1" };
+
+    (: : : :) => { "1 1 1 1" };
+    (:: : :) => { "2 1 1" };
+    (: :: :) => { "1 2 1" };
+    (: : ::) => { "1 1 2" };
+    (:: ::) => { "2 2" };
+}
+
+fn test() {
+    "2 1";
+    "1 2 1";
+    "2 2";
+}
+"#]],
+    );
+}

crates/mbe/src/benchmark.rs

@@ -141,7 +141,13 @@ fn invocation_fixtures(rules: &FxHashMap<String, DeclarativeMacro>) -> Vec<(Stri
                 None => (),
                 Some(kind) => panic!("Unhandled kind {kind:?}"),
             },
-            Op::Leaf(leaf) => parent.token_trees.push(leaf.clone().into()),
+            Op::Literal(it) => parent.token_trees.push(tt::Leaf::from(it.clone()).into()),
+            Op::Ident(it) => parent.token_trees.push(tt::Leaf::from(it.clone()).into()),
+            Op::Punct(puncts) => {
+                for punct in puncts {
+                    parent.token_trees.push(tt::Leaf::from(punct.clone()).into());
+                }
+            }
             Op::Repeat { tokens, kind, separator } => {
                 let max = 10;
                 let cnt = match kind {

crates/mbe/src/expander/matcher.rs

@@ -68,7 +68,7 @@ use crate::{
     expander::{Binding, Bindings, ExpandResult, Fragment},
     parser::{MetaVarKind, Op, RepeatKind, Separator},
     tt_iter::TtIter,
-    ExpandError, MetaTemplate,
+    ExpandError, MetaTemplate, ValueResult,
 };
 
 impl Bindings {
@@ -500,18 +500,69 @@ fn match_loop_inner<'t>(
                     }
                 }
             }
-            OpDelimited::Op(Op::Leaf(leaf)) => {
-                if let Err(err) = match_leaf(leaf, &mut src.clone()) {
-                    res.add_err(err);
+            OpDelimited::Op(Op::Literal(lhs)) => {
+                if let Ok(rhs) = src.clone().expect_leaf() {
+                    if matches!(rhs, tt::Leaf::Literal(it) if it.text == lhs.text) {
+                        item.dot.next();
+                    } else {
+                        res.add_err(ExpandError::UnexpectedToken);
+                        item.is_error = true;
+                    }
+                } else {
+                    res.add_err(ExpandError::binding_error(format!("expected literal: `{lhs}`")));
                     item.is_error = true;
+                }
+                try_push!(next_items, item);
+            }
+            OpDelimited::Op(Op::Ident(lhs)) => {
+                if let Ok(rhs) = src.clone().expect_leaf() {
+                    if matches!(rhs, tt::Leaf::Ident(it) if it.text == lhs.text) {
+                        item.dot.next();
+                    } else {
+                        res.add_err(ExpandError::UnexpectedToken);
+                        item.is_error = true;
+                    }
                 } else {
-                    item.dot.next();
+                    res.add_err(ExpandError::binding_error(format!("expected ident: `{lhs}`")));
+                    item.is_error = true;
                 }
                 try_push!(next_items, item);
             }
+            OpDelimited::Op(Op::Punct(lhs)) => {
+                let mut fork = src.clone();
+                let error = if let Ok(rhs) = fork.expect_glued_punct() {
+                    let first_is_single_quote = rhs[0].char == '\'';
+                    let lhs = lhs.iter().map(|it| it.char);
+                    let rhs = rhs.iter().map(|it| it.char);
+                    if lhs.clone().eq(rhs) {
+                        // HACK: here we use `meta_result` to pass `TtIter` back to caller because
+                        // it might have been advanced multiple times. `ValueResult` is
+                        // insignificant.
+                        item.meta_result = Some((fork, ValueResult::ok(None)));
+                        item.dot.next();
+                        next_items.push(item);
+                        continue;
+                    }
+
+                    if first_is_single_quote {
+                        // If the first punct token is a single quote, that's a part of a lifetime
+                        // ident, not a punct.
+                        ExpandError::UnexpectedToken
+                    } else {
+                        let lhs: SmolStr = lhs.collect();
+                        ExpandError::binding_error(format!("expected punct: `{lhs}`"))
+                    }
+                } else {
+                    ExpandError::UnexpectedToken
+                };
+
+                res.add_err(error);
+                item.is_error = true;
+                error_items.push(item);
+            }
             OpDelimited::Op(Op::Ignore { .. } | Op::Index { .. }) => {}
             OpDelimited::Open => {
-                if matches!(src.clone().next(), Some(tt::TokenTree::Subtree(..))) {
+                if matches!(src.peek_n(0), Some(tt::TokenTree::Subtree(..))) {
                     item.dot.next();
                     try_push!(next_items, item);
                 }
@@ -616,21 +667,33 @@ fn match_loop(pattern: &MetaTemplate, src: &tt::Subtree) -> Match {
         }
         // Dump all possible `next_items` into `cur_items` for the next iteration.
         else if !next_items.is_empty() {
-            // Now process the next token
-            cur_items.extend(next_items.drain(..));
-
-            match src.next() {
-                Some(tt::TokenTree::Subtree(subtree)) => {
-                    stack.push(src.clone());
-                    src = TtIter::new(subtree);
+            if let Some((iter, _)) = next_items[0].meta_result.take() {
+                // We've matched a possibly "glued" punct. The matched punct (hence
+                // `meta_result` also) must be the same for all items.
+                // FIXME: If there are multiple items, it's definitely redundant (and it's hacky!
+                // `meta_result` isn't supposed to be used this way).
+
+                // We already bumped, so no need to call `.next()` like in the other branch.
+                src = iter;
+                for item in next_items.iter_mut() {
+                    item.meta_result = None;
                 }
-                None => {
-                    if let Some(iter) = stack.pop() {
-                        src = iter;
+            } else {
+                match src.next() {
+                    Some(tt::TokenTree::Subtree(subtree)) => {
+                        stack.push(src.clone());
+                        src = TtIter::new(subtree);
+                    }
+                    None => {
+                        if let Some(iter) = stack.pop() {
+                            src = iter;
+                        }
                     }
+                    _ => (),
                 }
-                _ => (),
             }
+            // Now process the next token
+            cur_items.extend(next_items.drain(..));
         }
         // Finally, we have the case where we need to call the black-box parser to get some
         // nonterminal.
@@ -663,27 +726,6 @@ fn match_loop(pattern: &MetaTemplate, src: &tt::Subtree) -> Match {
     }
 }
 
-fn match_leaf(lhs: &tt::Leaf, src: &mut TtIter<'_>) -> Result<(), ExpandError> {
-    let rhs = src
-        .expect_leaf()
-        .map_err(|()| ExpandError::binding_error(format!("expected leaf: `{lhs}`")))?;
-    match (lhs, rhs) {
-        (
-            tt::Leaf::Punct(tt::Punct { char: lhs, .. }),
-            tt::Leaf::Punct(tt::Punct { char: rhs, .. }),
-        ) if lhs == rhs => Ok(()),
-        (
-            tt::Leaf::Ident(tt::Ident { text: lhs, .. }),
-            tt::Leaf::Ident(tt::Ident { text: rhs, .. }),
-        ) if lhs == rhs => Ok(()),
-        (
-            tt::Leaf::Literal(tt::Literal { text: lhs, .. }),
-            tt::Leaf::Literal(tt::Literal { text: rhs, .. }),
-        ) if lhs == rhs => Ok(()),
-        _ => Err(ExpandError::UnexpectedToken),
-    }
-}
-
 fn match_meta_var(kind: MetaVarKind, input: &mut TtIter<'_>) -> ExpandResult<Option<Fragment>> {
     let fragment = match kind {
         MetaVarKind::Path => parser::PrefixEntryPoint::Path,
@@ -756,10 +798,10 @@ fn collect_vars(collector_fun: &mut impl FnMut(SmolStr), pattern: &MetaTemplate)
     for op in pattern.iter() {
         match op {
             Op::Var { name, .. } => collector_fun(name.clone()),
-            Op::Leaf(_) => (),
             Op::Subtree { tokens, .. } => collect_vars(collector_fun, tokens),
             Op::Repeat { tokens, .. } => collect_vars(collector_fun, tokens),
-            Op::Ignore { .. } | Op::Index { .. } => {}
+            Op::Ignore { .. } | Op::Index { .. } | Op::Literal(_) | Op::Ident(_) | Op::Punct(_) => {
+            }
         }
     }
 }

crates/mbe/src/expander/transcriber.rs

@@ -134,7 +134,13 @@ fn expand_subtree(
     let mut err = None;
     for op in template.iter() {
         match op {
-            Op::Leaf(tt) => arena.push(tt.clone().into()),
+            Op::Literal(it) => arena.push(tt::Leaf::from(it.clone()).into()),
+            Op::Ident(it) => arena.push(tt::Leaf::from(it.clone()).into()),
+            Op::Punct(puncts) => {
+                for punct in puncts {
+                    arena.push(tt::Leaf::from(punct.clone()).into());
+                }
+            }
             Op::Subtree { tokens, delimiter } => {
                 let ExpandResult { value: tt, err: e } =
                     expand_subtree(ctx, tokens, *delimiter, arena);

crates/mbe/src/parser.rs

@@ -1,7 +1,7 @@
 //! Parser recognizes special macro syntax, `$var` and `$(repeat)*`, in token
 //! trees.
 
-use smallvec::SmallVec;
+use smallvec::{smallvec, SmallVec};
 use syntax::SmolStr;
 
 use crate::{tt_iter::TtIter, ParseError};
@@ -39,7 +39,7 @@ impl MetaTemplate {
         let mut src = TtIter::new(tt);
 
         let mut res = Vec::new();
-        while let Some(first) = src.next() {
+        while let Some(first) = src.peek_n(0) {
             let op = next_op(first, &mut src, mode)?;
             res.push(op);
         }
@@ -54,8 +54,10 @@ pub(crate) enum Op {
     Ignore { name: SmolStr, id: tt::TokenId },
     Index { depth: u32 },
     Repeat { tokens: MetaTemplate, kind: RepeatKind, separator: Option<Separator> },
-    Leaf(tt::Leaf),
     Subtree { tokens: MetaTemplate, delimiter: Option<tt::Delimiter> },
+    Literal(tt::Literal),
+    Punct(SmallVec<[tt::Punct; 3]>),
+    Ident(tt::Ident),
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -124,12 +126,17 @@ enum Mode {
     Template,
 }
 
-fn next_op<'a>(first: &tt::TokenTree, src: &mut TtIter<'a>, mode: Mode) -> Result<Op, ParseError> {
-    let res = match first {
-        tt::TokenTree::Leaf(leaf @ tt::Leaf::Punct(tt::Punct { char: '$', .. })) => {
+fn next_op<'a>(
+    first_peeked: &tt::TokenTree,
+    src: &mut TtIter<'a>,
+    mode: Mode,
+) -> Result<Op, ParseError> {
+    let res = match first_peeked {
+        tt::TokenTree::Leaf(tt::Leaf::Punct(p @ tt::Punct { char: '$', .. })) => {
+            src.next().expect("first token already peeked");
             // Note that the '$' itself is a valid token inside macro_rules.
             let second = match src.next() {
-                None => return Ok(Op::Leaf(leaf.clone())),
+                None => return Ok(Op::Punct(smallvec![p.clone()])),
                 Some(it) => it,
             };
             match second {
@@ -160,7 +167,7 @@ fn next_op<'a>(first: &tt::TokenTree, src: &mut TtIter<'a>, mode: Mode) -> Resul
                 tt::TokenTree::Leaf(leaf) => match leaf {
                     tt::Leaf::Ident(ident) if ident.text == "crate" => {
                         // We simply produce identifier `$crate` here. And it will be resolved when lowering ast to Path.
-                        Op::Leaf(tt::Leaf::from(tt::Ident { text: "$crate".into(), id: ident.id }))
+                        Op::Ident(tt::Ident { text: "$crate".into(), id: ident.id })
                     }
                     tt::Leaf::Ident(ident) => {
                         let kind = eat_fragment_kind(src, mode)?;
@@ -180,16 +187,33 @@ fn next_op<'a>(first: &tt::TokenTree, src: &mut TtIter<'a>, mode: Mode) -> Resul
                                 "`$$` is not allowed on the pattern side",
                             ))
                         }
-                        Mode::Template => Op::Leaf(tt::Leaf::Punct(*punct)),
+                        Mode::Template => Op::Punct(smallvec![*punct]),
                     },
                     tt::Leaf::Punct(_) | tt::Leaf::Literal(_) => {
                         return Err(ParseError::expected("expected ident"))
                     }
                 },
             }
         }
-        tt::TokenTree::Leaf(tt) => Op::Leaf(tt.clone()),
+
+        tt::TokenTree::Leaf(tt::Leaf::Literal(it)) => {
+            src.next().expect("first token already peeked");
+            Op::Literal(it.clone())
+        }
+
+        tt::TokenTree::Leaf(tt::Leaf::Ident(it)) => {
+            src.next().expect("first token already peeked");
+            Op::Ident(it.clone())
+        }
+
+        tt::TokenTree::Leaf(tt::Leaf::Punct(_)) => {
+            // There's at least one punct so this shouldn't fail.
+            let puncts = src.expect_glued_punct().unwrap();
+            Op::Punct(puncts)
+        }
+
         tt::TokenTree::Subtree(subtree) => {
+            src.next().expect("first token already peeked");
             let tokens = MetaTemplate::parse(subtree, mode)?;
             Op::Subtree { tokens, delimiter: subtree.delimiter }
         }

crates/mbe/src/tt_iter.rs

@@ -88,6 +88,10 @@ impl<'a> TtIter<'a> {
         }
     }
 
+    /// Returns consecutive `Punct`s that can be glued together.
+    ///
+    /// This method currently may return a single quotation, which is part of lifetime ident and
+    /// conceptually not a punct in the context of mbe. Callers should handle this.
     pub(crate) fn expect_glued_punct(&mut self) -> Result<SmallVec<[tt::Punct; 3]>, ()> {
         let tt::TokenTree::Leaf(tt::Leaf::Punct(first)) = self.next().ok_or(())?.clone() else {
             return Err(());
@@ -182,7 +186,7 @@ impl<'a> TtIter<'a> {
         ExpandResult { value: res, err }
     }
 
-    pub(crate) fn peek_n(&self, n: usize) -> Option<&tt::TokenTree> {
+    pub(crate) fn peek_n(&self, n: usize) -> Option<&'a tt::TokenTree> {
         self.inner.as_slice().get(n)
     }
 }