Rollup merge of rust-lang#141513 - nia-e:allocbytes-extend, r=RalfJung

interpret: add allocation parameters to `AllocBytes`

Necessary for a better implementation of [rust-lang/miri#4343](rust-lang/miri#4343). Also included here is the code from that PR, adapted to this new interface for the sake of example and so that CI can run on them; the Miri changes can be reverted and merged separately, though.

r? `@RalfJung`

Author: jieyouxu

compiler/rustc_const_eval/src/const_eval/dummy_machine.rs

@@ -197,4 +197,9 @@ impl<'tcx> interpret::Machine<'tcx> for DummyMachine {
     ) -> &'a mut Vec<interpret::Frame<'tcx, Self::Provenance, Self::FrameExtra>> {
         unimplemented!()
     }
+
+    fn get_default_alloc_params(
+        &self,
+    ) -> <Self::Bytes as rustc_middle::mir::interpret::AllocBytes>::AllocParams {
+    }
 }

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -735,6 +735,9 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
             Cow::Owned(compute_range())
         }
     }
+
+    fn get_default_alloc_params(&self) -> <Self::Bytes as mir::interpret::AllocBytes>::AllocParams {
+    }
 }
 
 // Please do not add any code below the above `Machine` trait impl. I (oli-obk) plan more cleanups

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -26,7 +26,7 @@ use crate::fluent_generated as fluent;
 /// Directly returns an `Allocation` containing an absolute path representation of the given type.
 pub(crate) fn alloc_type_name<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> ConstAllocation<'tcx> {
     let path = crate::util::type_name(tcx, ty);
-    let alloc = Allocation::from_bytes_byte_aligned_immutable(path.into_bytes());
+    let alloc = Allocation::from_bytes_byte_aligned_immutable(path.into_bytes(), ());
     tcx.mk_const_alloc(alloc)
 }
 

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -626,6 +626,10 @@ pub trait Machine<'tcx>: Sized {
         // Default to no caching.
         Cow::Owned(compute_range())
     }
+
+    /// Compute the value passed to the constructors of the `AllocBytes` type for
+    /// abstract machine allocations.
+    fn get_default_alloc_params(&self) -> <Self::Bytes as AllocBytes>::AllocParams;
 }
 
 /// A lot of the flexibility above is just needed for `Miri`, but all "compile-time" machines

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -233,10 +233,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         kind: MemoryKind<M::MemoryKind>,
         init: AllocInit,
     ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
+        let params = self.machine.get_default_alloc_params();
         let alloc = if M::PANIC_ON_ALLOC_FAIL {
-            Allocation::new(size, align, init)
+            Allocation::new(size, align, init, params)
         } else {
-            Allocation::try_new(size, align, init)?
+            Allocation::try_new(size, align, init, params)?
         };
         self.insert_allocation(alloc, kind)
     }
@@ -248,7 +249,8 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         kind: MemoryKind<M::MemoryKind>,
         mutability: Mutability,
     ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
-        let alloc = Allocation::from_bytes(bytes, align, mutability);
+        let params = self.machine.get_default_alloc_params();
+        let alloc = Allocation::from_bytes(bytes, align, mutability, params);
         self.insert_allocation(alloc, kind)
     }
 

compiler/rustc_const_eval/src/interpret/util.rs

@@ -38,7 +38,7 @@ pub(crate) fn create_static_alloc<'tcx>(
     static_def_id: LocalDefId,
     layout: TyAndLayout<'tcx>,
 ) -> InterpResult<'tcx, MPlaceTy<'tcx>> {
-    let alloc = Allocation::try_new(layout.size, layout.align.abi, AllocInit::Uninit)?;
+    let alloc = Allocation::try_new(layout.size, layout.align.abi, AllocInit::Uninit, ())?;
     let alloc_id = ecx.tcx.reserve_and_set_static_alloc(static_def_id.into());
     assert_eq!(ecx.machine.static_root_ids, None);
     ecx.machine.static_root_ids = Some((alloc_id, static_def_id));

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

@@ -27,12 +27,21 @@ use crate::ty;
 
 /// Functionality required for the bytes of an `Allocation`.
 pub trait AllocBytes: Clone + fmt::Debug + Deref<Target = [u8]> + DerefMut<Target = [u8]> {
+    /// The type of extra parameters passed in when creating an allocation.
+    /// Can be used by `interpret::Machine` instances to make runtime-configuration-dependent
+    /// decisions about the allocation strategy.
+    type AllocParams;
+
     /// Create an `AllocBytes` from a slice of `u8`.
-    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align) -> Self;
+    fn from_bytes<'a>(
+        slice: impl Into<Cow<'a, [u8]>>,
+        _align: Align,
+        _params: Self::AllocParams,
+    ) -> Self;
 
     /// Create a zeroed `AllocBytes` of the specified size and alignment.
     /// Returns `None` if we ran out of memory on the host.
