Provenance forging for miri_pointer_name

src/borrow_tracker/mod.rs

@@ -318,7 +318,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     fn give_pointer_debug_name(
         &mut self,
         ptr: Pointer<Option<Provenance>>,
-        nth_parent: u8,
         name: &str,
     ) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
@@ -328,8 +327,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 this.tcx.tcx.sess.warn("Stacked Borrows does not support named pointers; `miri_pointer_name` is a no-op");
                 Ok(())
             }
-            BorrowTrackerMethod::TreeBorrows =>
-                this.tb_give_pointer_debug_name(ptr, nth_parent, name),
+            BorrowTrackerMethod::TreeBorrows => this.tb_give_pointer_debug_name(ptr, name),
         }
     }
 
@@ -341,6 +339,38 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             BorrowTrackerMethod::TreeBorrows => this.print_tree(alloc_id, show_unnamed),
         }
     }
+
+    /// SB: returns `ptr`.
+    /// TB: forges a pointer with the provenance of the `nb`'th parent of `ptr`
+    /// in the same allocation.
+    fn forge_ptr_nth_parent(
+        &mut self,
+        ptr: Pointer<Option<Provenance>>,
+        nb: u8,
+    ) -> InterpResult<'tcx, Pointer<Option<Provenance>>> {
+        let this = self.eval_context_mut();
+        let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
+        match method {
+            BorrowTrackerMethod::StackedBorrows => Ok(ptr),
+            BorrowTrackerMethod::TreeBorrows => this.tb_forge_ptr_nth_parent(ptr, nb),
+        }
+    }
+
+    /// SB: returns `ptr1`.
+    /// TB: forges a pointer with the provenance of the nearest common ancestor of
+    /// `ptr1` and `ptr2` if they belong to the same allocation.
+    fn forge_ptr_common_ancestor(
+        &mut self,
+        ptr1: Pointer<Option<Provenance>>,
+        ptr2: Pointer<Option<Provenance>>,
+    ) -> InterpResult<'tcx, Pointer<Option<Provenance>>> {
+        let this = self.eval_context_mut();
+        let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
+        match method {
+            BorrowTrackerMethod::StackedBorrows => Ok(ptr1),
+            BorrowTrackerMethod::TreeBorrows => this.tb_forge_ptr_common_ancestor(ptr1, ptr2),
+        }
+    }
 }
 
 /// Extra per-allocation data for borrow tracking

src/borrow_tracker/tree_borrows/diagnostics.rs

@@ -1,7 +1,7 @@
 use std::fmt;
 use std::ops::Range;
 
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_span::{Span, SpanData};
 
 use crate::borrow_tracker::tree_borrows::{
@@ -181,35 +181,67 @@ impl fmt::Display for NodeDebugInfo {
     }
 }
 
