1 files changed, 84 insertions, 57 deletions

diff --git a/crates/proc_macro_srv/src/rustc_server.rs b/crates/proc_macro_srv/src/rustc_server.rs
index 14c853c77..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()
         }
     }
 
@@ -185,28 +193,22 @@ pub mod token_stream {
                 mbe::parse_to_token_tree(src).ok_or("Failed to parse from mbe")?;
 
             let subtree = subtree_replace_token_ids_with_unspecified(subtree);
-            Ok(TokenStream { 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 {
@@ -215,15 +217,25 @@ 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)
+                    }
+                }
             }
         }
     }
@@ -433,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)
@@ -444,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 {
@@ -755,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);
+    }
 }