proc_macro: Avoid cached TokenStream more often

src/libsyntax/parse/token.rs

@@ -519,8 +519,17 @@ impl Token {
         // all span information.
         //
         // As a result, some AST nodes are annotated with the token
-        // stream they came from. Attempt to extract these lossless
-        // token streams before we fall back to the stringification.
+        // stream they came from. Here we attempt to extract these
+        // lossless token streams before we fall back to the
+        // stringification.
+        //
+        // During early phases of the compiler, though, the AST could
+        // get modified directly (e.g. attributes added or removed) and
+        // the internal cache of tokens my not be invalidated or
+        // updated. Consequently if the "lossless" token stream
+        // disagrees with our actuall stringification (which has
+        // historically been much more battle-tested) then we go with
+        // the lossy stream anyway (losing span information).
         let mut tokens = None;
 
         match nt.0 {
@@ -547,13 +556,17 @@ impl Token {
             _ => {}
         }
 
-        tokens.unwrap_or_else(|| {
-            nt.1.force(|| {
-                // FIXME(jseyfried): Avoid this pretty-print + reparse hack
-                let source = pprust::token_to_string(self);
-                parse_stream_from_source_str(FileName::MacroExpansion, source, sess, Some(span))
-            })
-        })
+        let tokens_for_real = nt.1.force(|| {
+            // FIXME(#43081): Avoid this pretty-print + reparse hack
+            let source = pprust::token_to_string(self);
+            parse_stream_from_source_str(FileName::MacroExpansion, source, sess, Some(span))
+        });
+        if let Some(tokens) = tokens {
+            if tokens.eq_unspanned(&tokens_for_real) {
+                return tokens
+            }
+        }
+        return tokens_for_real
     }
 }
 

src/libsyntax/tokenstream.rs

@@ -118,7 +118,7 @@ impl TokenTree {
             (&TokenTree::Token(_, ref tk), &TokenTree::Token(_, ref tk2)) => tk == tk2,
             (&TokenTree::Delimited(_, ref dl), &TokenTree::Delimited(_, ref dl2)) => {
                 dl.delim == dl2.delim &&
-                dl.stream().trees().zip(dl2.stream().trees()).all(|(tt, tt2)| tt.eq_unspanned(&tt2))
+                dl.stream().eq_unspanned(&dl2.stream())
             }
             (_, _) => false,
         }
@@ -240,12 +240,14 @@ impl TokenStream {
 
     /// Compares two TokenStreams, checking equality without regarding span information.
     pub fn eq_unspanned(&self, other: &TokenStream) -> bool {
-        for (t1, t2) in self.trees().zip(other.trees()) {
+        let mut t1 = self.trees();
+        let mut t2 = other.trees();
+        for (t1, t2) in t1.by_ref().zip(t2.by_ref()) {
             if !t1.eq_unspanned(&t2) {
                 return false;
             }
         }
-        true
+        t1.next().is_none() && t2.next().is_none()
     }
 
     /// Precondition: `self` consists of a single token tree.

src/test/run-pass-fulldeps/proc-macro/auxiliary/modify-ast.rs

@@ -0,0 +1,57 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![feature(proc_macro)]
+
+extern crate proc_macro;
+
+use proc_macro::*;
+
+#[proc_macro_attribute]
+pub fn assert1(_a: TokenStream, b: TokenStream) -> TokenStream {
+    assert_eq(b.clone(), "pub fn foo() {}".parse().unwrap());
+    b
+}
+
+#[proc_macro_derive(Foo, attributes(foo))]
+pub fn assert2(a: TokenStream) -> TokenStream {
+    assert_eq(a, "pub struct MyStructc { _a: i32, }".parse().unwrap());
+    TokenStream::empty()
+}
+
+fn assert_eq(a: TokenStream, b: TokenStream) {
+    let mut a = a.into_iter();
+    let mut b = b.into_iter();
+    for (a, b) in a.by_ref().zip(&mut b) {
+        match (a, b) {
+            (TokenTree::Group(a), TokenTree::Group(b)) => {
+                assert_eq!(a.delimiter(), b.delimiter());
+                assert_eq(a.stream(), b.stream());
+            }
+            (TokenTree::Op(a), TokenTree::Op(b)) => {
+                assert_eq!(a.op(), b.op());
+                assert_eq!(a.spacing(), b.spacing());
+            }
+            (TokenTree::Literal(a), TokenTree::Literal(b)) => {
+                assert_eq!(a.to_string(), b.to_string());
+            }
+            (TokenTree::Term(a), TokenTree::Term(b)) => {
+                assert_eq!(a.to_string(), b.to_string());
+            }
+            (a, b) => panic!("{:?} != {:?}", a, b),
+        }
+    }
+
+    assert!(a.next().is_none());
+    assert!(b.next().is_none());
+}

src/test/run-pass-fulldeps/proc-macro/modify-ast.rs

@@ -0,0 +1,37 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// aux-build:modify-ast.rs
+
+#![feature(proc_macro)]
+
+extern crate modify_ast;
+
+use modify_ast::*;
+
+#[derive(Foo)]
+pub struct MyStructc {
+    #[cfg_attr(my_cfg, foo)]
+    _a: i32,
+}
+
+macro_rules! a {
+    ($i:item) => ($i)
+}
+
+a! {
+    #[assert1]
+    pub fn foo() {}
+}
+
+fn main() {
+    let _a = MyStructc { _a: 0 };
+    foo();
+}