Lift #[macro_export] macros to top level in HIR

Author: inquisitivecrystal

compiler/rustc_ast_lowering/src/item.rs

@@ -10,7 +10,7 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
-use rustc_hir::def_id::LocalDefId;
+use rustc_hir::def_id::{LocalDefId, CRATE_DEF_ID};
 use rustc_span::source_map::{respan, DesugaringKind};
 use rustc_span::symbol::{kw, sym, Ident};
 use rustc_span::Span;
@@ -26,44 +26,44 @@ pub(super) struct ItemLowerer<'a, 'lowering, 'hir> {
 }
 
 impl ItemLowerer<'_, '_, '_> {
-    fn with_trait_impl_ref(&mut self, impl_ref: &Option<TraitRef>, f: impl FnOnce(&mut Self)) {
+    fn with_trait_impl_ref<T>(
+        &mut self,
+        impl_ref: &Option<TraitRef>,
+        f: impl FnOnce(&mut Self) -> T,
+    ) -> T {
         let old = self.lctx.is_in_trait_impl;
         self.lctx.is_in_trait_impl = impl_ref.is_some();
-        f(self);
+        let ret = f(self);
         self.lctx.is_in_trait_impl = old;
+        ret
     }
 }
 
 impl<'a> Visitor<'a> for ItemLowerer<'a, '_, '_> {
     fn visit_item(&mut self, item: &'a Item) {
-        let mut item_hir_id = None;
-        self.lctx.with_hir_id_owner(item.id, |lctx| {
-            lctx.without_in_scope_lifetime_defs(|lctx| {
-                if let Some(hir_item) = lctx.lower_item(item) {
-                    let id = lctx.insert_item(hir_item);
-                    item_hir_id = Some(id);
-                }
-            })
+        let hir_item = self.lctx.with_hir_id_owner(item.id, |lctx| {
+            lctx.without_in_scope_lifetime_defs(|lctx| lctx.lower_item(item))
         });
-
-        if let Some(hir_id) = item_hir_id {
-            self.lctx.with_parent_item_lifetime_defs(hir_id, |this| {
-                let this = &mut ItemLowerer { lctx: this };
-                match item.kind {
-                    ItemKind::Mod(..) => {
-                        let def_id = this.lctx.lower_node_id(item.id).expect_owner();
-                        let old_current_module =
-                            mem::replace(&mut this.lctx.current_module, def_id);
-                        visit::walk_item(this, item);
-                        this.lctx.current_module = old_current_module;
-                    }
-                    ItemKind::Impl(box ImplKind { ref of_trait, .. }) => {
-                        this.with_trait_impl_ref(of_trait, |this| visit::walk_item(this, item));
-                    }
-                    _ => visit::walk_item(this, item),
+        let force_top_level = self.lctx.check_force_top_level(item);
+        let hir_id = self.lctx.insert_item(hir_item, force_top_level);
+
+        let was_at_top_level = mem::replace(&mut self.lctx.is_at_top_level, false);
+        self.lctx.with_parent_item_lifetime_defs(hir_id, |this| {
+            let this = &mut ItemLowerer { lctx: this };
+            match item.kind {
+                ItemKind::Mod(..) => {
+                    let def_id = this.lctx.lower_node_id(item.id).expect_owner();
+                    let old_current_module = mem::replace(&mut this.lctx.current_module, def_id);
+                    visit::walk_item(this, item);
+                    this.lctx.current_module = old_current_module;
                 }
-            });
-        }
+                ItemKind::Impl(box ImplKind { ref of_trait, .. }) => {
+                    this.with_trait_impl_ref(of_trait, |this| visit::walk_item(this, item));
+                }
+                _ => visit::walk_item(this, item),
+            }
+        });
+        self.lctx.is_at_top_level = was_at_top_level;
     }
 
     fn visit_fn(&mut self, fk: FnKind<'a>, sp: Span, _: NodeId) {
@@ -113,7 +113,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     fn with_parent_item_lifetime_defs<T>(
         &mut self,
         parent_hir_id: hir::ItemId,
-        f: impl FnOnce(&mut LoweringContext<'_, '_>) -> T,
+        f: impl FnOnce(&mut Self) -> T,
     ) -> T {
         let old_len = self.in_scope_lifetimes.len();
 
@@ -137,10 +137,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     // Clears (and restores) the `in_scope_lifetimes` field. Used when
     // visiting nested items, which never inherit in-scope lifetimes
     // from their surrounding environment.
-    fn without_in_scope_lifetime_defs<T>(
-        &mut self,
-        f: impl FnOnce(&mut LoweringContext<'_, '_>) -> T,
-    ) -> T {
+    fn without_in_scope_lifetime_defs<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> T {
         let old_in_scope_lifetimes = mem::replace(&mut self.in_scope_lifetimes, vec![]);
 
         // this vector is only used when walking over impl headers,
@@ -157,10 +154,23 @@ impl<'hir> LoweringContext<'_, 'hir> {
     }
 
     pub(super) fn lower_mod(&mut self, items: &[P<Item>], inner: Span) -> hir::Mod<'hir> {
-        hir::Mod {
-            inner,
-            item_ids: self.arena.alloc_from_iter(items.iter().flat_map(|x| self.lower_item_id(x))),
+        let mut items: Vec<_> = items
+            .iter()
+            .filter_map(|x| {
+                if self.check_force_top_level(&x) { None } else { Some(self.lower_item_id(&x)) }
+            })
+            .flatten()
+            .collect();
+
+        if self.current_hir_id_owner.0 == CRATE_DEF_ID {
+            let top_level_items = self.top_level_items.clone();
+            for item in top_level_items {
+                let id = self.lower_item_id(&item)[0];
+                items.push(id);
+            }
         }
+
+        hir::Mod { inner, item_ids: self.arena.alloc_from_iter(items.into_iter()) }
     }
 
     pub(super) fn lower_item_id(&mut self, i: &Item) -> SmallVec<[hir::ItemId; 1]> {
@@ -208,13 +218,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }
     }
 
-    pub fn lower_item(&mut self, i: &Item) -> Option<hir::Item<'hir>> {
+    pub fn lower_item(&mut self, i: &Item) -> hir::Item<'hir> {
+        let was_at_top_level = mem::replace(&mut self.is_at_top_level, false);
         let mut ident = i.ident;
         let mut vis = self.lower_visibility(&i.vis, None);
         let hir_id = self.lower_node_id(i.id);
         let attrs = self.lower_attrs(hir_id, &i.attrs);
         let kind = self.lower_item_kind(i.span, i.id, hir_id, &mut ident, attrs, &mut vis, &i.kind);
-        Some(hir::Item { def_id: hir_id.expect_owner(), ident, kind, vis, span: i.span })
+        self.is_at_top_level = was_at_top_level;
+        hir::Item { def_id: hir_id.expect_owner(), ident, kind, vis, span: i.span }
     }
 
     fn lower_item_kind(
@@ -285,7 +297,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
             ItemKind::Mod(_, ref mod_kind) => match mod_kind {
                 ModKind::Loaded(items, _, inner_span) => {
-                    hir::ItemKind::Mod(self.lower_mod(items, *inner_span))
+                    let module = self.lower_mod(items, *inner_span);
+                    hir::ItemKind::Mod(module)
                 }
                 ModKind::Unloaded => panic!("`mod` items should have been loaded by now"),
             },
@@ -537,13 +550,10 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             this.attrs.insert(new_id, attrs);
                         }
 
-                        this.insert_item(hir::Item {
-                            def_id: new_id.expect_owner(),
-                            ident,
-                            kind,
-                            vis,
-                            span,
-                        });
+                        this.insert_item(
+                            hir::Item { def_id: new_id.expect_owner(), ident, kind, vis, span },
+                            false,
+                        );
                     });
                 }
 
@@ -611,13 +621,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             this.attrs.insert(new_hir_id, attrs);
                         }
 
