share more alloc printing code between Miri and MIR dumping

Author: RalfJung

src/librustc_mir/interpret/memory.rs

@@ -20,6 +20,7 @@ use super::{
     AllocId, AllocMap, Allocation, AllocationExtra, CheckInAllocMsg, ErrorHandled, GlobalAlloc,
     GlobalId, InterpResult, Machine, MayLeak, Pointer, PointerArithmetic, Scalar,
 };
+use crate::util::pretty;
 
 #[derive(Debug, PartialEq, Copy, Clone)]
 pub enum MemoryKind<T> {
@@ -644,22 +645,22 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
         self.dump_allocs(vec![id]);
     }
 
-    fn dump_alloc_helper<Tag, Extra>(
-        &self,
-        allocs_to_print: &mut VecDeque<AllocId>,
-        alloc: &Allocation<Tag, Extra>,
-    ) {
-        for &(_, target_id) in alloc.relocations().values() {
-            allocs_to_print.push_back(target_id);
-        }
-        crate::util::pretty::write_allocation(self.tcx.tcx, alloc, &mut std::io::stderr(), "")
-            .unwrap();
-    }
-
     /// Print a list of allocations and all allocations they point to, recursively.
     /// This prints directly to stderr, ignoring RUSTC_LOG! It is up to the caller to
     /// control for this.
     pub fn dump_allocs(&self, mut allocs: Vec<AllocId>) {
+        // Cannot be a closure because it is generic in `Tag`, `Extra`.
+        fn write_allocation_track_relocs<'tcx, Tag, Extra>(
+            tcx: TyCtxtAt<'tcx>,
+            allocs_to_print: &mut VecDeque<AllocId>,
+            alloc: &Allocation<Tag, Extra>,
+        ) {
+            for &(_, target_id) in alloc.relocations().values() {
+                allocs_to_print.push_back(target_id);
+            }
+            pretty::write_allocation(tcx.tcx, alloc, &mut std::io::stderr()).unwrap();
+        }
+
         allocs.sort();
         allocs.dedup();
         let mut allocs_to_print = VecDeque::from(allocs);
@@ -671,46 +672,42 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
                 // Already printed, so skip this.
                 continue;
             }
-            eprint!("Alloc {:<5}: ", id);
-            fn msg<Tag, Extra>(alloc: &Allocation<Tag, Extra>, extra: &str) {
-                eprintln!(
-                    "({} bytes, alignment {}){}",
-                    alloc.size.bytes(),
-                    alloc.align.bytes(),
-                    extra
-                )
-            };
 
-            // normal alloc?
-            match self.alloc_map.get_or(id, || Err(())) {
-                Ok((kind, alloc)) => {
+            eprint!("{}", id);
+            match self.alloc_map.get(id) {
+                Some(&(kind, ref alloc)) => {
+                    // normal alloc
                     match kind {
-                        MemoryKind::Stack => msg(alloc, " (stack)"),
-                        MemoryKind::Vtable => msg(alloc, " (vtable)"),
-                        MemoryKind::CallerLocation => msg(alloc, " (caller_location)"),
-                        MemoryKind::Machine(m) => msg(alloc, &format!(" ({:?})", m)),
+                        MemoryKind::Stack => eprint!(" (stack variable, "),
+                        MemoryKind::Vtable => eprint!(" (vtable, "),
+                        MemoryKind::CallerLocation => eprint!(" (caller_location, "),
+                        MemoryKind::Machine(m) if Some(m) == M::GLOBAL_KIND => {
+                            eprint!(" (global, ")
+                        }
+                        MemoryKind::Machine(m) => eprint!(" ({:?}, ", m),
                     };
-                    self.dump_alloc_helper(&mut allocs_to_print, alloc);
+                    write_allocation_track_relocs(self.tcx, &mut allocs_to_print, alloc);
                 }
-                Err(()) => {
-                    // global alloc?
+                None => {
+                    // global alloc
                     match self.tcx.alloc_map.lock().get(id) {
                         Some(GlobalAlloc::Memory(alloc)) => {
-                            msg(alloc, " (immutable)");
-                            self.dump_alloc_helper(&mut allocs_to_print, alloc);
+                            eprint!(" (global, ");
+                            write_allocation_track_relocs(self.tcx, &mut allocs_to_print, alloc);
                         }
                         Some(GlobalAlloc::Function(func)) => {
-                            eprintln!("{}", func);
+                            eprint!(" (fn: {})", func);
                         }
                         Some(GlobalAlloc::Static(did)) => {
-                            eprintln!("{:?}", did);
+                            eprint!(" (static: {})", self.tcx.def_path_str(did));
                         }
                         None => {
-                            eprintln!("(deallocated)");
+                            eprint!(" (deallocated)");
                         }
                     }
                 }
-            };
+            }
+            eprintln!();
         }
     }
 

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,