From 79cd6f38f3546505f5e7a6e5ce4f696ea6b19d19 Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Sun, 3 Jul 2022 21:54:45 -0400
Subject: [PATCH] Implement a crude stack printing mechanism

---
 src/tools/miri/src/range_map.rs           |  4 ++++
 src/tools/miri/src/shims/foreign_items.rs | 13 +++++++++++++
 src/tools/miri/src/stacked_borrows/mod.rs | 15 +++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/src/tools/miri/src/range_map.rs b/src/tools/miri/src/range_map.rs
index c77ea63b0873..4742a365ec38 100644
--- a/src/tools/miri/src/range_map.rs
+++ b/src/tools/miri/src/range_map.rs
@@ -91,6 +91,10 @@ impl<T> RangeMap<T> {
         self.v.iter_mut().map(|elem| &mut elem.data)
     }
 
+    pub fn iter_all(&self) -> impl Iterator<Item = (ops::Range<u64>, &T)> {
+        self.v.iter().map(|elem| (elem.range.clone(), &elem.data))
+    }
+
     // Splits the element situated at the given `index`, such that the 2nd one starts at offset
     // `split_offset`. Do nothing if the element already starts there.
     // Returns whether a split was necessary.
diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs
index 26184fdc3c09..7e6a9595161f 100644
--- a/src/tools/miri/src/shims/foreign_items.rs
+++ b/src/tools/miri/src/shims/foreign_items.rs
@@ -417,6 +417,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // shim, add it to the corresponding submodule.
         match link_name.as_str() {
             // Miri-specific extern functions
+            "miri_get_alloc_id" => {
+                let [ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let ptr = this.read_pointer(ptr)?;
+                let (alloc_id, _, _) = this.ptr_get_alloc_id(ptr)?;
+                this.write_scalar(Scalar::from_u64(alloc_id.0.get()), dest)?;
+            }
+            "miri_print_stacks" => {
+                let [id] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let id = this.read_scalar(id)?.to_u64()?;
+                if let Some(id) = std::num::NonZeroU64::new(id) {
+                    this.print_stacks(AllocId(id))?;
+                }
+            }
             "miri_static_root" => {
                 let [ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;
                 let ptr = this.read_pointer(ptr)?;
diff --git a/src/tools/miri/src/stacked_borrows/mod.rs b/src/tools/miri/src/stacked_borrows/mod.rs
index 09d36ca9dfdb..f1dd38e5fc1e 100644
--- a/src/tools/miri/src/stacked_borrows/mod.rs
+++ b/src/tools/miri/src/stacked_borrows/mod.rs
@@ -1123,4 +1123,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
         Ok(())
     }
+
+    fn print_stacks(&mut self, alloc_id: AllocId) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let alloc_extra = this.get_alloc_extra(alloc_id)?;
+        let stacks = alloc_extra.stacked_borrows.as_ref().unwrap().borrow();
+        for (range, stack) in stacks.stacks.iter_all() {
+            print!("{:?}: [", range);
+            for i in 0..stack.len() {
+                let item = stack.get(i).unwrap();
+                print!(" {:?}{:?}", item.perm(), item.tag());
+            }
+            println!(" ]");
+        }
+        Ok(())
+    }
 }