--emit=mir now emits both `mir_for_ctfe` and `optimized_mir` for `const fn`

4 files changed, 86 insertions, 11 deletions

diff --git a/compiler/rustc_mir/src/util/pretty.rs b/compiler/rustc_mir/src/util/pretty.rs
index ca8bcffa896..7fc1c3a73af 100644
--- a/compiler/rustc_mir/src/util/pretty.rs
+++ b/compiler/rustc_mir/src/util/pretty.rs
@@ -273,13 +273,6 @@ pub fn write_mir_pretty<'tcx>(
 
     let mut first = true;
     for def_id in dump_mir_def_ids(tcx, single) {
-        let body = match tcx.hir().body_const_context(def_id.expect_local()) {
-            // For `const fn` we want to render the optimized MIR. If you want the mir used in
-            // ctfe, you can dump the MIR after the `Deaggregator` optimization pass.
-            None | Some(rustc_hir::ConstContext::ConstFn) => tcx.optimized_mir(def_id),
-            Some(_) => tcx.mir_for_ctfe(def_id),
-        };
-
         if first {
             first = false;
         } else {
@@ -287,11 +280,28 @@ pub fn write_mir_pretty<'tcx>(
             writeln!(w)?;
         }
 
-        write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
-
-        for body in tcx.promoted_mir(def_id) {
-            writeln!(w)?;
+        let render_body = |w: &mut dyn Write, body| -> io::Result<()> {
             write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
+
+            for body in tcx.promoted_mir(def_id) {
+                writeln!(w)?;
+                write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
+            }
+            Ok(())
+        };
+        match tcx.hir().body_const_context(def_id.expect_local()) {
+            None => render_body(w, tcx.optimized_mir(def_id))?,
+            // For `const fn` we want to render the optimized MIR. If you want the mir used in
+            // ctfe, you can dump the MIR after the `Deaggregator` optimization pass.
+            Some(rustc_hir::ConstContext::ConstFn) => {
+                render_body(w, tcx.optimized_mir(def_id))?;
+                writeln!(w)?;
+                writeln!(w, "// MIR FOR CTFE")?;
+                // Do not use `render_body`, as that would render the promoteds again, but these
+                // are shared between mir_for_ctfe and optimized_mir
+                write_mir_fn(tcx, tcx.mir_for_ctfe(def_id), &mut |_, _| Ok(()), w)?;
+            }
+            Some(_) => render_body(w, tcx.mir_for_ctfe(def_id))?,
         }
     }
     Ok(())
diff --git a/src/test/run-make/const_fn_mir/Makefile b/src/test/run-make/const_fn_mir/Makefile
new file mode 100644
index 00000000000..2aa0bc9d45d
--- /dev/null
+++ b/src/test/run-make/const_fn_mir/Makefile
@@ -0,0 +1,10 @@
+-include ../../run-make-fulldeps/tools.mk
+
+all:
+	$(RUSTC) main.rs --emit=mir -o "$(TMPDIR)"/dump.mir
+
+ifdef RUSTC_BLESS_TEST
+	cp "$(TMPDIR)"/dump.mir dump.mir
+else
+	$(DIFF) dump.mir "$(TMPDIR)"/dump.mir
+endif
diff --git a/src/test/run-make/const_fn_mir/dump.mir b/src/test/run-make/const_fn_mir/dump.mir
new file mode 100644
index 00000000000..3dac42c6782
--- /dev/null
+++ b/src/test/run-make/const_fn_mir/dump.mir
@@ -0,0 +1,45 @@
+// WARNING: This output format is intended for human consumers only
+// and is subject to change without notice. Knock yourself out.
+fn main() -> () {
+    let mut _0: ();                      // return place in scope 0 at main.rs:8:11: 8:11
+    let _1: i32;                         // in scope 0 at main.rs:9:5: 9:10
+
+    bb0: {
+        StorageLive(_1);                 // scope 0 at main.rs:9:5: 9:10
+        _1 = foo() -> bb1;               // scope 0 at main.rs:9:5: 9:10
+                                         // mir::Constant
+                                         // + span: main.rs:9:5: 9:8
+                                         // + literal: Const { ty: fn() -> i32 {foo}, val: Value(Scalar(<ZST>)) }
+    }
+
+    bb1: {
+        StorageDead(_1);                 // scope 0 at main.rs:9:10: 9:11
+        _0 = const ();                   // scope 0 at main.rs:8:11: 10:2
+        return;                          // scope 0 at main.rs:10:2: 10:2
+    }
+}
+
+fn foo() -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at main.rs:4:19: 4:22
+
+    bb0: {
+        _0 = const 11_i32;               // scope 0 at main.rs:5:5: 5:10
+        return;                          // scope 0 at main.rs:6:2: 6:2
+    }
+}
+
+// MIR FOR CTFE
+fn foo() -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at main.rs:4:19: 4:22
+    let mut _1: (i32, bool);             // in scope 0 at main.rs:5:5: 5:10
+
+    bb0: {
+        _1 = CheckedAdd(const 5_i32, const 6_i32); // scope 0 at main.rs:5:5: 5:10
+        assert(!move (_1.1: bool), "attempt to compute `{} + {}`, which would overflow", const 5_i32, const 6_i32) -> bb1; // scope 0 at main.rs:5:5: 5:10
+    }
+
+    bb1: {
+        _0 = move (_1.0: i32);           // scope 0 at main.rs:5:5: 5:10
+        return;                          // scope 0 at main.rs:6:2: 6:2
+    }
+}
diff --git a/src/test/run-make/const_fn_mir/main.rs b/src/test/run-make/const_fn_mir/main.rs
new file mode 100644
index 00000000000..e8552bd285a
--- /dev/null
+++ b/src/test/run-make/const_fn_mir/main.rs
@@ -0,0 +1,10 @@
+// emit-mir
+// check-pass
+
+const fn foo() -> i32 {
+    5 + 6
+}
+
+fn main() {
+    foo();
+}