Rollup merge of rust-lang#73103 - ecstatic-morse:replace-body-with-loop, r=pnkfelix

Preserve `Expr`s that have `DefId`s in `ReplaceBodyWithLoop`

This PR fixes the last part of rust-lang#71104 by preserving expressions that are assigned their own `DefId`s (closures and `async` blocks) when passing them to `rustdoc`. This avoids having a `DefId` without a corresponding `HirId`.

The first commit in this PR makes `-Zunpretty=everybody_loops` actually work again, and the subsequent two are miscellaneous cleanup. They should probably get merged regardless of what we end up doing here.

Sample input:
```rust
fn foo() -> Box<i32> {
    let x = |a: i64| {
        const FOO: i64 = 1;
    };

    let a = 4;
    Box::new(a)
}
```

Sample output:
```rust
fn foo() -> Box<i32> {
    || -> !
        {
            const FOO: i64 = 1;
            loop  { }
        };
    loop  { }
}
```

r? @ghost

Author: Dylan-DPC

src/librustc_ast/ast.rs

@@ -25,6 +25,7 @@ pub use UnsafeSource::*;
 use crate::ptr::P;
 use crate::token::{self, DelimToken};
 use crate::tokenstream::{DelimSpan, TokenStream, TokenTree};
+use crate::visit::{walk_ty, Visitor};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::Lrc;
@@ -1795,6 +1796,26 @@ pub struct Ty {
     pub span: Span,
 }
 
