Rollup merge of #147718 - RalfJung:miri-allocator_shim_contents, r=bjorn3

miri: use allocator_shim_contents codegen helper

r? `@bjorn3`

Follow-up to rust-lang/rust#147526, also using that new infrastructure in Miri.

Author: matthiaskrgr

src/lib.rs

@@ -53,6 +53,7 @@
 extern crate rustc_abi;
 extern crate rustc_apfloat;
 extern crate rustc_ast;
+extern crate rustc_codegen_ssa;
 extern crate rustc_const_eval;
 extern crate rustc_data_structures;
 extern crate rustc_errors;

src/machine.rs

@@ -12,6 +12,8 @@ use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
 use rustc_abi::{Align, ExternAbi, Size};
 use rustc_apfloat::{Float, FloatConvert};
+use rustc_ast::expand::allocator::{self, SpecialAllocatorMethod};
+use rustc_data_structures::either::Either;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 #[allow(unused)]
 use rustc_data_structures::static_assert_size;
@@ -27,6 +29,7 @@ use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_session::config::InliningThreshold;
 use rustc_span::def_id::{CrateNum, DefId};
 use rustc_span::{Span, SpanData, Symbol};
+use rustc_symbol_mangling::mangle_internal_symbol;
 use rustc_target::callconv::FnAbi;
 
 use crate::alloc_addresses::EvalContextExt;
@@ -652,6 +655,10 @@ pub struct MiriMachine<'tcx> {
     pub(crate) pthread_rwlock_sanity: Cell<bool>,
     pub(crate) pthread_condvar_sanity: Cell<bool>,
 
+    /// (Foreign) symbols that are synthesized as part of the allocator shim: the key indicates the
+    /// name of the symbol being synthesized; the value indicates whether this should invoke some
+    /// other symbol or whether this has special allocator semantics.
+    pub(crate) allocator_shim_symbols: FxHashMap<Symbol, Either<Symbol, SpecialAllocatorMethod>>,
     /// Cache for `mangle_internal_symbol`.
     pub(crate) mangle_internal_symbol_cache: FxHashMap<&'static str, String>,
 
@@ -819,6 +826,7 @@ impl<'tcx> MiriMachine<'tcx> {
             pthread_mutex_sanity: Cell::new(false),
             pthread_rwlock_sanity: Cell::new(false),
             pthread_condvar_sanity: Cell::new(false),
+            allocator_shim_symbols: Self::allocator_shim_symbols(tcx),
             mangle_internal_symbol_cache: Default::default(),
             force_intrinsic_fallback: config.force_intrinsic_fallback,
             float_nondet: config.float_nondet,
@@ -827,6 +835,36 @@ impl<'tcx> MiriMachine<'tcx> {
         }
     }
 
+    fn allocator_shim_symbols(
+        tcx: TyCtxt<'tcx>,
+    ) -> FxHashMap<Symbol, Either<Symbol, SpecialAllocatorMethod>> {
+        use rustc_codegen_ssa::base::allocator_shim_contents;
+
+        // codegen uses `allocator_kind_for_codegen` here, but that's only needed to deal with
+        // dylibs which we do not support.
+        let Some(kind) = tcx.allocator_kind(()) else {
+            return Default::default();
+        };
+        let methods = allocator_shim_contents(tcx, kind);
+        let mut symbols = FxHashMap::default();
+        for method in methods {
+            let from_name = Symbol::intern(&mangle_internal_symbol(
+                tcx,
+                &allocator::global_fn_name(method.name),
+            ));
+            let to = match method.special {
+                Some(special) => Either::Right(special),
+                None =>
+                    Either::Left(Symbol::intern(&mangle_internal_symbol(
+                        tcx,
+                        &allocator::default_fn_name(method.name),
+                    ))),
+            };
+            symbols.try_insert(from_name, to).unwrap();
+        }
+        symbols
+    }
+
     pub(crate) fn late_init(
         ecx: &mut MiriInterpCx<'tcx>,
         config: &MiriConfig,
@@ -992,6 +1030,7 @@ impl VisitProvenance for MiriMachine<'_> {
             pthread_mutex_sanity: _,
             pthread_rwlock_sanity: _,
             pthread_condvar_sanity: _,
+            allocator_shim_symbols: _,
             mangle_internal_symbol_cache: _,
             force_intrinsic_fallback: _,
             float_nondet: _,

src/shims/alloc.rs

@@ -1,5 +1,8 @@
-use rustc_abi::{Align, Size};
-use rustc_ast::expand::allocator::AllocatorKind;
+use rustc_abi::{Align, AlignFromBytesError, CanonAbi, Size};
+use rustc_ast::expand::allocator::SpecialAllocatorMethod;
+use rustc_middle::ty::Ty;
+use rustc_span::Symbol;
+use rustc_target::callconv::FnAbi;
 
 use crate::*;
 
@@ -54,30 +57,100 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         Align::from_bytes(prev_power_of_two(size)).unwrap()
     }
 
