Do not remove closures during ReplaceBodyWithLoop pass

Fixes rust-lang#71820

Author: marmeladema

src/librustc_ast/ast.rs

@@ -899,6 +899,15 @@ impl Stmt {
             _ => false,
         }
     }
+
+    pub fn as_expr(&self) -> Option<&Expr> {
+        match self.kind {
+            StmtKind::Local(ref l) => l.init.as_ref().map(|e| &**e),
+            StmtKind::Expr(ref e) => Some(e),
+            StmtKind::Semi(ref e) => Some(e),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Clone, RustcEncodable, RustcDecodable, Debug)]

src/librustc_interface/util.rs

@@ -719,6 +719,19 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
             }
         }
 
+        // We need to keep certain Expr nodes because they can contain
+        // exported macros. If we were to remove a parent of a MacroDef
+        // node, we would end up with in a case where the macro is properly
+        // lowered but not its parent, causing some surprising bug while
+        // trying to access the parent HirId. See issue #71820.
+        fn expr_to_keep(e: &ast::Expr) -> bool {
+            match e.kind {
+                ast::ExprKind::Closure(..) => true,
+                ast::ExprKind::Async(..) => true,
+                _ => false,
+            }
+        }
+
         let empty_block = stmt_to_block(BlockCheckMode::Default, None, self.resolver);
         let loop_expr = P(ast::Expr {
             kind: ast::ExprKind::Loop(P(empty_block), None),
@@ -741,7 +754,11 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
                 for s in b.stmts {
                     let old_blocks = self.nested_blocks.replace(vec![]);
 
-                    stmts.extend(self.flat_map_stmt(s).into_iter().filter(|s| s.is_item()));
+                    stmts.extend(
+                        self.flat_map_stmt(s).into_iter().filter(|s| {
+                            s.is_item() || s.as_expr().map_or(false, |e| expr_to_keep(e))
+                        }),
+                    );
 
                     // we put a Some in there earlier with that replace(), so this is valid
                     let new_blocks = self.nested_blocks.take().unwrap();

src/librustc_privacy/lib.rs

@@ -944,11 +944,8 @@ impl Visitor<'tcx> for EmbargoVisitor<'tcx> {
         let macro_module_def_id =
             ty::DefIdTree::parent(self.tcx, self.tcx.hir().local_def_id(md.hir_id).to_def_id())
                 .unwrap();
-        // FIXME(#71104) Should really be using just `as_local_hir_id` but
-        // some `DefId` do not seem to have a corresponding HirId.
-        let hir_id = macro_module_def_id
-            .as_local()
-            .and_then(|def_id| self.tcx.hir().opt_local_def_id_to_hir_id(def_id));
+        let hir_id =
+            macro_module_def_id.as_local().map(|def_id| self.tcx.hir().as_local_hir_id(def_id));
         let mut module_id = match hir_id {
             Some(module_id) if self.tcx.hir().is_hir_id_module(module_id) => module_id,
             // `module_id` doesn't correspond to a `mod`, return early (#63164, #65252).

src/test/rustdoc/macro-in-async-block.rs

@@ -0,0 +1,9 @@
+// Regression issue for rustdoc ICE encountered in PR #72088.
+// edition:2018
+#![feature(decl_macro)]
+
+fn main() {
+    async {
+        macro m() {}
+    };
+}

src/test/rustdoc/macro-in-closure.rs

@@ -6,4 +6,7 @@ fn main() {
     || {
         macro m() {}
     };
+    let _ = || {
+        macro n() {}
+    };
 }