Rollup merge of #98099 - RalfJung:convert_tag_add_extra, r=oli-obk

interpret: convert_tag_add_extra: allow tagger to raise errors

Needed for rust-lang/miri#2234

r? `@oli-obk`

Author: JohnTitor

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -334,12 +334,14 @@ pub trait Machine<'mir, 'tcx>: Sized {
     /// allocation (because a copy had to be done to add tags or metadata), machine memory will
     /// cache the result. (This relies on `AllocMap::get_or` being able to add the
     /// owned allocation to the map even when the map is shared.)
+    ///
+    /// This must only fail if `alloc` contains relocations.
     fn init_allocation_extra<'b>(
         ecx: &InterpCx<'mir, 'tcx, Self>,
         id: AllocId,
         alloc: Cow<'b, Allocation>,
         kind: Option<MemoryKind<Self::MemoryKind>>,
-    ) -> Cow<'b, Allocation<Self::PointerTag, Self::AllocExtra>>;
+    ) -> InterpResult<'tcx, Cow<'b, Allocation<Self::PointerTag, Self::AllocExtra>>>;
 
     /// Hook for performing extra checks on a memory read access.
     ///
@@ -485,9 +487,9 @@ pub macro compile_time_machine(<$mir: lifetime, $tcx: lifetime>) {
         _id: AllocId,
         alloc: Cow<'b, Allocation>,
         _kind: Option<MemoryKind<Self::MemoryKind>>,
-    ) -> Cow<'b, Allocation<Self::PointerTag>> {
+    ) -> InterpResult<$tcx, Cow<'b, Allocation<Self::PointerTag>>> {
         // We do not use a tag so we can just cheaply forward the allocation
-        alloc
+        Ok(alloc)
     }
 
     fn extern_static_base_pointer(

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -199,7 +199,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         kind: MemoryKind<M::MemoryKind>,
     ) -> InterpResult<'tcx, Pointer<M::PointerTag>> {
         let alloc = Allocation::uninit(size, align, M::PANIC_ON_ALLOC_FAIL)?;
-        Ok(self.allocate_raw_ptr(alloc, kind))
+        // We can `unwrap` since `alloc` contains no pointers.
+        Ok(self.allocate_raw_ptr(alloc, kind).unwrap())
     }
 
     pub fn allocate_bytes_ptr(
@@ -210,23 +211,25 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         mutability: Mutability,
     ) -> Pointer<M::PointerTag> {
         let alloc = Allocation::from_bytes(bytes, align, mutability);
-        self.allocate_raw_ptr(alloc, kind)
+        // We can `unwrap` since `alloc` contains no pointers.
+        self.allocate_raw_ptr(alloc, kind).unwrap()
     }
 
+    /// This can fail only of `alloc` contains relocations.
     pub fn allocate_raw_ptr(
         &mut self,
         alloc: Allocation,
         kind: MemoryKind<M::MemoryKind>,
-    ) -> Pointer<M::PointerTag> {
+    ) -> InterpResult<'tcx, Pointer<M::PointerTag>> {
         let id = self.tcx.reserve_alloc_id();
         debug_assert_ne!(
             Some(kind),
             M::GLOBAL_KIND.map(MemoryKind::Machine),
             "dynamically allocating global memory"
         );
-        let alloc = M::init_allocation_extra(self, id, Cow::Owned(alloc), Some(kind));
+        let alloc = M::init_allocation_extra(self, id, Cow::Owned(alloc), Some(kind))?;
         self.memory.alloc_map.insert(id, (kind, alloc.into_owned()));
-        M::tag_alloc_base_pointer(self, Pointer::from(id))
+        Ok(M::tag_alloc_base_pointer(self, Pointer::from(id)))
     }
 
     pub fn reallocate_ptr(
@@ -510,13 +513,12 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         };
         M::before_access_global(*self.tcx, &self.machine, id, alloc, def_id, is_write)?;
         // We got tcx memory. Let the machine initialize its "extra" stuff.
-        let alloc = M::init_allocation_extra(
+        M::init_allocation_extra(
             self,
             id, // always use the ID we got as input, not the "hidden" one.
             Cow::Borrowed(alloc.inner()),
             M::GLOBAL_KIND.map(MemoryKind::Machine),
-        );
-        Ok(alloc)
+        )
     }
 
     /// Gives raw access to the `Allocation`, without bounds or alignment checks.

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -244,12 +244,12 @@ impl<Tag> Allocation<Tag> {
 
 impl Allocation {
     /// Convert Tag and add Extra fields
-    pub fn convert_tag_add_extra<Tag, Extra>(
+    pub fn convert_tag_add_extra<Tag, Extra, Err>(
         self,
         cx: &impl HasDataLayout,
         extra: Extra,
-        mut tagger: impl FnMut(Pointer<AllocId>) -> Pointer<Tag>,
-    ) -> Allocation<Tag, Extra> {
+        mut tagger: impl FnMut(Pointer<AllocId>) -> Result<Pointer<Tag>, Err>,
+    ) -> Result<Allocation<Tag, Extra>, Err> {
         // Compute new pointer tags, which also adjusts the bytes.
         let mut bytes = self.bytes;
         let mut new_relocations = Vec::with_capacity(self.relocations.0.len());
@@ -260,19 +260,19 @@ impl Allocation {
             let ptr_bytes = &mut bytes[idx..idx + ptr_size];
             let bits = read_target_uint(endian, ptr_bytes).unwrap();
             let (ptr_tag, ptr_offset) =
-                tagger(Pointer::new(alloc_id, Size::from_bytes(bits))).into_parts();
+                tagger(Pointer::new(alloc_id, Size::from_bytes(bits)))?.into_parts();
             write_target_uint(endian, ptr_bytes, ptr_offset.bytes().into()).unwrap();
             new_relocations.push((offset, ptr_tag));
         }
         // Create allocation.
-        Allocation {
+        Ok(Allocation {
             bytes,
             relocations: Relocations::from_presorted(new_relocations),
             init_mask: self.init_mask,
             align: self.align,
             mutability: self.mutability,
             extra,
-        }
+        })
     }
 }