7 files changed, 134 insertions, 103 deletions

diff --git a/crates/proc_macro_srv/Cargo.toml b/crates/proc_macro_srv/Cargo.toml
index 4c1b3036a..63b3f1532 100644
--- a/crates/proc_macro_srv/Cargo.toml
+++ b/crates/proc_macro_srv/Cargo.toml
@@ -17,9 +17,9 @@ memmap2 = "0.2.0"
 tt = { path = "../tt", version = "0.0.0" }
 mbe = { path = "../mbe", version = "0.0.0" }
 proc_macro_api = { path = "../proc_macro_api", version = "0.0.0" }
-test_utils = { path = "../test_utils", version = "0.0.0" }
 
 [dev-dependencies]
+test_utils = { path = "../test_utils" }
 cargo_metadata = "0.13"
 
 # used as proc macro test targets
diff --git a/crates/proc_macro_srv/src/dylib.rs b/crates/proc_macro_srv/src/dylib.rs
index 28a6ee547..baf10fea9 100644
--- a/crates/proc_macro_srv/src/dylib.rs
+++ b/crates/proc_macro_srv/src/dylib.rs
@@ -138,7 +138,7 @@ impl Expander {
                         parsed_body,
                         false,
                     );
-                    return res.map(|it| it.subtree);
+                    return res.map(|it| it.into_subtree());
                 }
                 bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                     let res = client.run(
@@ -147,7 +147,7 @@ impl Expander {
                         parsed_body,
                         false,
                     );
-                    return res.map(|it| it.subtree);
+                    return res.map(|it| it.into_subtree());
                 }
                 bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                     let res = client.run(
@@ -157,7 +157,7 @@ impl Expander {
                         parsed_body,
                         false,
                     );
-                    return res.map(|it| it.subtree);
+                    return res.map(|it| it.into_subtree());
                 }
                 _ => continue,
             }
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/client.rs b/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
index ca6749b9b..b036d4e20 100644
--- a/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
+++ b/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
@@ -238,7 +238,7 @@ macro_rules! define_client_side {
         $(impl $name {
             #[allow(unused)]
             $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)* {
-                panic!("hello");
+                panic!("crates should be linked against the sysroot version of proc_macro, not this one from rust-analyzer");
                 // Bridge::with(|bridge| {
                 //     let mut b = bridge.cached_buffer.take();
 
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs b/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
index e4006a5ab..e67902682 100644
--- a/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
+++ b/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
@@ -268,7 +268,7 @@ trait Mark {
     fn mark(unmarked: Self::Unmarked) -> Self;
 }
 
-/// Unwrap types wrapped by `Mark::mark` (see `Mark` for details).
+/// Unwrap types wrapped by `cov_mark::mark` (see `Mark` for details).
 trait Unmark {
     type Unmarked;
     fn unmark(self) -> Self::Unmarked;
diff --git a/crates/proc_macro_srv/src/rustc_server.rs b/crates/proc_macro_srv/src/rustc_server.rs
index 952b4a97f..ceefd187d 100644
--- a/crates/proc_macro_srv/src/rustc_server.rs
+++ b/crates/proc_macro_srv/src/rustc_server.rs
@@ -25,27 +25,35 @@ type Span = tt::TokenId;
 
 #[derive(Debug, Clone)]
 pub struct TokenStream {
-    pub subtree: tt::Subtree,
+    pub token_trees: Vec<TokenTree>,
 }
 
 impl TokenStream {
     pub fn new() -> Self {
-        TokenStream { subtree: Default::default() }
+        TokenStream { token_trees: Default::default() }
     }
 
     pub fn with_subtree(subtree: tt::Subtree) -> Self {
-        TokenStream { subtree }
+        if subtree.delimiter.is_some() {
+            TokenStream { token_trees: vec![TokenTree::Subtree(subtree)] }
+        } else {
+            TokenStream { token_trees: subtree.token_trees }
+        }
+    }
+
+    pub fn into_subtree(self) -> tt::Subtree {
+        tt::Subtree { delimiter: None, token_trees: self.token_trees }
     }
 
     pub fn is_empty(&self) -> bool {
-        self.subtree.token_trees.is_empty()
+        self.token_trees.is_empty()
     }
 }
 
 /// Creates a token stream containing a single token tree.
 impl From<TokenTree> for TokenStream {
     fn from(tree: TokenTree) -> TokenStream {
-        TokenStream { subtree: tt::Subtree { delimiter: None, token_trees: vec![tree] } }
+        TokenStream { token_trees: vec![tree] }
     }
 }
 
@@ -78,10 +86,10 @@ impl Extend<TokenStream> for TokenStream {
             for tkn in item {
                 match tkn {
                     tt::TokenTree::Subtree(subtree) if subtree.delimiter.is_none() => {
-                        self.subtree.token_trees.extend(subtree.token_trees);
+                        self.token_trees.extend(subtree.token_trees);
                     }
                     _ => {
-                        self.subtree.token_trees.push(tkn);
+                        self.token_trees.push(tkn);
                     }
                 }
             }
@@ -164,7 +172,7 @@ pub mod token_stream {
         type IntoIter = super::IntoIter<TokenTree>;
 
         fn into_iter(self) -> Self::IntoIter {
-            self.subtree.token_trees.into_iter()
+            self.token_trees.into_iter()
         }
     }
 
@@ -184,28 +192,23 @@ pub mod token_stream {
             let (subtree, _token_map) =
                 mbe::parse_to_token_tree(src).ok_or("Failed to parse from mbe")?;
 
-            Ok(TokenStream { subtree })
+            let subtree = subtree_replace_token_ids_with_unspecified(subtree);
+            Ok(TokenStream::with_subtree(subtree))
         }
     }
 
     impl ToString for TokenStream {
         fn to_string(&self) -> String {
-            let tt = self.subtree.clone().into();
-            to_text(&tt)
-        }
-    }
+            return tokentrees_to_text(&self.token_trees[..]);
 
-    fn to_text(tkn: &tt::TokenTree) -> String {
-        match tkn {
-            tt::TokenTree::Leaf(tt::Leaf::Ident(ident)) => ident.text.clone().into(),
-            tt::TokenTree::Leaf(tt::Leaf::Literal(literal)) => literal.text.clone().into(),
-            tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) => format!("{}", punct.char),
-            tt::TokenTree::Subtree(subtree) => {
-                let content = subtree
-                    .token_trees
-                    .iter()
+            fn tokentrees_to_text(tkns: &[tt::TokenTree]) -> String {
+                tkns.iter()
                     .fold((String::new(), true), |(last, last_to_joint), tkn| {
-                        let s = [last, to_text(tkn)].join(if last_to_joint { "" } else { " " });
+                        let s = [last, tokentree_to_text(tkn)].join(if last_to_joint {
+                            ""
+                        } else {
+                            " "
+                        });
                         let mut is_joint = false;
                         if let tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) = tkn {
                             if punct.spacing == tt::Spacing::Joint {
@@ -214,15 +217,63 @@ pub mod token_stream {
                         }
                         (s, is_joint)
                     })
-                    .0;
-
-                let (open, close) = match subtree.delimiter.map(|it| it.kind) {
-                    None => ("", ""),
-                    Some(tt::DelimiterKind::Brace) => ("{", "}"),
-                    Some(tt::DelimiterKind::Parenthesis) => ("(", ")"),
-                    Some(tt::DelimiterKind::Bracket) => ("[", "]"),
-                };
-                format!("{}{}{}", open, content, close)
+                    .0
+            }
+
+            fn tokentree_to_text(tkn: &tt::TokenTree) -> String {
+                match tkn {
+                    tt::TokenTree::Leaf(tt::Leaf::Ident(ident)) => ident.text.clone().into(),
+                    tt::TokenTree::Leaf(tt::Leaf::Literal(literal)) => literal.text.clone().into(),
+                    tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) => format!("{}", punct.char),
+                    tt::TokenTree::Subtree(subtree) => {
+                        let content = tokentrees_to_text(&subtree.token_trees);
+                        let (open, close) = match subtree.delimiter.map(|it| it.kind) {
+                            None => ("", ""),
+                            Some(tt::DelimiterKind::Brace) => ("{", "}"),
+                            Some(tt::DelimiterKind::Parenthesis) => ("(", ")"),
+                            Some(tt::DelimiterKind::Bracket) => ("[", "]"),
+                        };
+                        format!("{}{}{}", open, content, close)
+                    }
+                }
+            }
+        }
+    }
+
+    fn subtree_replace_token_ids_with_unspecified(subtree: tt::Subtree) -> tt::Subtree {
+        tt::Subtree {
+            delimiter: subtree
+                .delimiter
+                .map(|d| tt::Delimiter { id: tt::TokenId::unspecified(), ..d }),
+            token_trees: subtree
+                .token_trees
+                .into_iter()
+                .map(|t| token_tree_replace_token_ids_with_unspecified(t))
+                .collect(),
+        }
+    }
+
+    fn token_tree_replace_token_ids_with_unspecified(tt: tt::TokenTree) -> tt::TokenTree {
+        match tt {
+            tt::TokenTree::Leaf(leaf) => {
+                tt::TokenTree::Leaf(leaf_replace_token_ids_with_unspecified(leaf))
+            }
+            tt::TokenTree::Subtree(subtree) => {
+                tt::TokenTree::Subtree(subtree_replace_token_ids_with_unspecified(subtree))
+            }
+        }
+    }
+
+    fn leaf_replace_token_ids_with_unspecified(leaf: tt::Leaf) -> tt::Leaf {
+        match leaf {
+            tt::Leaf::Literal(lit) => {
+                tt::Leaf::Literal(tt::Literal { id: tt::TokenId::unspecified(), ..lit })
+            }
+            tt::Leaf::Punct(punct) => {
+                tt::Leaf::Punct(tt::Punct { id: tt::TokenId::unspecified(), ..punct })
+            }
+            tt::Leaf::Ident(ident) => {
+                tt::Leaf::Ident(tt::Ident { id: tt::TokenId::unspecified(), ..ident })
             }
         }
     }
@@ -277,42 +328,6 @@ impl server::FreeFunctions for Rustc {
     }
 }
 
-fn subtree_replace_token_ids_with_unspecified(subtree: tt::Subtree) -> tt::Subtree {
-    tt::Subtree {
-        delimiter: subtree.delimiter.map(|d| tt::Delimiter { id: tt::TokenId::unspecified(), ..d }),
-        token_trees: subtree
-            .token_trees
-            .into_iter()
-            .map(|t| token_tree_replace_token_ids_with_unspecified(t))
-            .collect(),
-    }
-}
-
-fn token_tree_replace_token_ids_with_unspecified(tt: tt::TokenTree) -> tt::TokenTree {
-    match tt {
-        tt::TokenTree::Leaf(leaf) => {
-            tt::TokenTree::Leaf(leaf_replace_token_ids_with_unspecified(leaf))
-        }
-        tt::TokenTree::Subtree(subtree) => {
-            tt::TokenTree::Subtree(subtree_replace_token_ids_with_unspecified(subtree))
-        }
-    }
-}
-
-fn leaf_replace_token_ids_with_unspecified(leaf: tt::Leaf) -> tt::Leaf {
-    match leaf {
-        tt::Leaf::Literal(lit) => {
-            tt::Leaf::Literal(tt::Literal { id: tt::TokenId::unspecified(), ..lit })
-        }
-        tt::Leaf::Punct(punct) => {
-            tt::Leaf::Punct(tt::Punct { id: tt::TokenId::unspecified(), ..punct })
-        }
-        tt::Leaf::Ident(ident) => {
-            tt::Leaf::Ident(tt::Ident { id: tt::TokenId::unspecified(), ..ident })
-        }
-    }
-}
-
 impl server::TokenStream for Rustc {
     fn new(&mut self) -> Self::TokenStream {
         Self::TokenStream::new()
@@ -322,8 +337,9 @@ impl server::TokenStream for Rustc {
         stream.is_empty()
     }
     fn from_str(&mut self, src: &str) -> Self::TokenStream {
-        let (subtree, _) = mbe::parse_to_token_tree(src).expect("cannot parse string");
-        TokenStream::with_subtree(subtree_replace_token_ids_with_unspecified(subtree))
+        use std::str::FromStr;
+
+        Self::TokenStream::from_str(src).expect("cannot parse string")
     }
     fn to_string(&mut self, stream: &Self::TokenStream) -> String {
         stream.to_string()
@@ -429,10 +445,7 @@ fn spacing_to_external(spacing: Spacing) -> bridge::Spacing {
 
 impl server::Group for Rustc {
     fn new(&mut self, delimiter: bridge::Delimiter, stream: Self::TokenStream) -> Self::Group {
-        Self::Group {
-            delimiter: delim_to_internal(delimiter),
-            token_trees: stream.subtree.token_trees,
-        }
+        Self::Group { delimiter: delim_to_internal(delimiter), token_trees: stream.token_trees }
     }
     fn delimiter(&mut self, group: &Self::Group) -> bridge::Delimiter {
         delim_to_external(group.delimiter)
@@ -440,9 +453,7 @@ impl server::Group for Rustc {
 
     // NOTE: Return value of do not include delimiter
     fn stream(&mut self, group: &Self::Group) -> Self::TokenStream {
-        TokenStream {
-            subtree: tt::Subtree { delimiter: None, token_trees: group.token_trees.clone() },
-        }
+        TokenStream { token_trees: group.token_trees.clone() }
     }
 
     fn span(&mut self, group: &Self::Group) -> Self::Span {
@@ -751,28 +762,48 @@ mod tests {
     #[test]
     fn test_rustc_server_to_string() {
         let s = TokenStream {
-            subtree: tt::Subtree {
-                delimiter: None,
-                token_trees: vec![
-                    tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
-                        text: "struct".into(),
-                        id: tt::TokenId::unspecified(),
-                    })),
-                    tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
-                        text: "T".into(),
+            token_trees: vec![
+                tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
+                    text: "struct".into(),
+                    id: tt::TokenId::unspecified(),
+                })),
+                tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
+                    text: "T".into(),
+                    id: tt::TokenId::unspecified(),
+                })),
+                tt::TokenTree::Subtree(tt::Subtree {
+                    delimiter: Some(tt::Delimiter {
                         id: tt::TokenId::unspecified(),
-                    })),
-                    tt::TokenTree::Subtree(tt::Subtree {
-                        delimiter: Some(tt::Delimiter {
-                            id: tt::TokenId::unspecified(),
-                            kind: tt::DelimiterKind::Brace,
-                        }),
-                        token_trees: vec![],
+                        kind: tt::DelimiterKind::Brace,
                     }),
-                ],
-            },
+                    token_trees: vec![],
+                }),
+            ],
         };
 
         assert_eq!(s.to_string(), "struct T {}");
     }
+
+    #[test]
+    fn test_rustc_server_from_str() {
+        use std::str::FromStr;
+        let subtree_paren_a = tt::TokenTree::Subtree(tt::Subtree {
+            delimiter: Some(tt::Delimiter {
+                id: tt::TokenId::unspecified(),
+                kind: tt::DelimiterKind::Parenthesis,
+            }),
+            token_trees: vec![tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
+                text: "a".into(),
+                id: tt::TokenId::unspecified(),
+            }))],
+        });
+
+        let t1 = TokenStream::from_str("(a)").unwrap();
+        assert_eq!(t1.token_trees.len(), 1);
+        assert_eq!(t1.token_trees[0], subtree_paren_a);
+
+        let t2 = TokenStream::from_str("(a);").unwrap();
+        assert_eq!(t2.token_trees.len(), 2);
+        assert_eq!(t2.token_trees[0], subtree_paren_a);
+    }
 }