-    fn zeroed(size: Size, _align: Align) -> Option<Self>;
+    fn zeroed(size: Size, _align: Align, _params: Self::AllocParams) -> Option<Self>;
 
     /// Gives direct access to the raw underlying storage.
     ///
@@ -51,11 +60,13 @@ pub trait AllocBytes: Clone + fmt::Debug + Deref<Target = [u8]> + DerefMut<Targe
 
 /// Default `bytes` for `Allocation` is a `Box<u8>`.
 impl AllocBytes for Box<[u8]> {
-    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align) -> Self {
+    type AllocParams = ();
+
+    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align, _params: ()) -> Self {
         Box::<[u8]>::from(slice.into())
     }
 
-    fn zeroed(size: Size, _align: Align) -> Option<Self> {
+    fn zeroed(size: Size, _align: Align, _params: ()) -> Option<Self> {
         let bytes = Box::<[u8]>::try_new_zeroed_slice(size.bytes().try_into().ok()?).ok()?;
         // SAFETY: the box was zero-allocated, which is a valid initial value for Box<[u8]>
         let bytes = unsafe { bytes.assume_init() };
@@ -172,9 +183,8 @@ fn all_zero(buf: &[u8]) -> bool {
 }
 
 /// Custom encoder for [`Allocation`] to more efficiently represent the case where all bytes are 0.
-impl<Prov: Provenance, Extra, Bytes, E: Encoder> Encodable<E> for Allocation<Prov, Extra, Bytes>
+impl<Prov: Provenance, Extra, E: Encoder> Encodable<E> for Allocation<Prov, Extra, Box<[u8]>>
 where
-    Bytes: AllocBytes,
     ProvenanceMap<Prov>: Encodable<E>,
     Extra: Encodable<E>,
 {
@@ -192,9 +202,8 @@ where
     }
 }
 
-impl<Prov: Provenance, Extra, Bytes, D: Decoder> Decodable<D> for Allocation<Prov, Extra, Bytes>
+impl<Prov: Provenance, Extra, D: Decoder> Decodable<D> for Allocation<Prov, Extra, Box<[u8]>>
 where
-    Bytes: AllocBytes,
     ProvenanceMap<Prov>: Decodable<D>,
     Extra: Decodable<D>,
 {
@@ -203,7 +212,7 @@ where
 
         let len = decoder.read_usize();
         let bytes = if all_zero { vec![0u8; len] } else { decoder.read_raw_bytes(len).to_vec() };
-        let bytes = Bytes::from_bytes(bytes, align);
+        let bytes = <Box<[u8]> as AllocBytes>::from_bytes(bytes, align, ());
 
         let provenance = Decodable::decode(decoder);
         let init_mask = Decodable::decode(decoder);
@@ -395,8 +404,9 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         slice: impl Into<Cow<'a, [u8]>>,
         align: Align,
         mutability: Mutability,
+        params: <Bytes as AllocBytes>::AllocParams,
     ) -> Self {
-        let bytes = Bytes::from_bytes(slice, align);
+        let bytes = Bytes::from_bytes(slice, align, params);
         let size = Size::from_bytes(bytes.len());
         Self {
             bytes,
@@ -408,14 +418,18 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         }
     }
 
-    pub fn from_bytes_byte_aligned_immutable<'a>(slice: impl Into<Cow<'a, [u8]>>) -> Self {
-        Allocation::from_bytes(slice, Align::ONE, Mutability::Not)
+    pub fn from_bytes_byte_aligned_immutable<'a>(
+        slice: impl Into<Cow<'a, [u8]>>,
+        params: <Bytes as AllocBytes>::AllocParams,
+    ) -> Self {
+        Allocation::from_bytes(slice, Align::ONE, Mutability::Not, params)
     }
 
     fn new_inner<R>(
         size: Size,
         align: Align,
         init: AllocInit,
+        params: <Bytes as AllocBytes>::AllocParams,
         fail: impl FnOnce() -> R,
     ) -> Result<Self, R> {
         // We raise an error if we cannot create the allocation on the host.
@@ -424,7 +438,7 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         // deterministic. However, we can be non-deterministic here because all uses of const
         // evaluation (including ConstProp!) will make compilation fail (via hard error
         // or ICE) upon encountering a `MemoryExhausted` error.
-        let bytes = Bytes::zeroed(size, align).ok_or_else(fail)?;
+        let bytes = Bytes::zeroed(size, align, params).ok_or_else(fail)?;
 
         Ok(Allocation {
             bytes,
@@ -444,8 +458,13 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
 
     /// Try to create an Allocation of `size` bytes, failing if there is not enough memory
     /// available to the compiler to do so.
-    pub fn try_new<'tcx>(size: Size, align: Align, init: AllocInit) -> InterpResult<'tcx, Self> {
-        Self::new_inner(size, align, init, || {
+    pub fn try_new<'tcx>(
+        size: Size,
+        align: Align,
+        init: AllocInit,
+        params: <Bytes as AllocBytes>::AllocParams,
+    ) -> InterpResult<'tcx, Self> {
+        Self::new_inner(size, align, init, params, || {
             ty::tls::with(|tcx| tcx.dcx().delayed_bug("exhausted memory during interpretation"));
             InterpErrorKind::ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted)
         })
@@ -457,8 +476,13 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
     ///
     /// Example use case: To obtain an Allocation filled with specific data,
     /// first call this function and then call write_scalar to fill in the right data.
-    pub fn new(size: Size, align: Align, init: AllocInit) -> Self {
-        match Self::new_inner(size, align, init, || {
+    pub fn new(
+        size: Size,
+        align: Align,
+        init: AllocInit,
+        params: <Bytes as AllocBytes>::AllocParams,
+    ) -> Self {
+        match Self::new_inner(size, align, init, params, || {
             panic!(
                 "interpreter ran out of memory: cannot create allocation of {} bytes",
                 size.bytes()

compiler/rustc_middle/src/ty/context.rs

@@ -1582,7 +1582,7 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns the same `AllocId` if called again with the same bytes.
     pub fn allocate_bytes_dedup(self, bytes: &[u8], salt: usize) -> interpret::AllocId {
         // Create an allocation that just contains these bytes.
-        let alloc = interpret::Allocation::from_bytes_byte_aligned_immutable(bytes);
+        let alloc = interpret::Allocation::from_bytes_byte_aligned_immutable(bytes, ());
         let alloc = self.mk_const_alloc(alloc);
         self.reserve_and_set_memory_dedup(alloc, salt)
     }

compiler/rustc_middle/src/ty/vtable.rs

@@ -110,7 +110,7 @@ pub(super) fn vtable_allocation_provider<'tcx>(
     let ptr_align = tcx.data_layout.pointer_align.abi;
 
     let vtable_size = ptr_size * u64::try_from(vtable_entries.len()).unwrap();
-    let mut vtable = Allocation::new(vtable_size, ptr_align, AllocInit::Uninit);
+    let mut vtable = Allocation::new(vtable_size, ptr_align, AllocInit::Uninit, ());
 
     // No need to do any alignment checks on the memory accesses below, because we know the
     // allocation is correctly aligned as we created it above. Also we're only offsetting by

compiler/rustc_mir_build/src/builder/expr/as_constant.rs

@@ -121,14 +121,14 @@ fn lit_to_mir_constant<'tcx>(tcx: TyCtxt<'tcx>, lit_input: LitToConstInput<'tcx>
     let value = match (lit, lit_ty.kind()) {
         (ast::LitKind::Str(s, _), ty::Ref(_, inner_ty, _)) if inner_ty.is_str() => {
             let s = s.as_str();
-            let allocation = Allocation::from_bytes_byte_aligned_immutable(s.as_bytes());
+            let allocation = Allocation::from_bytes_byte_aligned_immutable(s.as_bytes(), ());
             let allocation = tcx.mk_const_alloc(allocation);
             ConstValue::Slice { data: allocation, meta: allocation.inner().size().bytes() }
         }
         (ast::LitKind::ByteStr(data, _), ty::Ref(_, inner_ty, _))
             if matches!(inner_ty.kind(), ty::Slice(_)) =>
         {
-            let allocation = Allocation::from_bytes_byte_aligned_immutable(data as &[u8]);
+            let allocation = Allocation::from_bytes_byte_aligned_immutable(data as &[u8], ());
             let allocation = tcx.mk_const_alloc(allocation);
             ConstValue::Slice { data: allocation, meta: allocation.inner().size().bytes() }
         }
@@ -138,7 +138,7 @@ fn lit_to_mir_constant<'tcx>(tcx: TyCtxt<'tcx>, lit_input: LitToConstInput<'tcx>
         }
         (ast::LitKind::CStr(data, _), ty::Ref(_, inner_ty, _)) if matches!(inner_ty.kind(), ty::Adt(def, _) if tcx.is_lang_item(def.did(), LangItem::CStr)) =>
         {
-            let allocation = Allocation::from_bytes_byte_aligned_immutable(data as &[u8]);
+            let allocation = Allocation::from_bytes_byte_aligned_immutable(data as &[u8], ());
             let allocation = tcx.mk_const_alloc(allocation);
             ConstValue::Slice { data: allocation, meta: allocation.inner().size().bytes() }
         }

compiler/rustc_mir_transform/src/large_enums.rs

@@ -241,6 +241,7 @@ impl EnumSizeOpt {
             data,
             tcx.data_layout.ptr_sized_integer().align(&tcx.data_layout).abi,
             Mutability::Not,
+            (),
         );
         let alloc = tcx.reserve_and_set_memory_alloc(tcx.mk_const_alloc(alloc));
         Some((*adt_def, num_discrs, *alloc_cache.entry(ty).or_insert(alloc)))

compiler/rustc_smir/src/rustc_smir/alloc.rs

@@ -48,6 +48,7 @@ pub(crate) fn try_new_allocation<'tcx>(
                 size,
                 layout.align.abi,
                 AllocInit::Uninit,
+                (),
             );
             allocation
                 .write_scalar(&tables.tcx, alloc_range(Size::ZERO, size), scalar)
@@ -65,6 +66,7 @@ pub(crate) fn try_new_allocation<'tcx>(
                 layout.size,
                 layout.align.abi,
                 AllocInit::Uninit,
+                (),
             );
             allocation
                 .write_scalar(

src/tools/miri/src/alloc_addresses/mod.rs

@@ -139,7 +139,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 AllocKind::LiveData => {
                     if memory_kind == MiriMemoryKind::Global.into() {
                         // For new global allocations, we always pre-allocate the memory to be able use the machine address directly.
-                        let prepared_bytes = MiriAllocBytes::zeroed(info.size, info.align)
+                        let prepared_bytes = MiriAllocBytes::zeroed(info.size, info.align, ())
                             .unwrap_or_else(|| {
                                 panic!("Miri ran out of memory: cannot create allocation of {size:?} bytes", size = info.size)
                             });
@@ -159,7 +159,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 AllocKind::Function | AllocKind::VTable => {
                     // Allocate some dummy memory to get a unique address for this function/vtable.
                     let alloc_bytes =
-                        MiriAllocBytes::from_bytes(&[0u8; 1], Align::from_bytes(1).unwrap());
+                        MiriAllocBytes::from_bytes(&[0u8; 1], Align::from_bytes(1).unwrap(), ());
                     let ptr = alloc_bytes.as_ptr();
                     // Leak the underlying memory to ensure it remains unique.
                     std::mem::forget(alloc_bytes);
@@ -429,7 +429,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             prepared_alloc_bytes.copy_from_slice(bytes);
             interp_ok(prepared_alloc_bytes)
         } else {
-            interp_ok(MiriAllocBytes::from_bytes(std::borrow::Cow::Borrowed(bytes), align))
+            interp_ok(MiriAllocBytes::from_bytes(std::borrow::Cow::Borrowed(bytes), align, ()))
         }
     }
 

src/tools/miri/src/alloc_bytes.rs

@@ -24,7 +24,7 @@ impl Clone for MiriAllocBytes {
     fn clone(&self) -> Self {
         let bytes: Cow<'_, [u8]> = Cow::Borrowed(self);
         let align = Align::from_bytes(self.layout.align().to_u64()).unwrap();
-        MiriAllocBytes::from_bytes(bytes, align)
+        MiriAllocBytes::from_bytes(bytes, align, ())
     }
 }
 
@@ -86,7 +86,10 @@ impl MiriAllocBytes {
 }
 
 impl AllocBytes for MiriAllocBytes {
-    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align) -> Self {
+    /// Placeholder!
+    type AllocParams = ();
+
+    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align, _params: ()) -> Self {
         let slice = slice.into();
         let size = slice.len();
         let align = align.bytes();
@@ -102,7 +105,7 @@ impl AllocBytes for MiriAllocBytes {
         alloc_bytes
     }
 
-    fn zeroed(size: Size, align: Align) -> Option<Self> {
+    fn zeroed(size: Size, align: Align, _params: ()) -> Option<Self> {
         let size = size.bytes();
         let align = align.bytes();
         // SAFETY: `alloc_fn` will only be used with `size != 0`.

src/tools/miri/src/concurrency/thread.rs

@@ -902,7 +902,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             let mut alloc = alloc.inner().adjust_from_tcx(
                 &this.tcx,
                 |bytes, align| {
-                    interp_ok(MiriAllocBytes::from_bytes(std::borrow::Cow::Borrowed(bytes), align))
+                    interp_ok(MiriAllocBytes::from_bytes(std::borrow::Cow::Borrowed(bytes), align, ()))
                 },
                 |ptr| this.global_root_pointer(ptr),
             )?;

src/tools/miri/src/machine.rs

@@ -1803,6 +1803,9 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
     ) -> Cow<'e, RangeSet> {
         Cow::Borrowed(ecx.machine.union_data_ranges.entry(ty).or_insert_with(compute_range))
     }
+
+    /// Placeholder!
+    fn get_default_alloc_params(&self) -> <Self::Bytes as AllocBytes>::AllocParams { () }
 }
 
 /// Trait for callbacks handling asynchronous machine operations.