Rollup merge of #71508 - oli-obk:alloc_map_unlock, r=RalfJung

Simplify the `tcx.alloc_map` API

This PR changes all functions that require manually locking the `alloc_map` to functions on `TyCtxt` that lock the map internally. In the same step we make the `TyCtxt::alloc_map` field private.

r? @RalfJung

src/librustc_codegen_llvm/common.rs

@@ -244,9 +244,8 @@ impl ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
                 }
             }
             Scalar::Ptr(ptr) => {
-                let alloc_kind = self.tcx.alloc_map.lock().get(ptr.alloc_id);
-                let base_addr = match alloc_kind {
-                    Some(GlobalAlloc::Memory(alloc)) => {
+                let base_addr = match self.tcx.global_alloc(ptr.alloc_id) {
+                    GlobalAlloc::Memory(alloc) => {
                         let init = const_alloc_to_llvm(self, alloc);
                         let value = match alloc.mutability {
                             Mutability::Mut => self.static_addr_of_mut(init, alloc.align, None),
@@ -257,12 +256,11 @@ impl ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
                         }
                         value
                     }
-                    Some(GlobalAlloc::Function(fn_instance)) => self.get_fn_addr(fn_instance),
-                    Some(GlobalAlloc::Static(def_id)) => {
+                    GlobalAlloc::Function(fn_instance) => self.get_fn_addr(fn_instance),
+                    GlobalAlloc::Static(def_id) => {
                         assert!(self.tcx.is_static(def_id));
                         self.get_static(def_id)
                     }
-                    None => bug!("missing allocation {:?}", ptr.alloc_id),
                 };
                 let llval = unsafe {
                     llvm::LLVMConstInBoundsGEP(

src/librustc_codegen_ssa/mir/operand.rs

@@ -92,7 +92,7 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                     _ => bug!("from_const: invalid ScalarPair layout: {:#?}", layout),
                 };
                 let a = Scalar::from(Pointer::new(
-                    bx.tcx().alloc_map.lock().create_memory_alloc(data),
+                    bx.tcx().create_memory_alloc(data),
                     Size::from_bytes(start),
                 ));
                 let a_llval = bx.scalar_to_backend(

src/librustc_middle/ich/impls_ty.rs

@@ -136,8 +136,7 @@ impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::AllocId {
         ty::tls::with_opt(|tcx| {
             trace!("hashing {:?}", *self);
             let tcx = tcx.expect("can't hash AllocIds during hir lowering");
-            let alloc_kind = tcx.alloc_map.lock().get(*self);
-            alloc_kind.hash_stable(hcx, hasher);
+            tcx.get_global_alloc(*self).hash_stable(hcx, hasher);
         });
     }
 }

src/librustc_middle/lib.rs

@@ -42,6 +42,7 @@
 #![feature(or_patterns)]
 #![feature(range_is_empty)]
 #![feature(specialization)]
+#![feature(track_caller)]
 #![feature(trusted_len)]
 #![feature(vec_remove_item)]
 #![feature(stmt_expr_attributes)]

src/librustc_middle/mir/interpret/mod.rs

@@ -197,9 +197,7 @@ pub fn specialized_encode_alloc_id<'tcx, E: Encoder>(
     tcx: TyCtxt<'tcx>,
     alloc_id: AllocId,
 ) -> Result<(), E::Error> {
-    let alloc: GlobalAlloc<'tcx> =
-        tcx.alloc_map.lock().get(alloc_id).expect("no value for given alloc ID");
-    match alloc {
+    match tcx.global_alloc(alloc_id) {
         GlobalAlloc::Memory(alloc) => {
             trace!("encoding {:?} with {:#?}", alloc_id, alloc);
             AllocDiscriminant::Alloc.encode(encoder)?;
@@ -294,7 +292,7 @@ impl<'s> AllocDecodingSession<'s> {
                         AllocDiscriminant::Alloc => {
                             // If this is an allocation, we need to reserve an
                             // `AllocId` so we can decode cyclic graphs.
-                            let alloc_id = decoder.tcx().alloc_map.lock().reserve();
+                            let alloc_id = decoder.tcx().reserve_alloc_id();
                             *entry =
                                 State::InProgress(TinyList::new_single(self.session_id), alloc_id);
                             Some(alloc_id)
@@ -338,23 +336,23 @@ impl<'s> AllocDecodingSession<'s> {
                     // We already have a reserved `AllocId`.
                     let alloc_id = alloc_id.unwrap();
                     trace!("decoded alloc {:?}: {:#?}", alloc_id, alloc);
-                    decoder.tcx().alloc_map.lock().set_alloc_id_same_memory(alloc_id, alloc);
+                    decoder.tcx().set_alloc_id_same_memory(alloc_id, alloc);
                     Ok(alloc_id)
                 }
                 AllocDiscriminant::Fn => {
                     assert!(alloc_id.is_none());
                     trace!("creating fn alloc ID");
                     let instance = ty::Instance::decode(decoder)?;
                     trace!("decoded fn alloc instance: {:?}", instance);
-                    let alloc_id = decoder.tcx().alloc_map.lock().create_fn_alloc(instance);
+                    let alloc_id = decoder.tcx().create_fn_alloc(instance);
                     Ok(alloc_id)
                 }
                 AllocDiscriminant::Static => {
                     assert!(alloc_id.is_none());
                     trace!("creating extern static alloc ID");
                     let did = DefId::decode(decoder)?;
                     trace!("decoded static def-ID: {:?}", did);
-                    let alloc_id = decoder.tcx().alloc_map.lock().create_static_alloc(did);
+                    let alloc_id = decoder.tcx().create_static_alloc(did);
                     Ok(alloc_id)
                 }
             }
@@ -381,7 +379,29 @@ pub enum GlobalAlloc<'tcx> {
     Memory(&'tcx Allocation),
 }
 
-pub struct AllocMap<'tcx> {
+impl GlobalAlloc<'tcx> {
+    /// Panics if the `GlobalAlloc` does not refer to an `GlobalAlloc::Memory`
+    #[track_caller]
+    #[inline]
+    pub fn unwrap_memory(&self) -> &'tcx Allocation {
+        match *self {
+            GlobalAlloc::Memory(mem) => mem,
+            _ => bug!("expected memory, got {:?}", self),
+        }
+    }
+
+    /// Panics if the `GlobalAlloc` is not `GlobalAlloc::Function`
+    #[track_caller]
+    #[inline]
+    pub fn unwrap_fn(&self) -> Instance<'tcx> {
+        match *self {
+            GlobalAlloc::Function(instance) => instance,
+            _ => bug!("expected function, got {:?}", self),
+        }
+    }
+}
+
+crate struct AllocMap<'tcx> {
     /// Maps `AllocId`s to their corresponding allocations.
     alloc_map: FxHashMap<AllocId, GlobalAlloc<'tcx>>,
 
@@ -397,16 +417,10 @@ pub struct AllocMap<'tcx> {
 }
 
 impl<'tcx> AllocMap<'tcx> {
-    pub fn new() -> Self {
+    crate fn new() -> Self {
         AllocMap { alloc_map: Default::default(), dedup: Default::default(), next_id: AllocId(0) }
     }
-
-    /// Obtains a new allocation ID that can be referenced but does not
-    /// yet have an allocation backing it.
-    ///
-    /// Make sure to call `set_alloc_id_memory` or `set_alloc_id_same_memory` before returning such
-    /// an `AllocId` from a query.
-    pub fn reserve(&mut self) -> AllocId {
+    fn reserve(&mut self) -> AllocId {
         let next = self.next_id;
         self.next_id.0 = self.next_id.0.checked_add(1).expect(
             "You overflowed a u64 by incrementing by 1... \
@@ -415,34 +429,46 @@ impl<'tcx> AllocMap<'tcx> {
         );
         next
     }
+}
+
+impl<'tcx> TyCtxt<'tcx> {
+    /// Obtains a new allocation ID that can be referenced but does not
+    /// yet have an allocation backing it.
+    ///
+    /// Make sure to call `set_alloc_id_memory` or `set_alloc_id_same_memory` before returning such
+    /// an `AllocId` from a query.
+    pub fn reserve_alloc_id(&self) -> AllocId {
+        self.alloc_map.lock().reserve()
+    }
 
     /// Reserves a new ID *if* this allocation has not been dedup-reserved before.
     /// Should only be used for function pointers and statics, we don't want
     /// to dedup IDs for "real" memory!
-    fn reserve_and_set_dedup(&mut self, alloc: GlobalAlloc<'tcx>) -> AllocId {
+    fn reserve_and_set_dedup(&self, alloc: GlobalAlloc<'tcx>) -> AllocId {
+        let mut alloc_map = self.alloc_map.lock();
         match alloc {
             GlobalAlloc::Function(..) | GlobalAlloc::Static(..) => {}
             GlobalAlloc::Memory(..) => bug!("Trying to dedup-reserve memory with real data!"),
         }
-        if let Some(&alloc_id) = self.dedup.get(&alloc) {
+        if let Some(&alloc_id) = alloc_map.dedup.get(&alloc) {
             return alloc_id;
         }
-        let id = self.reserve();
+        let id = alloc_map.reserve();
         debug!("creating alloc {:?} with id {}", alloc, id);
-        self.alloc_map.insert(id, alloc.clone());
-        self.dedup.insert(alloc, id);
+        alloc_map.alloc_map.insert(id, alloc.clone());
+        alloc_map.dedup.insert(alloc, id);
         id
     }
 
     /// Generates an `AllocId` for a static or return a cached one in case this function has been
     /// called on the same static before.
-    pub fn create_static_alloc(&mut self, static_id: DefId) -> AllocId {
+    pub fn create_static_alloc(&self, static_id: DefId) -> AllocId {
         self.reserve_and_set_dedup(GlobalAlloc::Static(static_id))
     }
 
     /// Generates an `AllocId` for a function.  Depending on the function type,
     /// this might get deduplicated or assigned a new ID each time.
-    pub fn create_fn_alloc(&mut self, instance: Instance<'tcx>) -> AllocId {
+    pub fn create_fn_alloc(&self, instance: Instance<'tcx>) -> AllocId {
         // Functions cannot be identified by pointers, as asm-equal functions can get deduplicated
         // by the linker (we set the "unnamed_addr" attribute for LLVM) and functions can be
         // duplicated across crates.
@@ -456,8 +482,9 @@ impl<'tcx> AllocMap<'tcx> {
         });
         if is_generic {
             // Get a fresh ID.
-            let id = self.reserve();
-            self.alloc_map.insert(id, GlobalAlloc::Function(instance));
+            let mut alloc_map = self.alloc_map.lock();
+            let id = alloc_map.reserve();
+            alloc_map.alloc_map.insert(id, GlobalAlloc::Function(instance));
             id
         } else {
             // Deduplicate.
@@ -470,8 +497,8 @@ impl<'tcx> AllocMap<'tcx> {
     /// Statics with identical content will still point to the same `Allocation`, i.e.,
     /// their data will be deduplicated through `Allocation` interning -- but they
     /// are different places in memory and as such need different IDs.
-    pub fn create_memory_alloc(&mut self, mem: &'tcx Allocation) -> AllocId {
-        let id = self.reserve();
+    pub fn create_memory_alloc(&self, mem: &'tcx Allocation) -> AllocId {
+        let id = self.reserve_alloc_id();
         self.set_alloc_id_memory(id, mem);
         id
     }
@@ -482,38 +509,35 @@ impl<'tcx> AllocMap<'tcx> {
     /// This function exists to allow const eval to detect the difference between evaluation-
     /// local dangling pointers and allocations in constants/statics.
     #[inline]
-    pub fn get(&self, id: AllocId) -> Option<GlobalAlloc<'tcx>> {
-        self.alloc_map.get(&id).cloned()
-    }
-
-    /// Panics if the `AllocId` does not refer to an `Allocation`
-    pub fn unwrap_memory(&self, id: AllocId) -> &'tcx Allocation {
-        match self.get(id) {
-            Some(GlobalAlloc::Memory(mem)) => mem,
-            _ => bug!("expected allocation ID {} to point to memory", id),
-        }
+    pub fn get_global_alloc(&self, id: AllocId) -> Option<GlobalAlloc<'tcx>> {
+        self.alloc_map.lock().alloc_map.get(&id).cloned()
     }
 
-    /// Panics if the `AllocId` does not refer to a function
-    pub fn unwrap_fn(&self, id: AllocId) -> Instance<'tcx> {
-        match self.get(id) {
-            Some(GlobalAlloc::Function(instance)) => instance,
-            _ => bug!("expected allocation ID {} to point to a function", id),
+    #[inline]
+    #[track_caller]
+    /// Panics in case the `AllocId` is dangling. Since that is impossible for `AllocId`s in
+    /// constants (as all constants must pass interning and validation that check for dangling
+    /// ids), this function is frequently used throughout rustc, but should not be used within
+    /// the miri engine.
+    pub fn global_alloc(&self, id: AllocId) -> GlobalAlloc<'tcx> {
+        match self.get_global_alloc(id) {
+            Some(alloc) => alloc,
+            None => bug!("could not find allocation for {}", id),
         }
     }
 
     /// Freezes an `AllocId` created with `reserve` by pointing it at an `Allocation`. Trying to
     /// call this function twice, even with the same `Allocation` will ICE the compiler.
-    pub fn set_alloc_id_memory(&mut self, id: AllocId, mem: &'tcx Allocation) {
-        if let Some(old) = self.alloc_map.insert(id, GlobalAlloc::Memory(mem)) {
+    pub fn set_alloc_id_memory(&self, id: AllocId, mem: &'tcx Allocation) {
+        if let Some(old) = self.alloc_map.lock().alloc_map.insert(id, GlobalAlloc::Memory(mem)) {
             bug!("tried to set allocation ID {}, but it was already existing as {:#?}", id, old);
         }
     }
 
     /// Freezes an `AllocId` created with `reserve` by pointing it at an `Allocation`. May be called
     /// twice for the same `(AllocId, Allocation)` pair.
-    fn set_alloc_id_same_memory(&mut self, id: AllocId, mem: &'tcx Allocation) {
-        self.alloc_map.insert_same(id, GlobalAlloc::Memory(mem));
+    fn set_alloc_id_same_memory(&self, id: AllocId, mem: &'tcx Allocation) {
+        self.alloc_map.lock().alloc_map.insert_same(id, GlobalAlloc::Memory(mem));
     }
 }
 

src/librustc_middle/mir/mod.rs

@@ -2404,13 +2404,9 @@ pub struct Constant<'tcx> {
 impl Constant<'tcx> {
     pub fn check_static_ptr(&self, tcx: TyCtxt<'_>) -> Option<DefId> {
         match self.literal.val.try_to_scalar() {
-            Some(Scalar::Ptr(ptr)) => match tcx.alloc_map.lock().get(ptr.alloc_id) {
-                Some(GlobalAlloc::Static(def_id)) => Some(def_id),
-                Some(_) => None,
-                None => {
-                    tcx.sess.delay_span_bug(DUMMY_SP, "MIR cannot contain dangling const pointers");
-                    None
-                }
+            Some(Scalar::Ptr(ptr)) => match tcx.global_alloc(ptr.alloc_id) {
+                GlobalAlloc::Static(def_id) => Some(def_id),
+                _ => None,
             },
             _ => None,
         }

src/librustc_middle/ty/context.rs

@@ -979,7 +979,7 @@ pub struct GlobalCtxt<'tcx> {
     allocation_interner: ShardedHashMap<&'tcx Allocation, ()>,
 
     /// Stores memory for globals (statics/consts).
-    pub alloc_map: Lock<interpret::AllocMap<'tcx>>,
+    pub(crate) alloc_map: Lock<interpret::AllocMap<'tcx>>,
 
     layout_interner: ShardedHashMap<&'tcx Layout, ()>,
 
@@ -1017,7 +1017,7 @@ impl<'tcx> TyCtxt<'tcx> {
         // Create an allocation that just contains these bytes.
         let alloc = interpret::Allocation::from_byte_aligned_bytes(bytes);
         let alloc = self.intern_const_alloc(alloc);
-        self.alloc_map.lock().create_memory_alloc(alloc)
+        self.create_memory_alloc(alloc)
     }
 
     pub fn intern_stability(self, stab: attr::Stability) -> &'tcx attr::Stability {

src/librustc_middle/ty/print/pretty.rs

@@ -956,9 +956,8 @@ pub trait PrettyPrinter<'tcx>:
             ) => {
                 let byte_str = self
                     .tcx()
-                    .alloc_map
-                    .lock()
-                    .unwrap_memory(ptr.alloc_id)
+                    .global_alloc(ptr.alloc_id)
+                    .unwrap_memory()
                     .get_bytes(&self.tcx(), ptr, Size::from_bytes(*data))
                     .unwrap();
                 p!(pretty_print_byte_str(byte_str));
@@ -1021,10 +1020,7 @@ pub trait PrettyPrinter<'tcx>:
                 )?;
             }
             (Scalar::Ptr(ptr), ty::FnPtr(_)) => {
-                let instance = {
-                    let alloc_map = self.tcx().alloc_map.lock();
-                    alloc_map.unwrap_fn(ptr.alloc_id)
-                };
+                let instance = self.tcx().global_alloc(ptr.alloc_id).unwrap_fn();
                 self = self.typed_value(
                     |this| this.print_value_path(instance.def_id(), instance.substs),
                     |this| this.print_type(ty),

src/librustc_middle/ty/relate.rs

@@ -549,9 +549,8 @@ pub fn super_relate_consts<R: TypeRelation<'tcx>>(
                     if a_val == b_val {
                         Ok(ConstValue::Scalar(a_val))
                     } else if let ty::FnPtr(_) = a.ty.kind {
-                        let alloc_map = tcx.alloc_map.lock();
-                        let a_instance = alloc_map.unwrap_fn(a_val.assert_ptr().alloc_id);
-                        let b_instance = alloc_map.unwrap_fn(b_val.assert_ptr().alloc_id);
+                        let a_instance = tcx.global_alloc(a_val.assert_ptr().alloc_id).unwrap_fn();
+                        let b_instance = tcx.global_alloc(b_val.assert_ptr().alloc_id).unwrap_fn();
                         if a_instance == b_instance {
                             Ok(ConstValue::Scalar(a_val))
                         } else {

src/librustc_mir/const_eval/eval_queries.rs

@@ -130,7 +130,7 @@ pub(super) fn op_to_const<'tcx>(
 
     let to_const_value = |mplace: MPlaceTy<'_>| match mplace.ptr {
         Scalar::Ptr(ptr) => {
-            let alloc = ecx.tcx.alloc_map.lock().unwrap_memory(ptr.alloc_id);
+            let alloc = ecx.tcx.global_alloc(ptr.alloc_id).unwrap_memory();
             ConstValue::ByRef { alloc, offset: ptr.offset }
         }
         Scalar::Raw { data, .. } => {
@@ -155,7 +155,7 @@ pub(super) fn op_to_const<'tcx>(
             Immediate::ScalarPair(a, b) => {
                 let (data, start) = match a.not_undef().unwrap() {
                     Scalar::Ptr(ptr) => {
-                        (ecx.tcx.alloc_map.lock().unwrap_memory(ptr.alloc_id), ptr.offset.bytes())
+                        (ecx.tcx.global_alloc(ptr.alloc_id).unwrap_memory(), ptr.offset.bytes())
                     }
                     Scalar::Raw { .. } => (
                         ecx.tcx
@@ -203,7 +203,7 @@ fn validate_and_turn_into_const<'tcx>(
         if is_static || cid.promoted.is_some() {
             let ptr = mplace.ptr.assert_ptr();
             Ok(ConstValue::ByRef {
-                alloc: ecx.tcx.alloc_map.lock().unwrap_memory(ptr.alloc_id),
+                alloc: ecx.tcx.global_alloc(ptr.alloc_id).unwrap_memory(),
                 offset: ptr.offset,
             })
         } else {