proc_macro: Fix crate root detection

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;
     }
 

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
+}

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
+