-impl<'tcx> Tree {
+impl Tree {
     /// Climb the tree to get the tag of a distant ancestor.
     /// Allows operations on tags that are unreachable by the program
     /// but still exist in the tree. Not guaranteed to perform consistently
     /// if `tag-gc=1`.
-    fn nth_parent(&self, tag: BorTag, nth_parent: u8) -> Option<BorTag> {
-        let mut idx = self.tag_mapping.get(&tag).unwrap();
-        for _ in 0..nth_parent {
+    ///
+    /// NOTE: the `nb`'th parent is computed in a "saturating" way: `nth_parent(root, 1) == root`
+    pub fn nth_parent(&self, tag: BorTag, nth_parent: u8) -> BorTag {
+        let mut idx = self.tag_mapping.get(&tag).expect("Tag not found in the allocation");
+        for dist in 0..nth_parent {
             let node = self.nodes.get(idx).unwrap();
-            idx = node.parent?;
+            let Some(parent) = node.parent else {
+                eprintln!("{tag:?} has no {nth_parent}'th parent (reached the root) at distance {dist}, exiting early");
+                break;
+            };
+            idx = parent;
         }
-        Some(self.nodes.get(idx).unwrap().tag)
+        self.nodes.get(idx).unwrap().tag
     }
 
     /// Debug helper: assign name to tag.
-    pub fn give_pointer_debug_name(
-        &mut self,
-        tag: BorTag,
-        nth_parent: u8,
-        name: &str,
-    ) -> InterpResult<'tcx> {
-        let tag = self.nth_parent(tag, nth_parent).unwrap();
+    pub fn give_pointer_debug_name(&mut self, tag: BorTag, name: &str) {
         let idx = self.tag_mapping.get(&tag).unwrap();
         if let Some(node) = self.nodes.get_mut(idx) {
             node.debug_info.add_name(name);
         } else {
             eprintln!("Tag {tag:?} (to be named '{name}') not found!");
         }
-        Ok(())
+    }
+
+    /// Find the nearest common ancestor of two tags.
+    /// This function will panic if the tags do not exist or are not part of the
+    /// correct allocation.
+    /// Unlike `nth_parent` this should be more resilient to garbage collection
+    /// and implementation details.
+    pub fn common_ancestor(&self, tag1: BorTag, tag2: BorTag) -> BorTag {
+        // To find the nearest common ancestor without paying too much
+        // on pathological cases: climb both paths in parallel.
+        let mut path1 = FxHashSet::default();
+        let mut path2 = FxHashSet::default();
+        let mut idx1 = self.tag_mapping.get(&tag1).expect("Tag not found in the allocation");
+        let mut idx2 = self.tag_mapping.get(&tag2).expect("Tag not found in the allocation");
+        loop {
+            // Record path
+            path1.insert(idx1);
+            path2.insert(idx2);
+            // Check membership of one index in the other path: if so we found the ancestor.
+            if path2.contains(&idx1) {
+                return self.nodes.get(idx1).unwrap().tag;
+            }
+            if path1.contains(&idx2) {
+                return self.nodes.get(idx2).unwrap().tag;
+            }
+            // Not found yet, climb up one step.
+            if let Some(parent1) = self.nodes.get(idx1).unwrap().parent {
+                idx1 = parent1;
+            }
+            if let Some(parent2) = self.nodes.get(idx2).unwrap().parent {
+                idx2 = parent2;
+            }
+        }
     }
 
     /// Debug helper: determines if the tree contains a tag.

src/borrow_tracker/tree_borrows/mod.rs

@@ -536,7 +536,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     fn tb_give_pointer_debug_name(
         &mut self,
         ptr: Pointer<Option<Provenance>>,
-        nth_parent: u8,
         name: &str,
     ) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
@@ -549,6 +548,66 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         };
         let alloc_extra = this.get_alloc_extra(alloc_id)?;
         let mut tree_borrows = alloc_extra.borrow_tracker_tb().borrow_mut();
-        tree_borrows.give_pointer_debug_name(tag, nth_parent, name)
+        tree_borrows.give_pointer_debug_name(tag, name);
+        Ok(())
+    }
+
+    /// Forge a pointer with the provenance of the `nb`'th parent of `ptr`.
+    /// No-op if `ptr` is a wildcard pointer.
+    fn tb_forge_ptr_nth_parent(
+        &mut self,
+        ptr: Pointer<Option<Provenance>>,
+        nb: u8,
+    ) -> InterpResult<'tcx, Pointer<Option<Provenance>>> {
+        let this = self.eval_context_mut();
+        let (tag, alloc_id) = match ptr.provenance {
+            Some(Provenance::Concrete { tag, alloc_id }) => (tag, alloc_id),
+            _ => {
+                eprintln!("{ptr:?} has no provenance, cannot forge ancestor");
+                return Ok(ptr);
+            }
+        };
+        let alloc_extra = this.get_alloc_extra(alloc_id)?;
+        let tree_borrows = alloc_extra.borrow_tracker_tb().borrow();
+        let forged_tag = tree_borrows.nth_parent(tag, nb);
+        Ok(ptr.map_provenance(|_| Some(Provenance::Concrete { tag: forged_tag, alloc_id })))
+    }
+
+    /// Forge a pointer with the provenance of the nearest common ancestor of `ptr1` and `ptr2`.
+    /// No-op on `ptr1` if either pointer is a wildcard pointer or if they belong to different
+    /// allocations.
+    fn tb_forge_ptr_common_ancestor(
+        &mut self,
+        ptr1: Pointer<Option<Provenance>>,
+        ptr2: Pointer<Option<Provenance>>,
+    ) -> InterpResult<'tcx, Pointer<Option<Provenance>>> {
+        let this = self.eval_context_mut();
+        let (tag1, alloc_id1) = match ptr1.provenance {
+            Some(Provenance::Concrete { tag, alloc_id }) => (tag, alloc_id),
+            _ => {
+                eprintln!("{ptr1:?} has no provenance, cannot forge ancestor");
+                return Ok(ptr1);
+            }
+        };
+        let (tag2, alloc_id2) = match ptr2.provenance {
+            Some(Provenance::Concrete { tag, alloc_id }) => (tag, alloc_id),
+            _ => {
+                eprintln!("{ptr2:?} has no provenance, cannot forge ancestor");
+                return Ok(ptr1);
+            }
+        };
+        if alloc_id1 != alloc_id2 {
+            eprintln!(
+                "{tag1:?} and {tag2:?} are not part of the same allocation, cannot forge ancestor"
+            );
+            return Ok(ptr1);
+        }
+
+        let alloc_extra = this.get_alloc_extra(alloc_id1)?;
+        let tree_borrows = alloc_extra.borrow_tracker_tb().borrow();
+        let forged_tag = tree_borrows.common_ancestor(tag1, tag2);
+        Ok(ptr1.map_provenance(|_| {
+            Some(Provenance::Concrete { tag: forged_tag, alloc_id: alloc_id1 })
+        }))
     }
 }

