Simplify the tcx.alloc_map API

Author: oli-obk

src/librustc_codegen_llvm/common.rs

@@ -244,8 +244,7 @@ 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 {
+                let base_addr = match self.tcx.get_global_alloc(ptr.alloc_id) {
                     Some(GlobalAlloc::Memory(alloc)) => {
                         let init = const_alloc_to_llvm(self, alloc);
                         let value = match alloc.mutability {

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/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.get_global_alloc(alloc_id).expect("no value for given 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)
                 }
             }
@@ -400,12 +398,6 @@ impl<'tcx> AllocMap<'tcx> {
     pub 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 {
         let next = self.next_id;
         self.next_id.0 = self.next_id.0.checked_add(1).expect(
@@ -415,34 +407,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 +460,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 +475,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 +487,38 @@ 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()
+    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 an `Allocation`
     pub fn unwrap_memory(&self, id: AllocId) -> &'tcx Allocation {
-        match self.get(id) {
+        match self.get_global_alloc(id) {
             Some(GlobalAlloc::Memory(mem)) => mem,
             _ => bug!("expected allocation ID {} to point to memory", id),
         }
     }
 
     /// Panics if the `AllocId` does not refer to a function
     pub fn unwrap_fn(&self, id: AllocId) -> Instance<'tcx> {
-        match self.get(id) {
+        match self.get_global_alloc(id) {
             Some(GlobalAlloc::Function(instance)) => instance,
             _ => bug!("expected allocation ID {} to point to a function", 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

@@ -2410,7 +2410,7 @@ 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(Scalar::Ptr(ptr)) => match tcx.get_global_alloc(ptr.alloc_id) {
                 Some(GlobalAlloc::Static(def_id)) => Some(def_id),
                 Some(_) => None,
                 None => {

src/librustc_middle/ty/context.rs

@@ -975,7 +975,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, ()>,
 
@@ -1013,7 +1013,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,8 +956,6 @@ pub trait PrettyPrinter<'tcx>:
             ) => {
                 let byte_str = self
                     .tcx()
-                    .alloc_map
-                    .lock()
                     .unwrap_memory(ptr.alloc_id)
                     .get_bytes(&self.tcx(), ptr, Size::from_bytes(*data))
                     .unwrap();
@@ -1021,10 +1019,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().unwrap_fn(ptr.alloc_id);
                 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.unwrap_fn(a_val.assert_ptr().alloc_id);
+                        let b_instance = tcx.unwrap_fn(b_val.assert_ptr().alloc_id);
                         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.unwrap_memory(ptr.alloc_id);
             ConstValue::ByRef { alloc, offset: ptr.offset }
         }
         Scalar::Raw { data, .. } => {
@@ -154,9 +154,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())
-                    }
+                    Scalar::Ptr(ptr) => (ecx.tcx.unwrap_memory(ptr.alloc_id), ptr.offset.bytes()),
                     Scalar::Raw { .. } => (
                         ecx.tcx
                             .intern_const_alloc(Allocation::from_byte_aligned_bytes(b"" as &[u8])),
@@ -202,10 +200,7 @@ fn validate_and_turn_into_const<'tcx>(
         // whether they become immediates.
         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),
-                offset: ptr.offset,
-            })
+            Ok(ConstValue::ByRef { alloc: ecx.tcx.unwrap_memory(ptr.alloc_id), offset: ptr.offset })
         } else {
             Ok(op_to_const(&ecx, mplace.into()))
         }

src/librustc_mir/interpret/intern.rs

@@ -91,7 +91,7 @@ fn intern_shallow<'rt, 'mir, 'tcx, M: CompileTimeMachine<'mir, 'tcx>>(
             // If the pointer is dangling (neither in local nor global memory), we leave it
             // to validation to error. The `delay_span_bug` ensures that we don't forget such
             // a check in validation.
-            if tcx.alloc_map.lock().get(alloc_id).is_none() {
+            if tcx.get_global_alloc(alloc_id).is_none() {
                 tcx.sess.delay_span_bug(ecx.tcx.span, "tried to intern dangling pointer");
             }
             // treat dangling pointers like other statics
@@ -134,7 +134,7 @@ fn intern_shallow<'rt, 'mir, 'tcx, M: CompileTimeMachine<'mir, 'tcx>>(
     // link the alloc id to the actual allocation
     let alloc = tcx.intern_const_alloc(alloc);
     leftover_allocations.extend(alloc.relocations().iter().map(|&(_, ((), reloc))| reloc));
-    tcx.alloc_map.lock().set_alloc_id_memory(alloc_id, alloc);
+    tcx.set_alloc_id_memory(alloc_id, alloc);
     Ok(None)
 }
 
@@ -389,7 +389,7 @@ where
                 }
             }
             let alloc = tcx.intern_const_alloc(alloc);
-            tcx.alloc_map.lock().set_alloc_id_memory(alloc_id, alloc);
+            tcx.set_alloc_id_memory(alloc_id, alloc);
             for &(_, ((), reloc)) in alloc.relocations().iter() {
                 if leftover_allocations.insert(reloc) {
                     todo.push(reloc);
@@ -398,7 +398,7 @@ where
         } else if ecx.memory.dead_alloc_map.contains_key(&alloc_id) {
             // dangling pointer
             throw_ub_format!("encountered dangling pointer in final constant")
-        } else if ecx.tcx.alloc_map.lock().get(alloc_id).is_none() {
+        } else if ecx.tcx.get_global_alloc(alloc_id).is_none() {
             // We have hit an `AllocId` that is neither in local or global memory and isn't marked
             // as dangling by local memory.
             span_bug!(ecx.tcx.span, "encountered unknown alloc id {:?}", alloc_id);