Refactor proc_macro::TokenStream to use syntax::tokenstream::TokenStream; fix tests for changed semantics

Author: abonander

src/libproc_macro/lib.rs

@@ -39,9 +39,9 @@ extern crate syntax;
 use std::fmt;
 use std::str::FromStr;
 
-use syntax::ast;
+use syntax::errors::DiagnosticBuilder;
 use syntax::parse;
-use syntax::ptr::P;
+use syntax::tokenstream::TokenStream as TokenStream_;
 
 /// The main type provided by this crate, representing an abstract stream of
 /// tokens.
@@ -54,7 +54,7 @@ use syntax::ptr::P;
 /// time!
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 pub struct TokenStream {
-    inner: Vec<P<ast::Item>>,
+    inner: TokenStream_,
 }
 
 /// Error returned from `TokenStream::from_str`.
@@ -77,17 +77,41 @@ pub struct LexError {
 #[doc(hidden)]
 pub mod __internal {
     use std::cell::Cell;
+    use std::rc::Rc;
 
     use syntax::ast;
     use syntax::ptr::P;
-    use syntax::parse::ParseSess;
-    use super::TokenStream;
+    use syntax::parse::{self, token, ParseSess};
+    use syntax::tokenstream::TokenStream as TokenStream_;
+
+    use super::{TokenStream, LexError};
 
     pub fn new_token_stream(item: P<ast::Item>) -> TokenStream {
-        TokenStream { inner: vec![item] }
+        TokenStream { inner: TokenStream_::from_tokens(vec![
+            token::Interpolated(Rc::new(token::NtItem(item)))
+        ])}
+    }
+
+    pub fn token_stream_wrap(inner: TokenStream_) -> TokenStream {
+        TokenStream {
+            inner: inner
+        }
+    }
+
+    pub fn token_stream_parse_items(stream: TokenStream) -> Result<Vec<P<ast::Item>>, LexError> {
+        with_parse_sess(move |sess| {
+            let mut parser = parse::new_parser_from_ts(sess, stream.inner);
+            let mut items = Vec::new();
+
+            while let Some(item) = try!(parser.parse_item().map_err(super::parse_to_lex_err)) {
+                items.push(item)
+            }
+
+            Ok(items)
+        })
     }
 
-    pub fn token_stream_items(stream: TokenStream) -> Vec<P<ast::Item>> {
+    pub fn token_stream_inner(stream: TokenStream) -> TokenStream_ {
         stream.inner
     }
 
@@ -96,6 +120,10 @@ pub mod __internal {
                                   trait_name: &str,
                                   expand: fn(TokenStream) -> TokenStream,
                                   attributes: &[&'static str]);
+
+        fn register_attr_proc_macro(&mut self,
+                                    name: &str,
+                                    expand: fn(TokenStream, TokenStream) -> TokenStream);
     }
 
     // Emulate scoped_thread_local!() here essentially
@@ -125,11 +153,17 @@ pub mod __internal {
         where F: FnOnce(&ParseSess) -> R
     {
         let p = CURRENT_SESS.with(|p| p.get());
-        assert!(!p.is_null());
+        assert!(!p.is_null(), "proc_macro::__internal::with_parse_sess() called \
+                               before set_parse_sess()!");
         f(unsafe { &*p })
     }
 }
 
+fn parse_to_lex_err(mut err: DiagnosticBuilder) -> LexError {
+    err.cancel();
+    LexError { _inner: () }
+}
+
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl FromStr for TokenStream {
     type Err = LexError;
@@ -138,30 +172,17 @@ impl FromStr for TokenStream {
         __internal::with_parse_sess(|sess| {
             let src = src.to_string();
             let name = "<proc-macro source code>".to_string();
-            let mut parser = parse::new_parser_from_source_str(sess, name, src);
-            let mut ret = TokenStream { inner: Vec::new() };
-            loop {
-                match parser.parse_item() {
-                    Ok(Some(item)) => ret.inner.push(item),
-                    Ok(None) => return Ok(ret),
-                    Err(mut err) => {
-                        err.cancel();
-                        return Err(LexError { _inner: () })
-                    }
-                }
-            }
+            let tts = try!(parse::parse_tts_from_source_str(name, src, sess)
+                .map_err(parse_to_lex_err));
+
+            Ok(__internal::token_stream_wrap(TokenStream_::from_tts(tts)))
         })
     }
 }
 
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl fmt::Display for TokenStream {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        for item in self.inner.iter() {
-            let item = syntax::print::pprust::item_to_string(item);
-            try!(f.write_str(&item));
-            try!(f.write_str("\n"));
-        }
-        Ok(())
+        self.inner.fmt(f)
     }
 }

src/libsyntax_ext/deriving/custom.rs

@@ -77,8 +77,9 @@ impl MultiItemModifier for CustomDerive {
             let inner = self.inner;
             panic::catch_unwind(panic::AssertUnwindSafe(|| inner(input)))
         });
-        let new_items = match res {
-            Ok(stream) => __internal::token_stream_items(stream),
+
+        let stream = match res {
+            Ok(stream) => stream,
             Err(e) => {
                 let msg = "custom derive attribute panicked";
                 let mut err = ecx.struct_span_fatal(span, msg);
@@ -94,6 +95,18 @@ impl MultiItemModifier for CustomDerive {
             }
         };
 
+        let new_items = __internal::set_parse_sess(&ecx.parse_sess, || {
+            match __internal::token_stream_parse_items(stream) {
+                Ok(new_items) => new_items,
+                Err(_) => {
+                    // FIXME: handle this better
+                    let msg = "custom derive produced unparseable tokens";
+                    ecx.struct_span_fatal(span, msg).emit();
+                    panic!(FatalError);
+                }
+            }
+        });
+
         let mut res = vec![Annotatable::Item(item)];
         // Reassign spans of all expanded items to the input `item`
         // for better errors here.

src/test/compile-fail-fulldeps/proc-macro/derive-bad.rs

@@ -16,8 +16,7 @@ extern crate derive_bad;
 #[derive(
     A
 )]
-//~^^ ERROR: custom derive attribute panicked
-//~| HELP: called `Result::unwrap()` on an `Err` value: LexError
+//~^^ ERROR: custom derive produced unparseable tokens
 struct A;
 
 fn main() {}

src/test/run-pass-fulldeps/proc-macro/auxiliary/derive-atob.rs

@@ -19,6 +19,6 @@ use proc_macro::TokenStream;
 #[proc_macro_derive(AToB)]
 pub fn derive(input: TokenStream) -> TokenStream {
     let input = input.to_string();
-    assert_eq!(input, "#[derive(Copy, Clone)]\nstruct A;\n");
+    assert_eq!(input, "#[derive(Copy, Clone)]\nstruct A;");
     "struct B;".parse().unwrap()
 }

src/test/run-pass-fulldeps/proc-macro/auxiliary/derive-ctod.rs

@@ -19,6 +19,6 @@ use proc_macro::TokenStream;
 #[proc_macro_derive(CToD)]
 pub fn derive(input: TokenStream) -> TokenStream {
     let input = input.to_string();
-    assert_eq!(input, "struct C;\n");
+    assert_eq!(input, "struct C;");
     "struct D;".parse().unwrap()
 }

src/test/run-pass-fulldeps/proc-macro/auxiliary/derive-same-struct.rs

@@ -18,12 +18,12 @@ use proc_macro::TokenStream;
 #[proc_macro_derive(AToB)]
 pub fn derive1(input: TokenStream) -> TokenStream {
     println!("input1: {:?}", input.to_string());
-    assert_eq!(input.to_string(), "struct A;\n");
+    assert_eq!(input.to_string(), "struct A;");
     "#[derive(BToC)] struct B;".parse().unwrap()
 }
 
 #[proc_macro_derive(BToC)]
 pub fn derive2(input: TokenStream) -> TokenStream {
-    assert_eq!(input.to_string(), "struct B;\n");
+    assert_eq!(input.to_string(), "struct B;");
     "struct C;".parse().unwrap()
 }