-                        this.insert_item(hir::Item {
-                            def_id: new_hir_id.expect_owner(),
-                            ident,
-                            kind,
-                            vis,
-                            span: use_tree.span,
-                        });
+                        this.insert_item(
+                            hir::Item {
+                                def_id: new_hir_id.expect_owner(),
+                                ident,
+                                kind,
+                                vis,
+                                span: use_tree.span,
+                            },
+                            false,
+                        );
                     });
                 }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -36,6 +36,7 @@
 #![recursion_limit = "256"]
 
 use rustc_ast::node_id::NodeMap;
+use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Token};
 use rustc_ast::tokenstream::{CanSynthesizeMissingTokens, TokenStream, TokenTree};
 use rustc_ast::visit::{self, AssocCtxt, Visitor};
@@ -120,8 +121,11 @@ struct LoweringContext<'a, 'hir: 'a> {
     /// outside of an `async fn`.
     current_item: Option<Span>,
 
+    top_level_items: Vec<Item>,
+
     catch_scopes: Vec<NodeId>,
     loop_scopes: Vec<NodeId>,
+    is_at_top_level: bool,
     is_in_loop_condition: bool,
     is_in_trait_impl: bool,
     is_in_dyn_type: bool,
@@ -331,6 +335,7 @@ pub fn lower_crate<'a, 'hir>(
         attrs: BTreeMap::default(),
         catch_scopes: Vec::new(),
         loop_scopes: Vec::new(),
+        is_at_top_level: true,
         is_in_loop_condition: false,
         is_in_trait_impl: false,
         is_in_dyn_type: false,
@@ -343,6 +348,7 @@ pub fn lower_crate<'a, 'hir>(
         generator_kind: None,
         task_context: None,
         current_item: None,
+        top_level_items: Vec::new(),
         lifetimes_to_define: Vec::new(),
         is_collecting_in_band_lifetimes: false,
         in_scope_lifetimes: Vec::new(),
@@ -467,13 +473,22 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                             .count();
                         self.lctx.type_def_lifetime_params.insert(def_id.to_def_id(), count);
                     }
+                    ItemKind::MacroDef(..) => {
+                        if self.lctx.check_force_top_level(item) {
+                            self.lctx.top_level_items.push(item.clone());
+                        }
+                    }
                     ItemKind::Use(ref use_tree) => {
                         self.allocate_use_tree_hir_id_counters(use_tree);
                     }
                     _ => {}
                 }
 