+impl Ty {
+    /// Returns `true` if the kind of this type or any types contained within are `impl Trait`.
+    pub fn contains_impl_trait(&self) -> bool {
+        struct ContainsImplTrait(bool);
+
+        impl<'ast> Visitor<'ast> for ContainsImplTrait {
+            fn visit_ty(&mut self, t: &'ast Ty) {
+                self.0 |= matches!(t.kind, TyKind::ImplTrait(..));
+                if !self.0 {
+                    walk_ty(self, t);
+                }
+            }
+        }
+
+        let mut vis = ContainsImplTrait(false);
+        vis.visit_ty(self);
+        vis.0
+    }
+}
+
 #[derive(Clone, RustcEncodable, RustcDecodable, Debug)]
 pub struct BareFnTy {
     pub unsafety: Unsafe,

src/librustc_interface/util.rs

@@ -592,97 +592,54 @@ pub fn build_output_filenames(
 //
 // [#34511]: https://github.com/rust-lang/rust/issues/34511#issuecomment-322340401
 pub struct ReplaceBodyWithLoop<'a, 'b> {
-    within_static_or_const: bool,
+    ignore_item: bool,
     nested_blocks: Option<Vec<ast::Block>>,
     resolver: &'a mut Resolver<'b>,
 }
 
 impl<'a, 'b> ReplaceBodyWithLoop<'a, 'b> {
     pub fn new(resolver: &'a mut Resolver<'b>) -> ReplaceBodyWithLoop<'a, 'b> {
-        ReplaceBodyWithLoop { within_static_or_const: false, nested_blocks: None, resolver }
+        ReplaceBodyWithLoop { ignore_item: false, nested_blocks: None, resolver }
     }
 
-    fn run<R, F: FnOnce(&mut Self) -> R>(&mut self, is_const: bool, action: F) -> R {
-        let old_const = mem::replace(&mut self.within_static_or_const, is_const);
+    fn run<R, F: FnOnce(&mut Self) -> R>(&mut self, ignore_item: bool, action: F) -> R {
+        let old_ignore_item = mem::replace(&mut self.ignore_item, ignore_item);
         let old_blocks = self.nested_blocks.take();
         let ret = action(self);
-        self.within_static_or_const = old_const;
+        self.ignore_item = old_ignore_item;
         self.nested_blocks = old_blocks;
         ret
     }
 
-    fn should_ignore_fn(ret_ty: &ast::FnRetTy) -> bool {
-        if let ast::FnRetTy::Ty(ref ty) = ret_ty {
-            fn involves_impl_trait(ty: &ast::Ty) -> bool {
-                match ty.kind {
-                    ast::TyKind::ImplTrait(..) => true,
-                    ast::TyKind::Slice(ref subty)
-                    | ast::TyKind::Array(ref subty, _)
-                    | ast::TyKind::Ptr(ast::MutTy { ty: ref subty, .. })
-                    | ast::TyKind::Rptr(_, ast::MutTy { ty: ref subty, .. })
-                    | ast::TyKind::Paren(ref subty) => involves_impl_trait(subty),
-                    ast::TyKind::Tup(ref tys) => any_involves_impl_trait(tys.iter()),
-                    ast::TyKind::Path(_, ref path) => {
-                        path.segments.iter().any(|seg| match seg.args.as_deref() {
-                            None => false,
-                            Some(&ast::GenericArgs::AngleBracketed(ref data)) => {
-                                data.args.iter().any(|arg| match arg {
-                                    ast::AngleBracketedArg::Arg(arg) => match arg {
-                                        ast::GenericArg::Type(ty) => involves_impl_trait(ty),
-                                        ast::GenericArg::Lifetime(_)
-                                        | ast::GenericArg::Const(_) => false,
-                                    },
-                                    ast::AngleBracketedArg::Constraint(c) => match c.kind {
-                                        ast::AssocTyConstraintKind::Bound { .. } => true,
-                                        ast::AssocTyConstraintKind::Equality { ref ty } => {
-                                            involves_impl_trait(ty)
-                                        }
-                                    },
-                                })
-                            }
-                            Some(&ast::GenericArgs::Parenthesized(ref data)) => {
-                                any_involves_impl_trait(data.inputs.iter())
-                                    || ReplaceBodyWithLoop::should_ignore_fn(&data.output)
-                            }
-                        })
-                    }
-                    _ => false,
-                }
-            }
-
-            fn any_involves_impl_trait<'a, I: Iterator<Item = &'a P<ast::Ty>>>(mut it: I) -> bool {
-                it.any(|subty| involves_impl_trait(subty))
-            }
-
-            involves_impl_trait(ty)
-        } else {
-            false
-        }
+    fn should_ignore_fn(sig: &ast::FnSig) -> bool {
+        matches!(sig.header.constness, ast::Const::Yes(_))
+            || matches!(&sig.decl.output, ast::FnRetTy::Ty(ty) if ty.contains_impl_trait())
     }
 
-    fn is_sig_const(sig: &ast::FnSig) -> bool {
-        matches!(sig.header.constness, ast::Const::Yes(_))
-            || ReplaceBodyWithLoop::should_ignore_fn(&sig.decl.output)
+    /// Keep some `Expr`s that are ultimately assigned `DefId`s. This keeps the `HirId` parent
+    /// mappings correct even after this pass runs.
+    fn should_preserve_expr(e: &ast::Expr) -> bool {
+        matches!(&e.kind, ast::ExprKind::Closure(..) | ast::ExprKind::Async(..))
     }
 }
 
 impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
     fn visit_item_kind(&mut self, i: &mut ast::ItemKind) {
-        let is_const = match i {
+        let ignore_item = match i {
             ast::ItemKind::Static(..) | ast::ItemKind::Const(..) => true,
-            ast::ItemKind::Fn(_, ref sig, _, _) => Self::is_sig_const(sig),
+            ast::ItemKind::Fn(_, ref sig, _, _) => Self::should_ignore_fn(sig),
             _ => false,
         };
-        self.run(is_const, |s| noop_visit_item_kind(i, s))
+        self.run(ignore_item, |s| noop_visit_item_kind(i, s))
     }
 
     fn flat_map_trait_item(&mut self, i: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
-        let is_const = match i.kind {
+        let ignore_item = match i.kind {
             ast::AssocItemKind::Const(..) => true,
-            ast::AssocItemKind::Fn(_, ref sig, _, _) => Self::is_sig_const(sig),
+            ast::AssocItemKind::Fn(_, ref sig, _, _) => Self::should_ignore_fn(sig),
             _ => false,
         };
-        self.run(is_const, |s| noop_flat_map_assoc_item(i, s))
+        self.run(ignore_item, |s| noop_flat_map_assoc_item(i, s))
     }
 
     fn flat_map_impl_item(&mut self, i: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
@@ -693,6 +650,75 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
         self.run(true, |s| noop_visit_anon_const(c, s))
     }
 
+    fn filter_map_expr(&mut self, mut e: P<ast::Expr>) -> Option<P<ast::Expr>> {
+        if self.ignore_item {
+            return noop_filter_map_expr(e, self);
+        }
+
+        if let ast::ExprKind::Closure(.., decl, expr, _) = &mut e.kind {
+            // Replacing a closure body and removing its callsites will break inference. Give
+            // all closures the signature `|| -> !` to work around this.
+            decl.inputs = vec![];
+            if let ast::FnRetTy::Default(_) = decl.output {
+                decl.output = ast::FnRetTy::Ty(P(ast::Ty {
+                    id: self.resolver.next_node_id(),
+                    span: rustc_span::DUMMY_SP,
+                    kind: ast::TyKind::Never,
+                }));
+            }
+
+            // Replace `|| expr` with `|| { expr }` so that `visit_block` gets called on the
+            // closure body.
+            if !matches!(expr.kind, ast::ExprKind::Block(..)) {
+                let new_stmt = ast::Stmt {
+                    kind: ast::StmtKind::Expr(expr.clone()),
+                    id: self.resolver.next_node_id(),
+                    span: expr.span,
+                };
+
+                let new_block = ast::Block {
+                    stmts: vec![new_stmt],
+                    rules: BlockCheckMode::Default,
+                    id: self.resolver.next_node_id(),
+                    span: expr.span,
+                };
+
+                expr.kind = ast::ExprKind::Block(P(new_block), None);
+            }
+        }
+
+        if Self::should_preserve_expr(&e) {
+            self.run(false, |s| noop_filter_map_expr(e, s))
+        } else {
+            noop_filter_map_expr(e, self)
+        }
+    }
+
+    fn flat_map_stmt(&mut self, s: ast::Stmt) -> SmallVec<[ast::Stmt; 1]> {
+        if self.ignore_item {
+            return noop_flat_map_stmt(s, self);
+        }
+
+        let ast::Stmt { id, span, .. } = s;
+        match s.kind {
+            // Replace `let x = || {};` with `|| {};`
+            ast::StmtKind::Local(mut local) if local.init.is_some() => self
+                .filter_map_expr(local.init.take().unwrap())
+                .into_iter()
+                .map(|e| ast::Stmt { kind: ast::StmtKind::Semi(e), id, span })
+                .collect(),
+
+            // Replace `|| {}` with `|| {};`
+            ast::StmtKind::Expr(expr) => self
+                .filter_map_expr(expr)
+                .into_iter()
+                .map(|e| ast::Stmt { kind: ast::StmtKind::Semi(e), id, span })
+                .collect(),
+
+            _ => noop_flat_map_stmt(s, self),
+        }
+    }
+
     fn visit_block(&mut self, b: &mut P<ast::Block>) {
         fn stmt_to_block(
             rules: ast::BlockCheckMode,
@@ -723,6 +749,11 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
             }
         }
 
+        if self.ignore_item {
+            noop_visit_block(b, self);
+            return;
+        }
+
         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),
@@ -738,39 +769,41 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
             kind: ast::StmtKind::Expr(loop_expr),
         };
 
-        if self.within_static_or_const {
-            noop_visit_block(b, self)
-        } else {
-            visit_clobber(b.deref_mut(), |b| {
-                let mut stmts = vec![];
-                for s in b.stmts {
-                    let old_blocks = self.nested_blocks.replace(vec![]);
+        visit_clobber(b.deref_mut(), |b| {
+            let mut stmts = vec![];
+            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()
+                        || matches!(
+                            &s.kind,
+                            ast::StmtKind::Semi(expr) if Self::should_preserve_expr(expr)
+                        )
+                }));
+
+                // we put a Some in there earlier with that replace(), so this is valid
+                let new_blocks = self.nested_blocks.take().unwrap();
+                self.nested_blocks = old_blocks;
+                stmts.extend(new_blocks.into_iter().map(|b| block_to_stmt(b, self.resolver)));
+            }
 
-                    stmts.extend(self.flat_map_stmt(s).into_iter().filter(|s| s.is_item()));
+            let mut new_block = ast::Block { stmts, ..b };
 
-                    // we put a Some in there earlier with that replace(), so this is valid
-                    let new_blocks = self.nested_blocks.take().unwrap();
-                    self.nested_blocks = old_blocks;
-                    stmts.extend(new_blocks.into_iter().map(|b| block_to_stmt(b, self.resolver)));
+            if let Some(old_blocks) = self.nested_blocks.as_mut() {
+                //push our fresh block onto the cache and yield an empty block with `loop {}`
+                if !new_block.stmts.is_empty() {
+                    old_blocks.push(new_block);
                 }
 
-                let mut new_block = ast::Block { stmts, ..b };
-
-                if let Some(old_blocks) = self.nested_blocks.as_mut() {
-                    //push our fresh block onto the cache and yield an empty block with `loop {}`
-                    if !new_block.stmts.is_empty() {
-                        old_blocks.push(new_block);
-                    }
-
-                    stmt_to_block(b.rules, Some(loop_stmt), &mut self.resolver)
-                } else {
-                    //push `loop {}` onto the end of our fresh block and yield that
-                    new_block.stmts.push(loop_stmt);
+                stmt_to_block(b.rules, Some(loop_stmt), &mut self.resolver)
+            } else {
+                //push `loop {}` onto the end of our fresh block and yield that
+                new_block.stmts.push(loop_stmt);
 
-                    new_block
-                }
-            })
-        }
+                new_block
+            }
+        })
     }
 
     // in general the pretty printer processes unexpanded code, so

src/librustc_privacy/lib.rs

@@ -943,11 +943,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/librustc_session/config.rs

@@ -1954,9 +1954,11 @@ impl PpMode {
         use PpMode::*;
         use PpSourceMode::*;
         match *self {
-            PpmSource(PpmNormal | PpmEveryBodyLoops | PpmIdentified) => false,
+            PpmSource(PpmNormal | PpmIdentified) => false,
 
-            PpmSource(PpmExpanded | PpmExpandedIdentified | PpmExpandedHygiene)
+            PpmSource(
+                PpmExpanded | PpmEveryBodyLoops | PpmExpandedIdentified | PpmExpandedHygiene,
+            )
             | PpmHir(_)
             | PpmHirTree(_)
             | PpmMir

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,11 @@ fn main() {
     || {
         macro m() {}
     };
+
+    let _ = || {
+        macro n() {}
+    };
+
+    let cond = true;
+    let _ = || if cond { macro n() {} } else { panic!() };
 }