add miri magic function to configure allocation tracking at runtime

src/tools/miri/README.md

@@ -474,7 +474,8 @@ to Miri failing to detect cases of undefined behavior in a program.
 * `-Zmiri-track-alloc-id=<id1>,<id2>,...` shows a backtrace when the given allocations are
   being allocated or freed.  This helps in debugging memory leaks and
   use after free bugs. Specifying this argument multiple times does not overwrite the previous
-  values, instead it appends its values to the list. Listing an id multiple times has no effect.
+  values, instead it appends its values to the list. Listing an ID multiple times has no effect.
+  You can also add IDs at runtime using `miri_track_alloc`.
 * `-Zmiri-track-pointer-tag=<tag1>,<tag2>,...` shows a backtrace when a given pointer tag
   is created and when (if ever) it is popped from a borrow stack (which is where the tag becomes invalid
   and any future use of it will error).  This helps you in finding out why UB is

src/tools/miri/src/diagnostics.rs

@@ -126,7 +126,7 @@ pub enum NonHaltingDiagnostic {
     CreatedPointerTag(NonZero<u64>, Option<String>, Option<(AllocId, AllocRange, ProvenanceExtra)>),
     /// This `Item` was popped from the borrow stack. The string explains the reason.
     PoppedPointerTag(Item, String),
-    CreatedAlloc(AllocId, Size, Align, MemoryKind),
+    TrackingAlloc(AllocId, Size, Align),
     FreedAlloc(AllocId),
     AccessedAlloc(AllocId, AccessKind),
     RejectedIsolatedOp(String),
@@ -656,7 +656,7 @@ impl<'tcx> MiriMachine<'tcx> {
                 ("GenMC might miss possible behaviors of this code".to_string(), DiagLevel::Warning),
             CreatedPointerTag(..)
             | PoppedPointerTag(..)
-            | CreatedAlloc(..)
+            | TrackingAlloc(..)
             | AccessedAlloc(..)
             | FreedAlloc(..)
             | ProgressReport { .. }
@@ -673,9 +673,9 @@ impl<'tcx> MiriMachine<'tcx> {
                     "created tag {tag:?} with {perm} at {alloc_id:?}{range:?} derived from {orig_tag:?}"
                 ),
             PoppedPointerTag(item, cause) => format!("popped tracked tag for item {item:?}{cause}"),
-            CreatedAlloc(AllocId(id), size, align, kind) =>
+            TrackingAlloc(AllocId(id), size, align) =>
                 format!(
-                    "created {kind} allocation of {size} bytes (alignment {align} bytes) with id {id}",
+                    "now tracking allocation of {size} bytes (alignment {align} bytes) with id {id}",
                     size = size.bytes(),
                     align = align.bytes(),
                 ),

src/tools/miri/src/machine.rs

@@ -595,7 +595,7 @@ pub struct MiriMachine<'tcx> {
 
     /// The allocation IDs to report when they are being allocated
     /// (helps for debugging memory leaks and use after free bugs).
-    tracked_alloc_ids: FxHashSet<AllocId>,
+    pub(crate) tracked_alloc_ids: FxHashSet<AllocId>,
     /// For the tracked alloc ids, also report read/write accesses.
     track_alloc_accesses: bool,
 
@@ -928,7 +928,7 @@ impl<'tcx> MiriMachine<'tcx> {
         align: Align,
     ) -> InterpResult<'tcx, AllocExtra<'tcx>> {
         if ecx.machine.tracked_alloc_ids.contains(&id) {
-            ecx.emit_diagnostic(NonHaltingDiagnostic::CreatedAlloc(id, size, align, kind));
+            ecx.emit_diagnostic(NonHaltingDiagnostic::TrackingAlloc(id, size, align));
         }
 
         let borrow_tracker = ecx

src/tools/miri/src/shims/foreign_items.rs

@@ -350,6 +350,21 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     MiriMemoryKind::Miri.into(),
                 )?;
             }
+            "miri_track_alloc" => {
+                let [ptr] = this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
+                let ptr = this.read_pointer(ptr)?;
+                let (alloc_id, _, _) = this.ptr_get_alloc_id(ptr, 0).map_err_kind(|_e| {
+                    err_machine_stop!(TerminationInfo::Abort(format!(
+                        "pointer passed to `miri_get_alloc_id` must not be dangling, got {ptr:?}"
+                    )))
+                })?;
+                if this.machine.tracked_alloc_ids.insert(alloc_id) {
+                    let info = this.get_alloc_info(alloc_id);
+                    this.emit_diagnostic(NonHaltingDiagnostic::TrackingAlloc(
+                        alloc_id, info.size, info.align,
+                    ));
+                }
+            }
             "miri_start_unwind" => {
                 let [payload] =
                     this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;

src/tools/miri/tests/pass/alloc-access-tracking.rs

@@ -1,26 +1,15 @@
-#![no_std]
-#![no_main]
-//@compile-flags: -Zmiri-track-alloc-id=19 -Zmiri-track-alloc-accesses -Cpanic=abort
-//@normalize-stderr-test: "id 19" -> "id $$ALLOC"
-//@only-target: linux # alloc IDs differ between OSes (due to extern static allocations)
+//@compile-flags: -Zmiri-track-alloc-accesses
+//@normalize-stderr-test: "id \d+" -> "id $$ALLOC"
 
-extern "Rust" {
-    fn miri_alloc(size: usize, align: usize) -> *mut u8;
-    fn miri_dealloc(ptr: *mut u8, size: usize, align: usize);
-}
+#[path = "../utils/mod.rs"]
+mod utils;
 
-#[no_mangle]
-fn miri_start(_argc: isize, _argv: *const *const u8) -> isize {
+fn main() {
     unsafe {
-        let ptr = miri_alloc(123, 1);
+        let mut b = Box::<[u8; 123]>::new_uninit();
+        let ptr = b.as_mut_ptr() as *mut u8;
+        utils::miri_track_alloc(ptr);
         *ptr = 42; // Crucially, only a write is printed here, no read!
         assert_eq!(*ptr, 42);
-        miri_dealloc(ptr, 123, 1);
     }
-    0
-}
-
-#[panic_handler]
-fn panic_handler(_: &core::panic::PanicInfo) -> ! {
-    loop {}
 }

src/tools/miri/tests/pass/alloc-access-tracking.stderr

@@ -1,11 +1,11 @@
-note: created Miri bare-metal heap allocation of 123 bytes (alignment ALIGN bytes) with id $ALLOC
+note: now tracking allocation of 123 bytes (alignment ALIGN bytes) with id $ALLOC
   --> tests/pass/alloc-access-tracking.rs:LL:CC
    |
-LL |         let ptr = miri_alloc(123, 1);
-   |                   ^^^^^^^^^^^^^^^^^^ tracking was triggered here
+LL |         utils::miri_track_alloc(ptr);
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ tracking was triggered here
    |
    = note: BACKTRACE:
-   = note: inside `miri_start` at tests/pass/alloc-access-tracking.rs:LL:CC
+   = note: inside `main` at tests/pass/alloc-access-tracking.rs:LL:CC
 
 note: write access to allocation with id $ALLOC
   --> tests/pass/alloc-access-tracking.rs:LL:CC
@@ -14,7 +14,7 @@ LL |         *ptr = 42; // Crucially, only a write is printed here, no read!
    |         ^^^^^^^^^ tracking was triggered here
    |
    = note: BACKTRACE:
-   = note: inside `miri_start` at tests/pass/alloc-access-tracking.rs:LL:CC
+   = note: inside `main` at tests/pass/alloc-access-tracking.rs:LL:CC
 
 note: read access to allocation with id $ALLOC
   --> tests/pass/alloc-access-tracking.rs:LL:CC
@@ -23,15 +23,21 @@ LL |         assert_eq!(*ptr, 42);
    |         ^^^^^^^^^^^^^^^^^^^^ tracking was triggered here
    |
    = note: BACKTRACE:
-   = note: inside `miri_start` at RUSTLIB/core/src/macros/mod.rs:LL:CC
+   = note: inside `main` at RUSTLIB/core/src/macros/mod.rs:LL:CC
    = note: this note originates in the macro `assert_eq` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 note: freed allocation with id $ALLOC
-  --> tests/pass/alloc-access-tracking.rs:LL:CC
+  --> RUSTLIB/alloc/src/boxed.rs:LL:CC
    |
-LL |         miri_dealloc(ptr, 123, 1);
-   |         ^^^^^^^^^^^^^^^^^^^^^^^^^ tracking was triggered here
+LL |                 self.1.deallocate(From::from(ptr.cast()), layout);
+   |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ tracking was triggered here
    |
    = note: BACKTRACE:
-   = note: inside `miri_start` at tests/pass/alloc-access-tracking.rs:LL:CC
+   = note: inside `<std::boxed::Box<std::mem::MaybeUninit<[u8; 123]>> as std::ops::Drop>::drop` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+   = note: inside `std::ptr::drop_in_place::<std::boxed::Box<std::mem::MaybeUninit<[u8; 123]>>> - shim(Some(std::boxed::Box<std::mem::MaybeUninit<[u8; 123]>>))` at RUSTLIB/core/src/ptr/mod.rs:LL:CC
+note: inside `main`
+  --> tests/pass/alloc-access-tracking.rs:LL:CC
+   |
+LL |     }
+   |     ^
 

src/tools/miri/tests/utils/miri_extern.rs

@@ -119,6 +119,11 @@ extern "Rust" {
     /// Miri-provided extern function to deallocate memory.
     pub fn miri_dealloc(ptr: *mut u8, size: usize, align: usize);
 
+    /// Add the allocation that this pointer points to to the "tracked" allocations.
+    /// This is equivalent to `-Zmiri-track-allic-id=<id>`, but also works if the ID is
+    /// only known at runtime.
+    pub fn miri_track_alloc(ptr: *const u8);
+
     /// Convert a path from the host Miri runs on to the target Miri interprets.
     /// Performs conversion of path separators as needed.
     ///