rustc_driver/pretty: avoid using TypeckTables::empty for TypedAnnotation.

2 files changed, 20 insertions, 10 deletions

diff --git a/src/librustc_driver/lib.rs b/src/librustc_driver/lib.rs
index b45ab0f80ff..c2c19b6b405 100644
--- a/src/librustc_driver/lib.rs
+++ b/src/librustc_driver/lib.rs
@@ -6,6 +6,7 @@
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/")]
 #![feature(nll)]
+#![cfg_attr(bootstrap, feature(track_caller))]
 #![recursion_limit = "256"]
 
 #[macro_use]
diff --git a/src/librustc_driver/pretty.rs b/src/librustc_driver/pretty.rs
index 0a21eb8de05..6a34a310f73 100644
--- a/src/librustc_driver/pretty.rs
+++ b/src/librustc_driver/pretty.rs
@@ -80,8 +80,7 @@ where
         PpmTyped => {
             abort_on_err(tcx.analysis(LOCAL_CRATE), tcx.sess);
 
-            let empty_tables = ty::TypeckTables::empty(None);
-            let annotation = TypedAnnotation { tcx, tables: Cell::new(&empty_tables) };
+            let annotation = TypedAnnotation { tcx, maybe_typeck_tables: Cell::new(None) };
             tcx.dep_graph.with_ignore(|| f(&annotation, tcx.hir().krate()))
         }
         _ => panic!("Should use call_with_pp_support"),
@@ -304,12 +303,22 @@ impl<'a> pprust::PpAnn for HygieneAnnotation<'a> {
     }
 }
 
-struct TypedAnnotation<'a, 'tcx> {
+struct TypedAnnotation<'tcx> {
     tcx: TyCtxt<'tcx>,
-    tables: Cell<&'a ty::TypeckTables<'tcx>>,
+    maybe_typeck_tables: Cell<Option<&'tcx ty::TypeckTables<'tcx>>>,
 }
 
-impl<'b, 'tcx> HirPrinterSupport<'tcx> for TypedAnnotation<'b, 'tcx> {
+impl<'tcx> TypedAnnotation<'tcx> {
+    /// Gets the type-checking side-tables 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 tables(&self) -> &'tcx ty::TypeckTables<'tcx> {
+        self.maybe_typeck_tables.get().expect("`TypedAnnotation::tables` called outside of body")
+    }
+}
+
+impl<'tcx> HirPrinterSupport<'tcx> for TypedAnnotation<'tcx> {
     fn sess(&self) -> &Session {
         &self.tcx.sess
     }
@@ -327,15 +336,15 @@ impl<'b, 'tcx> HirPrinterSupport<'tcx> for TypedAnnotation<'b, 'tcx> {
     }
 }
 
-impl<'a, 'tcx> pprust_hir::PpAnn for TypedAnnotation<'a, 'tcx> {
+impl<'tcx> pprust_hir::PpAnn for TypedAnnotation<'tcx> {
     fn nested(&self, state: &mut pprust_hir::State<'_>, nested: pprust_hir::Nested) {
-        let old_tables = self.tables.get();
+        let old_maybe_typeck_tables = self.maybe_typeck_tables.get();
         if let pprust_hir::Nested::Body(id) = nested {
-            self.tables.set(self.tcx.body_tables(id));
+            self.maybe_typeck_tables.set(Some(self.tcx.body_tables(id)));
         }
         let pp_ann = &(&self.tcx.hir() as &dyn hir::intravisit::Map<'_>);
         pprust_hir::PpAnn::nested(pp_ann, state, nested);
-        self.tables.set(old_tables);
+        self.maybe_typeck_tables.set(old_maybe_typeck_tables);
     }
     fn pre(&self, s: &mut pprust_hir::State<'_>, node: pprust_hir::AnnNode<'_>) {
         if let pprust_hir::AnnNode::Expr(_) = node {
@@ -347,7 +356,7 @@ impl<'a, 'tcx> pprust_hir::PpAnn for TypedAnnotation<'a, 'tcx> {
             s.s.space();
             s.s.word("as");
             s.s.space();
-            s.s.word(self.tables.get().expr_ty(expr).to_string());
+            s.s.word(self.tables().expr_ty(expr).to_string());
             s.pclose();
         }
     }