From a9c4c6da4cf7f7907b0b0a00b896fcb4b128ec2d Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Mon, 22 Nov 2021 18:00:32 +0100
Subject: [PATCH] Fix mbe::Shift::new not accounting for non-ident token ids

---
 crates/hir_def/src/macro_expansion_tests/mbe.rs |  4 ++--
 crates/mbe/src/lib.rs                           | 13 ++++++++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/crates/hir_def/src/macro_expansion_tests/mbe.rs b/crates/hir_def/src/macro_expansion_tests/mbe.rs
index 962e1e7459..466c85fc5b 100644
--- a/crates/hir_def/src/macro_expansion_tests/mbe.rs
+++ b/crates/hir_def/src/macro_expansion_tests/mbe.rs
@@ -27,7 +27,7 @@ macro_rules! f {
 f!(struct MyTraitMap2);
 "#,
         expect![[r##"
-// call ids will be shifted by Shift(27)
+// call ids will be shifted by Shift(30)
 // +tokenids
 macro_rules! f {#0
     (#1 struct#2 $#3ident#4:#5ident#6 )#1 =#7>#8 {#9
@@ -39,7 +39,7 @@ macro_rules! f {#0
 
 // // +tokenids
 // f!(struct#1 MyTraitMap2#2);
-struct#10 MyTraitMap2#29 {#13
+struct#10 MyTraitMap2#32 {#13
     map#14:#15 ::std#18::collections#21::HashSet#24<#25(#26)#26>#27,#28
 }#13
 "##]],
diff --git a/crates/mbe/src/lib.rs b/crates/mbe/src/lib.rs
index 833a67ad99..b58b86b38d 100644
--- a/crates/mbe/src/lib.rs
+++ b/crates/mbe/src/lib.rs
@@ -120,12 +120,15 @@ impl Shift {
                             _ => tree_id,
                         }
                     }
-                    tt::TokenTree::Leaf(tt::Leaf::Ident(ident))
-                        if ident.id != tt::TokenId::unspecified() =>
-                    {
-                        Some(ident.id.0)
+                    tt::TokenTree::Leaf(leaf) => {
+                        let id = match leaf {
+                            tt::Leaf::Literal(it) => it.id,
+                            tt::Leaf::Punct(it) => it.id,
+                            tt::Leaf::Ident(it) => it.id,
+                        };
+
+                        (id != tt::TokenId::unspecified()).then(|| id.0)
                     }
-                    _ => None,
                 })
                 .max()
         }