3 files changed, 56 insertions, 65 deletions

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/rustc_server.rs b/crates/proc_macro_srv/src/rustc_server.rs
index f02b0af60..2798dbf0d 100644
--- a/crates/proc_macro_srv/src/rustc_server.rs
+++ b/crates/proc_macro_srv/src/rustc_server.rs
@@ -25,36 +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 {
         if subtree.delimiter.is_some() {
-            TokenStream {
-                subtree: tt::Subtree {
-                    token_trees: vec![TokenTree::Subtree(subtree)],
-                    delimiter: None,
-                },
-            }
+            TokenStream { token_trees: vec![TokenTree::Subtree(subtree)] }
         } else {
-            TokenStream { subtree }
+            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] }
     }
 }
 
@@ -87,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);
                     }
                 }
             }
@@ -173,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()
         }
     }
 
@@ -200,32 +199,32 @@ pub mod token_stream {
 
     impl ToString for TokenStream {
         fn to_string(&self) -> String {
-            let tt = self.subtree.clone().into();
-            to_text(&tt)
+            tokentrees_to_text(&self.token_trees[..])
         }
     }
 
-    fn to_text(tkn: &tt::TokenTree) -> String {
+    fn tokentrees_to_text(tkns: &[tt::TokenTree]) -> String {
+        tkns.iter()
+            .fold((String::new(), true), |(last, last_to_joint), tkn| {
+                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 {
+                        is_joint = true;
+                    }
+                }
+                (s, is_joint)
+            })
+            .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 = subtree
-                    .token_trees
-                    .iter()
-                    .fold((String::new(), true), |(last, last_to_joint), tkn| {
-                        let s = [last, 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 {
-                                is_joint = true;
-                            }
-                        }
-                        (s, is_joint)
-                    })
-                    .0;
-
+                let content = tokentrees_to_text(&subtree.token_trees);
                 let (open, close) = match subtree.delimiter.map(|it| it.kind) {
                     None => ("", ""),
                     Some(tt::DelimiterKind::Brace) => ("{", "}"),
@@ -442,10 +441,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)
@@ -453,9 +449,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 {
@@ -764,26 +758,23 @@ 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 {}");
@@ -804,11 +795,11 @@ mod tests {
         });
 
         let t1 = TokenStream::from_str("(a)").unwrap();
-        assert_eq!(t1.subtree.token_trees.len(), 1);
-        assert_eq!(t1.subtree.token_trees[0], subtree_paren_a);
+        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.subtree.token_trees.len(), 2);
-        assert_eq!(t2.subtree.token_trees[0], subtree_paren_a);
+        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/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));
 }