1 files changed, 43 insertions, 0 deletions

diff --git a/crates/ra_mbe/src/tests.rs b/crates/ra_mbe/src/tests.rs
index 6bcfedcac..ae7f3dfd4 100644
--- a/crates/ra_mbe/src/tests.rs
+++ b/crates/ra_mbe/src/tests.rs
@@ -89,6 +89,32 @@ macro_rules! foobar {
 }
 
 #[test]
+fn test_token_map() {
+    use ra_parser::SyntaxKind::*;
+    use ra_syntax::T;
+
+    let macro_definition = r#"
+macro_rules! foobar {
+    ($e:ident) => { fn $e() {} }
+}
+"#;
+    let rules = create_rules(macro_definition);
+    let (expansion, (token_map, content)) = expand_and_map(&rules, "foobar!(baz);");
+
+    let get_text = |id, kind| -> String {
+        content[token_map.range_by_token(id).unwrap().range(kind).unwrap()].to_string()
+    };
+
+    assert_eq!(expansion.token_trees.len(), 4);
+    // {($e:ident) => { fn $e() {} }}
+    // 012345      67 8 9  T12  3
+
+    assert_eq!(get_text(tt::TokenId(9), IDENT), "fn");
+    assert_eq!(get_text(tt::TokenId(12), T!['(']), "(");
+    assert_eq!(get_text(tt::TokenId(13), T!['{']), "{");
+}
+
+#[test]
 fn test_convert_tt() {
     let macro_definition = r#"
 macro_rules! impl_froms {
@@ -1443,6 +1469,23 @@ pub(crate) fn expand(rules: &MacroRules, invocation: &str) -> tt::Subtree {
     rules.expand(&invocation_tt).unwrap()
 }
 
+pub(crate) fn expand_and_map(
+    rules: &MacroRules,
+    invocation: &str,
+) -> (tt::Subtree, (TokenMap, String)) {
+    let source_file = ast::SourceFile::parse(invocation).ok().unwrap();
+    let macro_invocation =
+        source_file.syntax().descendants().find_map(ast::MacroCall::cast).unwrap();
+
+    let (invocation_tt, _) = ast_to_token_tree(&macro_invocation.token_tree().unwrap()).unwrap();
+    let expanded = rules.expand(&invocation_tt).unwrap();
+
+    let (node, expanded_token_tree) =
+        token_tree_to_syntax_node(&expanded, FragmentKind::Items).unwrap();
+
+    (expanded, (expanded_token_tree, node.syntax_node().to_string()))
+}
+
 pub(crate) enum MacroKind {
     Items,
     Stmts,