shims/foreign_items: Add hook for updating the canonical address

Where a method modifies the address that should be considered canonical - such as via TBI or when an MTE tag has been set - Miri will need to be notified.

This adds a hook to inform Miri of the new address.

src/alloc_addresses/mod.rs

@@ -460,6 +460,21 @@ impl<'tcx> MiriMachine<'tcx> {
             }
         })
     }
+
+    /// Updates the current canonical address for the allocation. Note that any access to this allocation *must* use this address from this point on.
+    pub fn set_alloc_address(&mut self, id: AllocId, new: u64) {
+        let global_state = self.alloc_addresses.get_mut();
+        if let Some(addr) = global_state.base_addr.insert(id, new) {
+            // Remove the old address' int->ptr mapping.
+            let pos =
+                global_state.int_to_ptr_map.binary_search_by_key(&addr, |(addr, _)| *addr).unwrap();
+            let removed = global_state.int_to_ptr_map.remove(pos);
+            assert_eq!(removed, (addr, id));
+        }
+        let new_pos =
+            global_state.int_to_ptr_map.binary_search_by_key(&new, |(addr, _)| *addr).unwrap_err();
+        global_state.int_to_ptr_map.insert(new_pos, (new, id));
+    }
 }
 
 #[cfg(test)]

src/shims/foreign_items.rs

@@ -423,6 +423,13 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     }
                 }
             }
+            "miri_set_canonical_address" => {
+                let [old_ptr, new_ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let old_ptr = this.read_pointer(old_ptr)?;
+                let new_ptr = this.read_pointer(new_ptr)?;
+                let (alloc_id, _, _) = this.ptr_get_alloc_id(old_ptr, 0)?;
+                this.machine.set_alloc_address(alloc_id, new_ptr.addr().bytes());
+            }
 
             // Aborting the process.
             "exit" => {

tests/utils/miri_extern.rs

@@ -147,4 +147,9 @@ extern "Rust" {
     /// "symbolic" alignment checks. Will fail if the pointer is not actually aligned or `align` is
     /// not a power of two. Has no effect when alignment checks are concrete (which is the default).
     pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+
+    /// Miri-provided extern function to specify that a new address is to be considered the
+    /// canonical address, where `new` is a valid alias to the `old` allocation,
+    /// usually due to them having different values in bits that are ignored by hardware.
+    pub fn miri_set_canonical_address(old: *const (), new: *const ());
 }