Return all ranges corresponding to a token id in TokenMap

crates/mbe/src/tests/expand.rs

@@ -58,8 +58,9 @@ 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, kind).unwrap()].to_string() };
+    let get_text = |id, kind| -> String {
+        content[token_map.first_range_by_token(id, kind).unwrap()].to_string()
+    };
 
     assert_eq!(expanded.token_trees.len(), 4);
     // {($e:ident) => { fn $e() {} }}

crates/mbe/src/token_map.rs

@@ -46,9 +46,23 @@ impl TokenMap {
         Some(token_id)
     }
 
-    pub fn range_by_token(&self, token_id: tt::TokenId, kind: SyntaxKind) -> Option<TextRange> {
-        let &(_, range) = self.entries.iter().find(|(tid, _)| *tid == token_id)?;
-        range.by_kind(kind)
+    pub fn ranges_by_token(
+        &self,
+        token_id: tt::TokenId,
+        kind: SyntaxKind,
+    ) -> impl Iterator<Item = TextRange> + '_ {
+        self.entries
+            .iter()
+            .filter(move |&&(tid, _)| tid == token_id)
+            .filter_map(move |(_, range)| range.by_kind(kind))
+    }
+
+    pub fn first_range_by_token(
+        &self,
+        token_id: tt::TokenId,
+        kind: SyntaxKind,
+    ) -> Option<TextRange> {
+        self.ranges_by_token(token_id, kind).next()
     }
 
     pub(crate) fn shrink_to_fit(&mut self) {