2 files changed, 6 insertions, 7 deletions

diff --git a/crates/mbe/src/tests/expand.rs b/crates/mbe/src/tests/expand.rs
index 3a1d840ea..5f173f513 100644
--- a/crates/mbe/src/tests/expand.rs
+++ b/crates/mbe/src/tests/expand.rs
@@ -58,9 +58,8 @@ macro_rules! foobar {
     let (node, token_map) = token_tree_to_syntax_node(&expanded, FragmentKind::Items).unwrap();
     let content = node.syntax_node().to_string();
 
-    let get_text = |id, kind| -> String {
-        content[token_map.range_by_token(id).unwrap().by_kind(kind).unwrap()].to_string()
-    };
+    let get_text =
+        |id, kind| -> String { content[token_map.range_by_token(id, kind).unwrap()].to_string() };
 
     assert_eq!(expanded.token_trees.len(), 4);
     // {($e:ident) => { fn $e() {} }}
diff --git a/crates/mbe/src/token_map.rs b/crates/mbe/src/token_map.rs
index 58c9f5aa5..0567475be 100644
--- a/crates/mbe/src/token_map.rs
+++ b/crates/mbe/src/token_map.rs
@@ -2,13 +2,13 @@ use parser::{SyntaxKind, T};
 use syntax::{TextRange, TextSize};
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
-pub enum TokenTextRange {
+enum TokenTextRange {
     Token(TextRange),
     Delimiter(TextRange),
 }
 
 impl TokenTextRange {
-    pub fn by_kind(self, kind: SyntaxKind) -> Option<TextRange> {
+    fn by_kind(self, kind: SyntaxKind) -> Option<TextRange> {
         match self {
             TokenTextRange::Token(it) => Some(it),
             TokenTextRange::Delimiter(it) => match kind {
@@ -42,9 +42,9 @@ impl TokenMap {
         Some(token_id)
     }
 
-    pub fn range_by_token(&self, token_id: tt::TokenId) -> Option<TokenTextRange> {
+    pub fn range_by_token(&self, token_id: tt::TokenId, kind: SyntaxKind) -> Option<TextRange> {
         let &(_, range) = self.entries.iter().find(|(tid, _)| *tid == token_id)?;
-        Some(range)
+        range.by_kind(kind)
     }
 
     pub(crate) fn shrink_to_fit(&mut self) {