2 files changed, 4 insertions, 9 deletions
diff --git a/compiler/rustc_mir_build/src/builder/mod.rs b/compiler/rustc_mir_build/src/builder/mod.rs
index 8400709740f..127b191e335 100644
--- a/compiler/rustc_mir_build/src/builder/mod.rs
+++ b/compiler/rustc_mir_build/src/builder/mod.rs
@@ -450,10 +450,6 @@ fn construct_fn<'tcx>(
     let span = tcx.def_span(fn_def);
     let fn_id = tcx.local_def_id_to_hir_id(fn_def);
 
-    // The representation of thir for `-Zunpretty=thir-tree` relies on
-    // the entry expression being the last element of `thir.exprs`.
-    assert_eq!(expr.as_usize(), thir.exprs.len() - 1);
-
     // Figure out what primary body this item has.
     let body = tcx.hir_body_owned_by(fn_def);
     let span_with_body = tcx.hir_span_with_body(fn_id);
diff --git a/compiler/rustc_mir_build/src/thir/print.rs b/compiler/rustc_mir_build/src/thir/print.rs
index 37248941e2c..db9547a481f 100644
--- a/compiler/rustc_mir_build/src/thir/print.rs
+++ b/compiler/rustc_mir_build/src/thir/print.rs
@@ -8,10 +8,10 @@ use rustc_span::def_id::LocalDefId;
 /// Create a THIR tree for debugging.
 pub fn thir_tree(tcx: TyCtxt<'_>, owner_def: LocalDefId) -> String {
     match super::cx::thir_body(tcx, owner_def) {
-        Ok((thir, _)) => {
+        Ok((thir, expr)) => {
             let thir = thir.steal();
             let mut printer = ThirPrinter::new(&thir);
-            printer.print();
+            printer.print(expr);
             printer.into_buffer()
         }
         Err(_) => "error".into(),
@@ -58,7 +58,7 @@ impl<'a, 'tcx> ThirPrinter<'a, 'tcx> {
         }
     }
 
-    fn print(&mut self) {
+    fn print(&mut self, body_expr: ExprId) {
         print_indented!(self, "params: [", 0);
         for param in self.thir.params.iter() {
             self.print_param(param, 1);
@@ -66,8 +66,7 @@ impl<'a, 'tcx> ThirPrinter<'a, 'tcx> {
         print_indented!(self, "]", 0);
 
         print_indented!(self, "body:", 0);
-        let expr = ExprId::from_usize(self.thir.exprs.len() - 1);
-        self.print_expr(expr, 1);
+        self.print_expr(body_expr, 1);
     }
 
     fn into_buffer(self) -> String {