+                let was_at_top_level = mem::replace(&mut self.lctx.is_at_top_level, false);
+
                 visit::walk_item(self, item);
+
+                self.lctx.is_at_top_level = was_at_top_level;
             }
 
             fn visit_assoc_item(&mut self, item: &'tcx AssocItem, ctxt: AssocCtxt) {
@@ -510,7 +525,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         visit::walk_crate(&mut MiscCollector { lctx: &mut self }, c);
         visit::walk_crate(&mut item::ItemLowerer { lctx: &mut self }, c);
 
-        let module = self.arena.alloc(self.lower_mod(&c.items, c.span));
+        let mut items: Vec<_> = self.top_level_items.drain(..).map(|item| P(item)).collect();
+        items.extend(c.items.clone());
+
+        let module = self.arena.alloc(self.lower_mod(items.as_slice(), c.span));
         self.lower_attrs(hir::CRATE_HIR_ID, &c.attrs);
         self.owners.ensure_contains_elem(CRATE_DEF_ID, || None);
         self.owners[CRATE_DEF_ID] = Some(hir::OwnerNode::Crate(module));
@@ -561,12 +579,13 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         self.arena.alloc(krate)
     }
 
-    fn insert_item(&mut self, item: hir::Item<'hir>) -> hir::ItemId {
+    fn insert_item(&mut self, item: hir::Item<'hir>, force_top_level: bool) -> hir::ItemId {
         let id = item.item_id();
         let item = self.arena.alloc(item);
+        let module = if force_top_level { CRATE_DEF_ID } else { self.current_module };
         self.owners.ensure_contains_elem(id.def_id, || None);
         self.owners[id.def_id] = Some(hir::OwnerNode::Item(item));
-        self.modules.entry(self.current_module).or_default().items.insert(id);
+        self.modules.entry(module).or_default().items.insert(id);
         id
     }
 
@@ -597,6 +616,19 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         id
     }
 
+    /// This checks if an item is a `#[macro_use]` macro, and thus always
+    /// defined at the top level. As a special case, if the macro is already
+    /// at the top level, it remains at its current level.
+    fn check_force_top_level(&self, item: &Item) -> bool {
+        if self.is_at_top_level {
+            false
+        } else if let ItemKind::MacroDef(MacroDef { macro_rules, .. }) = item.kind {
+            macro_rules && self.sess.contains_name(&item.attrs, sym::macro_export)
+        } else {
+            false
+        }
+    }
+
     fn allocate_hir_id_counter(&mut self, owner: NodeId) -> hir::HirId {
         // Set up the counter if needed.
         self.item_local_id_counters.entry(owner).or_insert(0);
@@ -1593,7 +1625,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         // Insert the item into the global item list. This usually happens
         // automatically for all AST items. But this opaque type item
         // does not actually exist in the AST.
-        self.insert_item(opaque_ty_item);
+        self.insert_item(opaque_ty_item, false);
     }
 
     fn lifetimes_from_impl_trait_bounds(
@@ -2382,7 +2414,22 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             [stmts @ .., Stmt { kind: StmtKind::Expr(e), .. }] => (stmts, Some(&*e)),
             stmts => (stmts, None),
         };
-        let stmts = self.arena.alloc_from_iter(stmts.iter().flat_map(|stmt| self.lower_stmt(stmt)));
+        let stmts: Vec<_> = stmts
+            .iter()
+            .filter_map(|stmt| {
+                if let StmtKind::Item(ref item) = stmt.kind {
+                    if self.check_force_top_level(item) {
+                        None
+                    } else {
+                        Some(self.lower_stmt(stmt))
+                    }
+                } else {
+                    Some(self.lower_stmt(stmt))
+                }
+            })
+            .flatten()
+            .collect();
+        let stmts = self.arena.alloc_from_iter(stmts.into_iter());
         let expr = expr.map(|e| self.lower_expr(e));
         let rules = self.lower_block_check_mode(&b.rules);
         let hir_id = self.lower_node_id(b.id);