Fix ICE with `-Zunpretty=hir,typed` when an expression occurs in a function signature

1 files changed, 18 insertions, 16 deletions

diff --git a/compiler/rustc_driver/src/pretty.rs b/compiler/rustc_driver/src/pretty.rs
index a2b4f3fcf73..2a11d62ac15 100644
--- a/compiler/rustc_driver/src/pretty.rs
+++ b/compiler/rustc_driver/src/pretty.rs
@@ -293,18 +293,6 @@ struct TypedAnnotation<'tcx> {
     maybe_typeck_results: Cell<Option<&'tcx ty::TypeckResults<'tcx>>>,
 }
 
-impl<'tcx> TypedAnnotation<'tcx> {
-    /// Gets the type-checking results for the current body.
-    /// As this will ICE if called outside bodies, only call when working with
-    /// `Expr` or `Pat` nodes (they are guaranteed to be found only in bodies).
-    #[track_caller]
-    fn typeck_results(&self) -> &'tcx ty::TypeckResults<'tcx> {
-        self.maybe_typeck_results
-            .get()
-            .expect("`TypedAnnotation::typeck_results` called outside of body")
-    }
-}
-
 impl<'tcx> HirPrinterSupport<'tcx> for TypedAnnotation<'tcx> {
     fn sess(&self) -> &Session {
         &self.tcx.sess
@@ -336,10 +324,24 @@ impl<'tcx> pprust_hir::PpAnn for TypedAnnotation<'tcx> {
     }
     fn post(&self, s: &mut pprust_hir::State<'_>, node: pprust_hir::AnnNode<'_>) {
         if let pprust_hir::AnnNode::Expr(expr) = node {
-            s.s.space();
-            s.s.word("as");
-            s.s.space();
-            s.s.word(self.typeck_results().expr_ty(expr).to_string());
+            let typeck_results =
+                self.maybe_typeck_results.get().or_else(|| {
+                    if let Some(body_id) = self.tcx.hir().maybe_body_owned_by(
+                        self.tcx.hir().local_def_id_to_hir_id(expr.hir_id.owner),
+                    ) {
+                        Some(self.tcx.typeck_body(body_id))
+                    } else {
+                        None
+                    }
+                });
+
+            if let Some(typeck_results) = typeck_results {
+                s.s.space();
+                s.s.word("as");
+                s.s.space();
+                s.s.word(typeck_results.expr_ty(expr).to_string());
+            }
+
             s.pclose();
         }
     }