From 29bb1cb37cb57dc7b14cf989cd533686aae6a9ad Mon Sep 17 00:00:00 2001
From: Aaron Hill <aa1ronham@gmail.com>
Date: Sun, 24 Jan 2021 14:16:45 -0500
Subject: [PATCH] When snapshoting a TokenCursor, preserve the top frame of the
 stack

If we start capturing on the first token of a frame,
we can legally pop the frame (e.g. we can capture something like
'{ my_tokens } other_tokens'. As a result, we need to snapshot the frame
at the top of the current stack, rather than just using an empty stack.

Fixes taiki-e/pin-project#312
---
 compiler/rustc_parse/src/parser/mod.rs        |  8 ++++---
 .../auxiliary/nonterminal-recollect-attr.rs   | 23 +++++++++++++++++++
 .../proc-macro/nonterminal-recollect-attr.rs  | 17 ++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 src/test/ui/proc-macro/auxiliary/nonterminal-recollect-attr.rs
 create mode 100644 src/test/ui/proc-macro/nonterminal-recollect-attr.rs

diff --git a/compiler/rustc_parse/src/parser/mod.rs b/compiler/rustc_parse/src/parser/mod.rs
index f11c6591de909..dff553d9d39b3 100644
--- a/compiler/rustc_parse/src/parser/mod.rs
+++ b/compiler/rustc_parse/src/parser/mod.rs
@@ -1260,9 +1260,11 @@ impl<'a> Parser<'a> {
         let start_token = (self.token.clone(), self.token_spacing);
         let cursor_snapshot = TokenCursor {
             frame: self.token_cursor.frame.clone(),
-            // We only ever capture tokens within our current frame,
-            // so we can just use an empty frame stack
-            stack: vec![],
+            // If we start collecting on the first token of a frame
+            // (e.g. the '{' in '{ my_tokens }'), then we can legally
+            // exit the frame (e.g. collecting '{ my_tokens } something_else'
+            // Therefore, we preserve the top of the stack, but nothing else
+            stack: self.token_cursor.stack.first().cloned().into_iter().collect(),
             desugar_doc_comments: self.token_cursor.desugar_doc_comments,
             num_next_calls: self.token_cursor.num_next_calls,
             append_unglued_token: self.token_cursor.append_unglued_token.clone(),
diff --git a/src/test/ui/proc-macro/auxiliary/nonterminal-recollect-attr.rs b/src/test/ui/proc-macro/auxiliary/nonterminal-recollect-attr.rs
new file mode 100644
index 0000000000000..a6903283aa108
--- /dev/null
+++ b/src/test/ui/proc-macro/auxiliary/nonterminal-recollect-attr.rs
@@ -0,0 +1,23 @@
+// force-host
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![feature(proc_macro_quote)]
+
+extern crate proc_macro;
+use proc_macro::{TokenStream, quote};
+
+#[proc_macro_attribute]
+pub fn first_attr(_: TokenStream, input: TokenStream) -> TokenStream {
+    let recollected: TokenStream = input.into_iter().collect();
+    quote! {
+        #[second_attr]
+        $recollected
+    }
+}
+
+#[proc_macro_attribute]
+pub fn second_attr(_: TokenStream, input: TokenStream) -> TokenStream {
+    let _recollected: TokenStream = input.into_iter().collect();
+    TokenStream::new()
+}
diff --git a/src/test/ui/proc-macro/nonterminal-recollect-attr.rs b/src/test/ui/proc-macro/nonterminal-recollect-attr.rs
new file mode 100644
index 0000000000000..5d4649b78c270
--- /dev/null
+++ b/src/test/ui/proc-macro/nonterminal-recollect-attr.rs
@@ -0,0 +1,17 @@
+// check-pass
+// aux-build:nonterminal-recollect-attr.rs
+
+extern crate nonterminal_recollect_attr;
+use nonterminal_recollect_attr::*;
+
+macro_rules! my_macro {
+    ($v:ident) => {
+        #[first_attr]
+        $v struct Foo {
+            field: u8
+        }
+    }
+}
+
+my_macro!(pub);
+fn main() {}