From edffb2fcf077bc944ef508fe1578786547d49dc8 Mon Sep 17 00:00:00 2001
From: Vadim Petrochenkov <vadim.petrochenkov@gmail.com>
Date: Fri, 13 Jul 2018 01:59:41 +0300
Subject: [PATCH] proc_macro: Fix crate root detection

---
 src/libsyntax_ext/proc_macro_registrar.rs     | 16 +++++-------
 src/test/ui-fulldeps/proc-macro/non-root.rs   | 25 +++++++++++++++++++
 .../ui-fulldeps/proc-macro/non-root.stderr    |  8 ++++++
 3 files changed, 39 insertions(+), 10 deletions(-)
 create mode 100644 src/test/ui-fulldeps/proc-macro/non-root.rs
 create mode 100644 src/test/ui-fulldeps/proc-macro/non-root.stderr

diff --git a/src/libsyntax_ext/proc_macro_registrar.rs b/src/libsyntax_ext/proc_macro_registrar.rs
index 85aa84acc4221..ab2bb446631b7 100644
--- a/src/libsyntax_ext/proc_macro_registrar.rs
+++ b/src/libsyntax_ext/proc_macro_registrar.rs
@@ -12,7 +12,7 @@ use std::mem;
 
 use errors;
 
-use syntax::ast::{self, Ident, NodeId};
+use syntax::ast::{self, Ident};
 use syntax::attr;
 use syntax::codemap::{ExpnInfo, MacroAttribute, hygiene, respan};
 use syntax::ext::base::ExtCtxt;
@@ -293,7 +293,10 @@ impl<'a> Visitor<'a> for CollectProcMacros<'a> {
         let attr = match found_attr {
             None => {
                 self.check_not_pub_in_root(&item.vis, item.span);
-                return visit::walk_item(self, item);
+                let prev_in_root = mem::replace(&mut self.in_root, false);
+                visit::walk_item(self, item);
+                self.in_root = prev_in_root;
+                return;
             },
             Some(attr) => attr,
         };
@@ -326,15 +329,8 @@ impl<'a> Visitor<'a> for CollectProcMacros<'a> {
             self.collect_bang_proc_macro(item, attr);
         };
 
+        let prev_in_root = mem::replace(&mut self.in_root, false);
         visit::walk_item(self, item);
-    }
-
-    fn visit_mod(&mut self, m: &'a ast::Mod, _s: Span, _a: &[ast::Attribute], id: NodeId) {
-        let mut prev_in_root = self.in_root;
-        if id != ast::CRATE_NODE_ID {
-            prev_in_root = mem::replace(&mut self.in_root, false);
-        }
-        visit::walk_mod(self, m);
         self.in_root = prev_in_root;
     }
 
diff --git a/src/test/ui-fulldeps/proc-macro/non-root.rs b/src/test/ui-fulldeps/proc-macro/non-root.rs
new file mode 100644
index 0000000000000..288c63b4c7781
--- /dev/null
+++ b/src/test/ui-fulldeps/proc-macro/non-root.rs
@@ -0,0 +1,25 @@
+// 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
+
+#![feature(proc_macro)]
+#![crate_type = "proc-macro"]
+
+extern crate proc_macro;
+use proc_macro::*;
+
+fn foo(arg: TokenStream) -> TokenStream {
+    #[proc_macro]
+    pub fn foo(arg: TokenStream) -> TokenStream { arg }
+    //~^ ERROR functions tagged with `#[proc_macro]` must currently reside in the root of the crate
+
+    arg
+}
diff --git a/src/test/ui-fulldeps/proc-macro/non-root.stderr b/src/test/ui-fulldeps/proc-macro/non-root.stderr
new file mode 100644
index 0000000000000..8c14f644d7a64
--- /dev/null
+++ b/src/test/ui-fulldeps/proc-macro/non-root.stderr
@@ -0,0 +1,8 @@
+error: functions tagged with `#[proc_macro]` must currently reside in the root of the crate
+  --> $DIR/non-root.rs:21:5
+   |
+LL |     pub fn foo(arg: TokenStream) -> TokenStream { arg }
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to previous error
+