src/borrow_tracker/tree_borrows/unimap.rs

@@ -17,7 +17,7 @@ use std::hash::Hash;
 use rustc_data_structures::fx::FxHashMap;
 
 /// Intermediate key between a UniKeyMap and a UniValMap.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct UniIndex {
     idx: u32,
 }

src/shims/foreign_items.rs

@@ -455,15 +455,28 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             "miri_pointer_name" => {
                 // This associates a name to a tag. Very useful for debugging, and also makes
                 // tests more strict.
-                let [ptr, nth_parent, name] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let [ptr, name] = this.check_shim(abi, Abi::Rust, link_name, args)?;
                 let ptr = this.read_pointer(ptr)?;
-                let nth_parent = this.read_scalar(nth_parent)?.to_u8()?;
                 let name = this.read_byte_slice(name)?;
                 // We must make `name` owned because we need to
                 // end the shared borrow from `read_byte_slice` before we can
                 // start the mutable borrow for `give_pointer_debug_name`.
                 let name = String::from_utf8_lossy(name).into_owned();
-                this.give_pointer_debug_name(ptr, nth_parent, &name)?;
+                this.give_pointer_debug_name(ptr, &name)?;
+            }
+            "miri_tree_nth_parent" => {
+                let [ptr, nb] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let ptr = this.read_pointer(ptr)?;
+                let nb = this.read_scalar(nb)?.to_u8()?;
+                let res = this.forge_ptr_nth_parent(ptr, nb)?;
+                this.write_pointer(res, dest)?;
+            }
+            "miri_tree_common_ancestor" => {
+                let [ptr1, ptr2] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let ptr1 = this.read_pointer(ptr1)?;
+                let ptr2 = this.read_pointer(ptr2)?;
+                let res = this.forge_ptr_common_ancestor(ptr1, ptr2)?;
+                this.write_pointer(res, dest)?;
             }
             "miri_static_root" => {
                 let [ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;

tests/fail/tree_borrows/reserved/cell-protected-write.rs

@@ -19,8 +19,8 @@ fn main() {
         write_second(x, y);
         unsafe fn write_second(x: &mut UnsafeCell<u8>, y: *mut u8) {
             let alloc_id = alloc_id!(x.get());
-            name!(x.get(), "callee:x");
-            name!(x.get()=>1, "caller:x");
+            name!(ancestor!(x.get(), x.get()), "callee:x");
+            name!(parent!(ancestor!(x.get(), x.get())), "caller:x");
             name!(y, "callee:y");
             name!(y, "caller:y");
             print_state!(alloc_id);

tests/fail/tree_borrows/reserved/int-protected-write.rs

@@ -18,7 +18,7 @@ fn main() {
         unsafe fn write_second(x: &mut u8, y: *mut u8) {
             let alloc_id = alloc_id!(x);
             name!(x, "callee:x");
-            name!(x=>1, "caller:x");
+            name!(parent!(x), "caller:x");
             name!(y, "callee:y");
             name!(y, "caller:y");
             print_state!(alloc_id);

tests/pass/tree_borrows/end-of-protector.rs

@@ -26,7 +26,7 @@ fn main() {
 
 unsafe fn do_nothing(x: &mut u8) {
     name!(x, "callee:x");
-    name!(x=>1, "caller:x");
+    name!(parent!(x), "caller:x");
     let alloc_id = alloc_id!(x);
     print_state!(alloc_id);
 }

tests/pass/tree_borrows/formatting.rs

@@ -16,7 +16,7 @@ fn main() {
 // decimal digits to verify proper padding.
 unsafe fn alignment_check() {
     let data: &mut [u8] = &mut [0; 1024];
-    name!(data.as_ptr()=>2, "data");
+    name!(ancestor!(data.as_ptr(), data.as_ptr()), "data");
     let alloc_id = alloc_id!(data.as_ptr());
     let x = &mut data[1];
     name!(x as *mut _, "data[1]");

tests/pass/tree_borrows/reserved.rs

@@ -33,7 +33,7 @@ unsafe fn print(msg: &str) {
 }
 
 unsafe fn read_second<T>(x: &mut T, y: *mut u8) {
-    name!(x as *mut T as *mut u8=>1, "caller:x");
+    name!(parent!(x as *mut T as *mut u8), "caller:x");
     name!(x as *mut T as *mut u8, "callee:x");
     name!(y, "caller:y");
     name!(y, "callee:y");

tests/utils/macros.rs

@@ -29,7 +29,30 @@ macro_rules! print_state {
     };
 }
 
-/// `name!(ptr => nth_parent, name)`: associate `name` to the `nth_parent` of `ptr`.
+/// `parent!(ptr)`: forges a pointer with the provenance of the direct parent of `ptr`.
+///
+/// Use in conjunction with `name!` to access an otherwise unreachable tag.
+macro_rules! parent {
+    ($ptr:expr) => {
+        crate::utils::miri_extern::miri_tree_nth_parent($ptr as *const u8 as *const (), 1)
+    };
+}
+
+/// `ancestor!(ptr1, ptr2)`: forges a pointer with the provenance of the
+/// nearest common ancestor of `ptr1` and `ptr2`.
+/// Caller should ensure that they are part of the same allocation.
+///
+/// Use in conjunction with `name!` to access an otherwise unreachable tag.
+macro_rules! ancestor {
+    ($ptr1:expr, $ptr2:expr) => {
+        crate::utils::miri_extern::miri_tree_common_ancestor(
+            $ptr1 as *const u8 as *const (),
+            $ptr2 as *const u8 as *const (),
+        )
+    };
+}
+
+/// `name!(tg, name)`: associate `name` to the `nth_parent` of `ptr`.
 ///
 /// `ptr` should be any pointer or reference that can be converted with `_ as *const u8`.
 ///
@@ -41,21 +64,17 @@ macro_rules! print_state {
 /// `name` is an optional string that will be used as the name. Defaults to
 /// `stringify!($ptr)` the name of `ptr` in the source code.
 macro_rules! name {
-    ($ptr:expr, $name:expr) => {
-        crate::utils::macros::name!($ptr => 0, $name);
-    };
     ($ptr:expr) => {
-        crate::utils::macros::name!($ptr => 0, stringify!($ptr));
-    };
-    ($ptr:expr => $nb:expr) => {
-        crate::utils::macros::name!($ptr => $nb, stringify!($ptr));
+        crate::utils::macros::name!($ptr, stringify!($ptr));
     };
-    ($ptr:expr => $nb:expr, $name:expr) => {
+    ($ptr:expr, $name:expr) => {
         let name = $name.as_bytes();
-        crate::utils::miri_extern::miri_pointer_name($ptr as *const u8 as *const (), $nb, name);
+        crate::utils::miri_extern::miri_pointer_name($ptr as *const u8 as *const (), name);
     };
 }
 
 pub(crate) use alloc_id;
+pub(crate) use ancestor;
 pub(crate) use name;
+pub(crate) use parent;
 pub(crate) use print_state;