proc_macro: Avoid cached TokenStream more often

This commit adds even more pessimization to use the cached `TokenStream` inside
of an AST node. As a reminder the `proc_macro` API requires taking an arbitrary
AST node and transforming it back into a `TokenStream` to hand off to a
procedural macro. Such functionality isn't actually implemented in rustc today,
so the way `proc_macro` works today is that it stringifies an AST node and then
reparses for a list of tokens.

This strategy unfortunately loses all span information, so we try to avoid it
whenever possible. Implemented in rust-lang#43230 some AST nodes have a `TokenStream`
cache representing the tokens they were originally parsed from. This
`TokenStream` cache, however, has turned out to not always reflect the current
state of the item when it's being tokenized. For example `#[cfg]` processing or
macro expansion could modify the state of an item. Consequently we've seen a
number of bugs (rust-lang#48644 and rust-lang#49846) related to using this stale cache.

This commit tweaks the usage of the cached `TokenStream` to compare it to our
lossy stringification of the token stream. If the tokens that make up the cache
and the stringified token stream are the same then we return the cached version
(which has correct span information). If they differ, however, then we will
return the stringified version as the cache has been invalidated and we just
haven't figured that out.

Closes rust-lang#48644
Closes rust-lang#49846

Author: alexcrichton

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();
+}