Auto merge of #2322 - saethlin:stack-inspection-tools, r=oli-obk

Ideas on getting information about borrow stacks during execution

From time to time people ask what some borrow stack looks like in some code. I just know that I am terrible at doing Stacked Borrows by hand, so I always toss together something like this.

I know that Miri has logging, but I've never found it particularly useful because there's just too much output. Also I personally don't think about exactly what the state of a borrow stack is, but this seems to be something that newcomers to Stacked Borrows always want.

Update: This has been sitting as S-waiting-on-author for a long time. I bring it out from time to time to explain Stacked Borrows to people, and just now `@JakobDegen` said
>  Can we please merge that btw? It's such a valuable teaching tool
>  Interfaces can be fixed later

I'm inclined to trust Jake's judgement here.

7 files changed, 79 insertions, 0 deletions

diff --git a/src/tools/miri/README.md b/src/tools/miri/README.md
index 7ad51520a42..32e616cb074 100644
--- a/src/tools/miri/README.md
+++ b/src/tools/miri/README.md
@@ -531,6 +531,17 @@ extern "Rust" {
     /// This is internal and unstable and should not be used; we give it here
     /// just to be complete.
     fn miri_start_panic(payload: *mut u8) -> !;
+
+    /// Miri-provided extern function to get the internal unique identifier for the allocation that a pointer
+    /// points to. This is only useful as an input to `miri_print_stacks`, and it is a separate call because
+    /// getting a pointer to an allocation at runtime can change the borrow stacks in the allocation.
+    fn miri_get_alloc_id(ptr: *const ()) -> u64;
+
+    /// Miri-provided extern function to print (from the interpreter, not the program) the contents of all
+    /// borrow stacks in an allocation. The format of what this emits is unstable and may change at any time.
+    /// In particular, users should be aware that Miri will periodically attempt to garbage collect the
+    /// contents of all stacks. Callers of this function may wish to pass `-Zmiri-tag-gc=0` to disable the GC.
+    fn miri_print_stacks(alloc_id: u64);
 }
 ```
 
diff --git a/src/tools/miri/src/range_map.rs b/src/tools/miri/src/range_map.rs
index c77ea63b087..4742a365ec3 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 26184fdc3c0..7e6a9595161 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 09d36ca9dfd..f1dd38e5fc1 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(())
+    }
 }
diff --git a/src/tools/miri/tests/compiletest.rs b/src/tools/miri/tests/compiletest.rs
index c80ec3625ef..65fd7c2eccb 100644
--- a/src/tools/miri/tests/compiletest.rs
+++ b/src/tools/miri/tests/compiletest.rs
@@ -139,6 +139,8 @@ regexes! {
     STDOUT:
     // Windows file paths
     r"\\"                           => "/",
+    // erase Stacked Borrows tags
+    "<[0-9]+>"                      => "<TAG>",
 }
 
 regexes! {
diff --git a/src/tools/miri/tests/pass/stacked-borrows/stack-printing.rs b/src/tools/miri/tests/pass/stacked-borrows/stack-printing.rs
new file mode 100644
index 00000000000..8d96a2e1ca9
--- /dev/null
+++ b/src/tools/miri/tests/pass/stacked-borrows/stack-printing.rs
@@ -0,0 +1,29 @@
+use std::{
+    alloc::{self, Layout},
+    mem::ManuallyDrop,
+};
+
+extern "Rust" {
+    fn miri_get_alloc_id(ptr: *const u8) -> u64;
+    fn miri_print_stacks(alloc_id: u64);
+}
+
+fn main() {
+    let ptr = unsafe { alloc::alloc(Layout::new::<u8>()) };
+    let alloc_id = unsafe { miri_get_alloc_id(ptr) };
+    unsafe { miri_print_stacks(alloc_id) };
+
+    assert!(!ptr.is_null());
+    unsafe { miri_print_stacks(alloc_id) };
+
+    unsafe { *ptr = 42 };
+    unsafe { miri_print_stacks(alloc_id) };
+
+    let _b = unsafe { ManuallyDrop::new(Box::from_raw(ptr)) };
+    unsafe { miri_print_stacks(alloc_id) };
+
+    let _ptr = unsafe { &*ptr };
+    unsafe { miri_print_stacks(alloc_id) };
+
+    unsafe { alloc::dealloc(ptr, Layout::new::<u8>()) };
+}
diff --git a/src/tools/miri/tests/pass/stacked-borrows/stack-printing.stdout b/src/tools/miri/tests/pass/stacked-borrows/stack-printing.stdout
new file mode 100644
index 00000000000..660ee71e6f5
--- /dev/null
+++ b/src/tools/miri/tests/pass/stacked-borrows/stack-printing.stdout
@@ -0,0 +1,5 @@
+0..1: [ SharedReadWrite<TAG> ]
+0..1: [ SharedReadWrite<TAG> ]
+0..1: [ SharedReadWrite<TAG> ]
+0..1: [ SharedReadWrite<TAG> Unique<TAG> Unique<TAG> Unique<TAG> Unique<TAG> Unique<TAG> ]
+0..1: [ SharedReadWrite<TAG> Disabled<TAG> Disabled<TAG> Disabled<TAG> Disabled<TAG> Disabled<TAG> SharedReadOnly<TAG> ]