-    /// Emulates calling the internal __rust_* allocator functions
-    fn emulate_allocator(
+    /// Check some basic requirements for this allocation request:
+    /// non-zero size, power-of-two alignment.
+    fn check_rust_alloc_request(&self, size: u64, align: u64) -> InterpResult<'tcx> {
+        let this = self.eval_context_ref();
+        if size == 0 {
+            throw_ub_format!("creating allocation with size 0");
+        }
+        if size > this.max_size_of_val().bytes() {
+            throw_ub_format!("creating an allocation larger than half the address space");
+        }
+        if let Err(e) = Align::from_bytes(align) {
+            match e {
+                AlignFromBytesError::TooLarge(_) => {
+                    throw_unsup_format!(
+                        "creating allocation with alignment {align} exceeding rustc's maximum \
+                         supported value"
+                    );
+                }
+                AlignFromBytesError::NotPowerOfTwo(_) => {
+                    throw_ub_format!("creating allocation with non-power-of-two alignment {align}");
+                }
+            }
+        }
+
+        interp_ok(())
+    }
+
+    fn rust_special_allocator_method(
         &mut self,
-        default: impl FnOnce(&mut MiriInterpCx<'tcx>) -> InterpResult<'tcx>,
-    ) -> InterpResult<'tcx, EmulateItemResult> {
+        method: SpecialAllocatorMethod,
+        link_name: Symbol,
+        abi: &FnAbi<'tcx, Ty<'tcx>>,
+        args: &[OpTy<'tcx>],
+        dest: &PlaceTy<'tcx>,
+    ) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
 
-        let Some(allocator_kind) = this.tcx.allocator_kind(()) else {
-            // in real code, this symbol does not exist without an allocator
-            return interp_ok(EmulateItemResult::NotSupported);
-        };
+        match method {
+            SpecialAllocatorMethod::Alloc | SpecialAllocatorMethod::AllocZeroed => {
+                let [size, align] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
+                let size = this.read_target_usize(size)?;
+                let align = this.read_target_usize(align)?;
+
+                this.check_rust_alloc_request(size, align)?;
 
-        match allocator_kind {
-            AllocatorKind::Global => {
-                // When `#[global_allocator]` is used, `__rust_*` is defined by the macro expansion
-                // of this attribute. As such we have to call an exported Rust function,
-                // and not execute any Miri shim. Somewhat unintuitively doing so is done
-                // by returning `NotSupported`, which triggers the `lookup_exported_symbol`
-                // fallback case in `emulate_foreign_item`.
-                interp_ok(EmulateItemResult::NotSupported)
+                let ptr = this.allocate_ptr(
+                    Size::from_bytes(size),
+                    Align::from_bytes(align).unwrap(),
+                    MiriMemoryKind::Rust.into(),
+                    if matches!(method, SpecialAllocatorMethod::AllocZeroed) {
+                        AllocInit::Zero
+                    } else {
+                        AllocInit::Uninit
+                    },
+                )?;
+
+                this.write_pointer(ptr, dest)
+            }
+            SpecialAllocatorMethod::Dealloc => {
+                let [ptr, old_size, align] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
+                let ptr = this.read_pointer(ptr)?;
+                let old_size = this.read_target_usize(old_size)?;
+                let align = this.read_target_usize(align)?;
+
+                // No need to check old_size/align; we anyway check that they match the allocation.
+                this.deallocate_ptr(
+                    ptr,
+                    Some((Size::from_bytes(old_size), Align::from_bytes(align).unwrap())),
+                    MiriMemoryKind::Rust.into(),
+                )
             }
-            AllocatorKind::Default => {
-                default(this)?;
-                interp_ok(EmulateItemResult::NeedsReturn)
+            SpecialAllocatorMethod::Realloc => {
+                let [ptr, old_size, align, new_size] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
+                let ptr = this.read_pointer(ptr)?;
+                let old_size = this.read_target_usize(old_size)?;
+                let align = this.read_target_usize(align)?;
+                let new_size = this.read_target_usize(new_size)?;
+                // No need to check old_size; we anyway check that they match the allocation.
+
+                this.check_rust_alloc_request(new_size, align)?;
+
+                let align = Align::from_bytes(align).unwrap();
+                let new_ptr = this.reallocate_ptr(
+                    ptr,
+                    Some((Size::from_bytes(old_size), align)),
+                    Size::from_bytes(new_size),
+                    align,
+                    MiriMemoryKind::Rust.into(),
+                    AllocInit::Uninit,
+                )?;
+                this.write_pointer(new_ptr, dest)
             }
         }
     }