Move unwrap_fn and unwrap_memory to GlobalAlloc

Author: oli-obk

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

@@ -379,6 +379,28 @@ pub enum GlobalAlloc<'tcx> {
     Memory(&'tcx Allocation),
 }
 
+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),
+        }
+    }
+}
+
 pub struct AllocMap<'tcx> {
     /// Maps `AllocId`s to their corresponding allocations.
     alloc_map: FxHashMap<AllocId, GlobalAlloc<'tcx>>,
@@ -491,22 +513,6 @@ impl<'tcx> TyCtxt<'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_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_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(&self, id: AllocId, mem: &'tcx Allocation) {

src/librustc_middle/ty/print/pretty.rs

@@ -956,7 +956,9 @@ pub trait PrettyPrinter<'tcx>:
             ) => {
                 let byte_str = self
                     .tcx()
-                    .unwrap_memory(ptr.alloc_id)
+                    .get_global_alloc(ptr.alloc_id)
+                    .unwrap()
+                    .unwrap_memory()
                     .get_bytes(&self.tcx(), ptr, Size::from_bytes(*data))
                     .unwrap();
                 p!(pretty_print_byte_str(byte_str));
@@ -1019,7 +1021,7 @@ pub trait PrettyPrinter<'tcx>:
                 )?;
             }
             (Scalar::Ptr(ptr), ty::FnPtr(_)) => {
-                let instance = self.tcx().unwrap_fn(ptr.alloc_id);
+                let instance = self.tcx().get_global_alloc(ptr.alloc_id).unwrap().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,8 +549,10 @@ 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 a_instance = tcx.unwrap_fn(a_val.assert_ptr().alloc_id);
-                        let b_instance = tcx.unwrap_fn(b_val.assert_ptr().alloc_id);
+                        let a_instance =
+                            tcx.get_global_alloc(a_val.assert_ptr().alloc_id).unwrap().unwrap_fn();
+                        let b_instance =
+                            tcx.get_global_alloc(b_val.assert_ptr().alloc_id).unwrap().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.unwrap_memory(ptr.alloc_id);
+            let alloc = ecx.tcx.get_global_alloc(ptr.alloc_id).unwrap().unwrap_memory();
             ConstValue::ByRef { alloc, offset: ptr.offset }
         }
         Scalar::Raw { data, .. } => {
@@ -154,7 +154,10 @@ pub(super) fn op_to_const<'tcx>(
             },
             Immediate::ScalarPair(a, b) => {
                 let (data, start) = match a.not_undef().unwrap() {
-                    Scalar::Ptr(ptr) => (ecx.tcx.unwrap_memory(ptr.alloc_id), ptr.offset.bytes()),
+                    Scalar::Ptr(ptr) => (
+                        ecx.tcx.get_global_alloc(ptr.alloc_id).unwrap().unwrap_memory(),
+                        ptr.offset.bytes(),
+                    ),
                     Scalar::Raw { .. } => (
                         ecx.tcx
                             .intern_const_alloc(Allocation::from_byte_aligned_bytes(b"" as &[u8])),
@@ -200,7 +203,10 @@ 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.unwrap_memory(ptr.alloc_id), offset: ptr.offset })
+            Ok(ConstValue::ByRef {
+                alloc: ecx.tcx.get_global_alloc(ptr.alloc_id).unwrap().unwrap_memory(),
+                offset: ptr.offset,
+            })
         } else {
             Ok(op_to_const(&ecx, mplace.into()))
         }

src/librustc_mir/interpret/memory.rs

@@ -467,7 +467,7 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
                     })?;
                 // Make sure we use the ID of the resolved memory, not the lazy one!
                 let id = raw_const.alloc_id;
-                let allocation = tcx.unwrap_memory(id);
+                let allocation = tcx.get_global_alloc(id).unwrap().unwrap_memory();
 
                 (allocation, Some(def_id))
             }

src/librustc_mir_build/hair/pattern/_match.rs

@@ -286,7 +286,7 @@ impl<'tcx> LiteralExpander<'tcx> {
             (ConstValue::Scalar(p), x, y) if x == y => {
                 match p {
                     Scalar::Ptr(p) => {
-                        let alloc = self.tcx.unwrap_memory(p.alloc_id);
+                        let alloc = self.tcx.get_global_alloc(p.alloc_id).unwrap().unwrap_memory();
                         ConstValue::ByRef { alloc, offset: p.offset }
                     }
                     Scalar::Raw { .. } => {
@@ -305,7 +305,7 @@ impl<'tcx> LiteralExpander<'tcx> {
             (ConstValue::Scalar(Scalar::Ptr(p)), ty::Array(t, n), ty::Slice(u)) => {
                 assert_eq!(t, u);
                 ConstValue::Slice {
-                    data: self.tcx.unwrap_memory(p.alloc_id),
+                    data: self.tcx.get_global_alloc(p.alloc_id).unwrap().unwrap_memory(),
                     start: p.offset.bytes().try_into().unwrap(),
                     end: n.eval_usize(self.tcx, ty::ParamEnv::empty()).try_into().unwrap(),
                 }