share more alloc printing code between Miri and MIR dumping

1 files changed, 44 insertions, 26 deletions

diff --git a/src/librustc_mir/util/pretty.rs b/src/librustc_mir/util/pretty.rs
index a81fcb54580..64221c41bff 100644
--- a/src/librustc_mir/util/pretty.rs
+++ b/src/librustc_mir/util/pretty.rs
@@ -567,26 +567,21 @@ pub fn write_allocations<'tcx>(
     }
     let mut visitor = CollectAllocIds(Default::default());
     body.visit_with(&mut visitor);
+    // `seen` contains all seen allocations, including the ones we have *not* printed yet.
+    // The protocol is to first `insert` into `seen`, and only if that returns `true`
+    // then push to `todo`.
     let mut seen = visitor.0;
     let mut todo: Vec<_> = seen.iter().copied().collect();
     while let Some(id) = todo.pop() {
-        let mut write_header_and_allocation =
+        let mut write_allocation_track_relocs =
             |w: &mut dyn Write, alloc: &Allocation| -> io::Result<()> {
-                write!(w, "size: {}, align: {})", alloc.size.bytes(), alloc.align.bytes())?;
-                if alloc.size == Size::ZERO {
-                    write!(w, " {{}}")?;
-                } else {
-                    writeln!(w, " {{")?;
-                    write_allocation(tcx, alloc, w, "    ")?;
-                    write!(w, "}}")?;
-                    // `.rev()` because we are popping them from the back of the `todo` vector.
-                    for id in alloc_ids_from_alloc(alloc).rev() {
-                        if seen.insert(id) {
-                            todo.push(id);
-                        }
+                // `.rev()` because we are popping them from the back of the `todo` vector.
+                for id in alloc_ids_from_alloc(alloc).rev() {
+                    if seen.insert(id) {
+                        todo.push(id);
                     }
                 }
-                Ok(())
+                write_allocation(tcx, alloc, w)
             };
         write!(w, "\n{}", id)?;
         let alloc = tcx.alloc_map.lock().get(id);
@@ -599,7 +594,7 @@ pub fn write_allocations<'tcx>(
                 match tcx.const_eval_poly(did) {
                     Ok(ConstValue::ByRef { alloc, .. }) => {
                         write!(w, " (static: {}, ", tcx.def_path_str(did))?;
-                        write_header_and_allocation(w, alloc)?;
+                        write_allocation_track_relocs(w, alloc)?;
                     }
                     Ok(_) => {
                         span_bug!(tcx.def_span(did), " static item without `ByRef` initializer")
@@ -616,15 +611,46 @@ pub fn write_allocations<'tcx>(
             }
             Some(GlobalAlloc::Memory(alloc)) => {
                 write!(w, " (")?;
-                write_header_and_allocation(w, alloc)?
+                write_allocation_track_relocs(w, alloc)?
             }
         }
-
         writeln!(w)?;
     }
     Ok(())
 }
 
+/// Dumps the size and metadata and content of an allocation to the given writer.
+/// The expectation is that the caller first prints other relevant metadata, so the exact
+/// format of this function is (*without* leading or trailing newline):
+/// ```
+/// size: {}, align: {}) {
+///     <bytes>
+/// }
+/// ```
+///
+/// The byte format is similar to how hex editors print bytes. Each line starts with the address of
+/// the start of the line, followed by all bytes in hex format (space separated).
+/// If the allocation is small enough to fit into a single line, no start address is given.
+/// After the hex dump, an ascii dump follows, replacing all unprintable characters (control
+/// characters or characters whose value is larger than 127) with a `.`
+/// This also prints relocations adequately.
+pub fn write_allocation<Tag, Extra>(
+    tcx: TyCtxt<'tcx>,
+    alloc: &Allocation<Tag, Extra>,
+    w: &mut dyn Write,
+) -> io::Result<()> {
+    write!(w, "size: {}, align: {})", alloc.size.bytes(), alloc.align.bytes())?;
+    if alloc.size == Size::ZERO {
+        // We are done.
+        return write!(w, " {{}}");
+    }
+    // Write allocation bytes.
+    writeln!(w, " {{")?;
+    write_allocation_bytes(tcx, alloc, w, "    ")?;
+    write!(w, "}}")?;
+    Ok(())
+}
+
 fn write_allocation_endline(w: &mut dyn Write, ascii: &str) -> io::Result<()> {
     for _ in 0..(BYTES_PER_LINE - ascii.chars().count()) {
         write!(w, "   ")?;
@@ -649,18 +675,10 @@ fn write_allocation_newline(
     Ok(line_start)
 }
 
-/// Dumps the bytes of an allocation to the given writer. This also prints relocations instead of
-/// the raw bytes where applicable.
-/// The byte format is similar to how hex editors print bytes. Each line starts with the address of
-/// the start of the line, followed by all bytes in hex format (space separated).
-/// If the allocation is small enough to fit into a single line, no start address is given.
-/// After the hex dump, an ascii dump follows, replacing all unprintable characters (control
-/// characters or characters whose value is larger than 127) with a `.`
-///
 /// The `prefix` argument allows callers to add an arbitrary prefix before each line (even if there
 /// is only one line). Note that your prefix should contain a trailing space as the lines are
 /// printed directly after it.
-pub fn write_allocation<Tag, Extra>(
+fn write_allocation_bytes<Tag, Extra>(
     tcx: TyCtxt<'tcx>,
     alloc: &Allocation<Tag, Extra>,
     w: &mut dyn Write,