From 0346e926a805a338754a192714cd7f11b29ba816 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Wed, 12 Mar 2025 18:39:06 +0100
Subject: [PATCH] alloc_addresses: use MemoryKind instead of tcx query to
 determine global allocations

---
 src/tools/miri/src/alloc_addresses/mod.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/tools/miri/src/alloc_addresses/mod.rs b/src/tools/miri/src/alloc_addresses/mod.rs
index ff3a25e94bd5..26d71dca3606 100644
--- a/src/tools/miri/src/alloc_addresses/mod.rs
+++ b/src/tools/miri/src/alloc_addresses/mod.rs
@@ -170,20 +170,22 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
             // This ensures the interpreted program and native code have the same view of memory.
             let base_ptr = match info.kind {
                 AllocKind::LiveData => {
-                    if this.tcx.try_get_global_alloc(alloc_id).is_some() {
+                    if memory_kind == MiriMemoryKind::Global.into() {
                         // For new global allocations, we always pre-allocate the memory to be able use the machine address directly.
                         let prepared_bytes = MiriAllocBytes::zeroed(info.size, info.align)
                             .unwrap_or_else(|| {
                                 panic!("Miri ran out of memory: cannot create allocation of {size:?} bytes", size = info.size)
                             });
                         let ptr = prepared_bytes.as_ptr();
-                        // Store prepared allocation space to be picked up for use later.
+                        // Store prepared allocation to be picked up for use later.
                         global_state
                             .prepared_alloc_bytes
                             .try_insert(alloc_id, prepared_bytes)
                             .unwrap();
                         ptr
                     } else {
+                        // Non-global allocations are already in memory at this point so
+                        // we can just get a pointer to where their data is stored.
                         this.get_alloc_bytes_unchecked_raw(alloc_id)?
                     }
                 }
@@ -382,6 +384,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         align: Align,
     ) -> InterpResult<'tcx, MiriAllocBytes> {
         let this = self.eval_context_ref();
+        assert!(this.tcx.try_get_global_alloc(id).is_some());
         if this.machine.native_lib.is_some() {
             // In native lib mode, MiriAllocBytes for global allocations are handled via `prepared_alloc_bytes`.
             // This additional call ensures that some `MiriAllocBytes` are always prepared, just in case