diff --git a/crates/proc_macro_srv/src/tests/fixtures/test_serialize_proc_macro.txt b/crates/proc_macro_srv/src/tests/fixtures/test_serialize_proc_macro.txt
index ea34e688f..fa581f110 100644
--- a/crates/proc_macro_srv/src/tests/fixtures/test_serialize_proc_macro.txt
+++ b/crates/proc_macro_srv/src/tests/fixtures/test_serialize_proc_macro.txt
@@ -101,7 +101,7 @@ SUBTREE $
     PUNCH   : [alone] 4294967295
     IDENT   Serialize 4294967295
     IDENT   for 4294967295
-    IDENT   Foo 1
+    IDENT   Foo 4294967295
     SUBTREE {} 4294967295
       IDENT   fn 4294967295
       IDENT   serialize 4294967295
diff --git a/crates/proc_macro_srv/src/tests/utils.rs b/crates/proc_macro_srv/src/tests/utils.rs
index 22813052d..0484c3af4 100644
--- a/crates/proc_macro_srv/src/tests/utils.rs
+++ b/crates/proc_macro_srv/src/tests/utils.rs
@@ -52,7 +52,7 @@ pub fn assert_expand(
     let expander = dylib::Expander::new(&path).unwrap();
     let fixture = parse_string(ra_fixture).unwrap();
 
-    let res = expander.expand(macro_name, &fixture.subtree, None).unwrap();
+    let res = expander.expand(macro_name, &fixture.into_subtree(), None).unwrap();
     assert_eq_text!(&expect.trim(), &format!("{:?}", res));
 }