From 59f9a918eddcf9d131f55c2ef96d9e82a5706cf0 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Wed, 20 Jul 2022 17:40:49 -0400
Subject: [PATCH] handle get_alloc_extra the same throughout Stacked Borrows

---
 src/intptrcast.rs          |  9 +++++--
 src/machine.rs             |  7 +++---
 src/stacked_borrows/mod.rs | 48 ++++++++++++++++++++++++--------------
 3 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/src/intptrcast.rs b/src/intptrcast.rs
index 4272966ae6c..442c201e8ec 100644
--- a/src/intptrcast.rs
+++ b/src/intptrcast.rs
@@ -94,16 +94,21 @@ fn alloc_id_from_addr(ecx: &MiriEvalContext<'mir, 'tcx>, addr: u64) -> Option<Al
         None
     }
 
-    pub fn expose_ptr(ecx: &mut MiriEvalContext<'mir, 'tcx>, alloc_id: AllocId, sb: SbTag) {
+    pub fn expose_ptr(
+        ecx: &mut MiriEvalContext<'mir, 'tcx>,
+        alloc_id: AllocId,
+        sb: SbTag,
+    ) -> InterpResult<'tcx> {
         let global_state = ecx.machine.intptrcast.get_mut();
         // In strict mode, we don't need this, so we can save some cycles by not tracking it.
         if global_state.provenance_mode != ProvenanceMode::Strict {
             trace!("Exposing allocation id {alloc_id:?}");
             global_state.exposed.insert(alloc_id);
             if ecx.machine.stacked_borrows.is_some() {
-                ecx.expose_tag(alloc_id, sb);
+                ecx.expose_tag(alloc_id, sb)?;
             }
         }
+        Ok(())
     }
 
     pub fn ptr_from_addr_transmute(
diff --git a/src/machine.rs b/src/machine.rs
index 1adfb837781..1ac60e2ad84 100644
--- a/src/machine.rs
+++ b/src/machine.rs
@@ -754,15 +754,14 @@ fn expose_ptr(
         ptr: Pointer<Self::Provenance>,
     ) -> InterpResult<'tcx> {
         match ptr.provenance {
-            Provenance::Concrete { alloc_id, sb } => {
-                intptrcast::GlobalStateInner::expose_ptr(ecx, alloc_id, sb);
-            }
+            Provenance::Concrete { alloc_id, sb } =>
+                intptrcast::GlobalStateInner::expose_ptr(ecx, alloc_id, sb),
             Provenance::Wildcard => {
                 // No need to do anything for wildcard pointers as
                 // their provenances have already been previously exposed.
+                Ok(())
             }
         }
-        Ok(())
     }
 
     /// Convert a pointer with provenance into an allocation-offset pair,
diff --git a/src/stacked_borrows/mod.rs b/src/stacked_borrows/mod.rs
index c20927013b0..38a73929a03 100644
--- a/src/stacked_borrows/mod.rs
+++ b/src/stacked_borrows/mod.rs
@@ -777,20 +777,31 @@ fn reborrow(
                 return Ok(())
             };
 
-            let extra = this.get_alloc_extra(alloc_id)?;
-            let mut stacked_borrows = extra
-                .stacked_borrows
-                .as_ref()
-                .expect("we should have Stacked Borrows data")
-                .borrow_mut();
-            stacked_borrows.history.log_creation(
-                Some(orig_tag),
-                new_tag,
-                alloc_range(base_offset, size),
-                current_span,
-            );
-            if protect {
-                stacked_borrows.history.log_protector(orig_tag, new_tag, current_span);
+            let (_size, _align, kind) = this.get_alloc_info(alloc_id);
+            match kind {
+                AllocKind::LiveData => {
+                    // This should have alloc_extra data, but `get_alloc_extra` can still fail
+                    // if converting this alloc_id from a global to a local one
+                    // uncovers a non-supported `extern static`.
+                    let extra = this.get_alloc_extra(alloc_id)?;
+                    let mut stacked_borrows = extra
+                        .stacked_borrows
+                        .as_ref()
+                        .expect("we should have Stacked Borrows data")
+                        .borrow_mut();
+                    stacked_borrows.history.log_creation(
+                        Some(orig_tag),
+                        new_tag,
+                        alloc_range(base_offset, size),
+                        current_span,
+                    );
+                    if protect {
+                        stacked_borrows.history.log_protector(orig_tag, new_tag, current_span);
+                    }
+                }
+                AllocKind::Function | AllocKind::Dead => {
+                    // No stacked borrows on these allocations.
+                }
             }
             Ok(())
         };
@@ -1116,7 +1127,7 @@ fn retag_return_place(&mut self) -> InterpResult<'tcx> {
     }
 
     /// Mark the given tag as exposed. It was found on a pointer with the given AllocId.
-    fn expose_tag(&mut self, alloc_id: AllocId, tag: SbTag) {
+    fn expose_tag(&mut self, alloc_id: AllocId, tag: SbTag) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
 
         // Function pointers and dead objects don't have an alloc_extra so we ignore them.
@@ -1125,8 +1136,10 @@ fn expose_tag(&mut self, alloc_id: AllocId, tag: SbTag) {
         let (_size, _align, kind) = this.get_alloc_info(alloc_id);
         match kind {
             AllocKind::LiveData => {
-                // This should have alloc_extra data.
-                let alloc_extra = this.get_alloc_extra(alloc_id).unwrap();
+                // This should have alloc_extra data, but `get_alloc_extra` can still fail
+                // if converting this alloc_id from a global to a local one
+                // uncovers a non-supported `extern static`.
+                let alloc_extra = this.get_alloc_extra(alloc_id)?;
                 trace!("Stacked Borrows tag {tag:?} exposed in {alloc_id:?}");
                 alloc_extra.stacked_borrows.as_ref().unwrap().borrow_mut().exposed_tags.insert(tag);
             }
@@ -1134,5 +1147,6 @@ fn expose_tag(&mut self, alloc_id: AllocId, tag: SbTag) {
                 // No stacked borrows on these allocations.
             }
         }
+        Ok(())
     }
 }