Handle oom strategy in a weak symbol like manner

bjorn3 · bjorn3 · commit c2a135b8f6e7 · 2025-10-09T14:58:18.000Z
diff --git a/compiler/rustc_builtin_macros/src/lib.rs b/compiler/rustc_builtin_macros/src/lib.rs
@@ -56,6 +56,7 @@ mod trace_macros;
 pub mod asm;
 pub mod cmdline_attrs;
 pub mod contracts;
+pub mod oom_strategy;
 pub mod proc_macro_harness;
 pub mod standard_library_imports;
 pub mod test_harness;
diff --git a/compiler/rustc_builtin_macros/src/oom_strategy.rs b/compiler/rustc_builtin_macros/src/oom_strategy.rs
@@ -0,0 +1,71 @@
+use rustc_ast::expand::allocator::global_fn_name;
+use rustc_ast::{ast, token};
+use rustc_expand::base::{ExtCtxt, ResolverExpand};
+use rustc_expand::expand::{AstFragment, ExpansionConfig};
+use rustc_feature::Features;
+use rustc_session::Session;
+use rustc_session::config::OomStrategy;
+use rustc_span::hygiene::AstPass;
+use rustc_span::{DUMMY_SP, Ident, Symbol, kw, sym};
+use smallvec::smallvec;
+use thin_vec::thin_vec;
+
+pub fn inject(
+    krate: &mut ast::Crate,
+    sess: &Session,
+    features: &Features,
+    resolver: &mut dyn ResolverExpand,
+) {
+    if sess.opts.unstable_opts.oom == OomStrategy::Abort {
+        // Handled in the allocator shim
+        return;
+    }
+
+    let econfig = ExpansionConfig::default(sym::alloc_error_handler, features);
+    let mut ecx = ExtCtxt::new(sess, econfig, resolver, None);
+
+    let expn_id = ecx.resolver.expansion_for_ast_pass(
+        DUMMY_SP,
+        AstPass::OomStrategy,
+        &[sym::rustc_attrs],
+        None,
+    );
+    let sp = DUMMY_SP.with_def_site_ctxt(expn_id.to_expn_id());
+
+    let decl = ecx.fn_decl(
+        thin_vec![],
+        ast::FnRetTy::Ty(
+            ecx.ty(sp, ast::TyKind::Path(None, ecx.path_ident(sp, Ident::new(sym::u8, sp)))),
+        ),
+    );
+    let sig = ast::FnSig { decl, header: ast::FnHeader::default(), span: sp };
+    let oom_strategy_fn = Box::new(ast::Fn {
+        defaultness: ast::Defaultness::Final,
+        ident: Ident::new(Symbol::intern(&global_fn_name(OomStrategy::SYMBOL)), sp),
+        generics: ast::Generics::default(),
+        sig,
+        contract: None,
+        define_opaque: None,
+        body: Some(ecx.block_expr({
+            let suffix = Some(ast::UintTy::U8.name());
+            let lit = token::Lit::new(token::Integer, sym::integer(1), suffix);
+            ecx.expr(sp, ast::ExprKind::Lit(lit))
+        })),
+    });
+
+    let oom_strategy_item = ecx.item(
+        sp,
+        thin_vec![ecx.attr_word(sym::rustc_std_internal_symbol, sp)],
+        ast::ItemKind::Fn(oom_strategy_fn),
+    );
+
+    let items = AstFragment::Items(smallvec![ecx.item_const(
+        sp,
+        Ident::new(kw::Underscore, sp),
+        ecx.ty(sp, ast::TyKind::Tup(thin_vec![])),
+        ecx.expr_block(ecx.block(sp, thin_vec![ecx.stmt_item(sp, oom_strategy_item)])),
+    )]);
+    krate
+        .items
+        .push(ecx.monotonic_expander().fully_expand_fragment(items).make_items().pop().unwrap());
+}
diff --git a/compiler/rustc_codegen_cranelift/src/allocator.rs b/compiler/rustc_codegen_cranelift/src/allocator.rs
@@ -15,13 +15,7 @@ use crate::prelude::*;
 /// Returns whether an allocator shim was created
 pub(crate) fn codegen(tcx: TyCtxt<'_>, module: &mut dyn Module) -> bool {
     let Some(kind) = allocator_kind_for_codegen(tcx) else { return false };
-    codegen_inner(
-        tcx,
-        module,
-        kind,
-        tcx.alloc_error_handler_kind(()).unwrap(),
-        tcx.sess.opts.unstable_opts.oom,
-    );
+    codegen_inner(tcx, module, kind, tcx.alloc_error_handler_kind(()).unwrap());
     true
 }
 
@@ -30,7 +24,6 @@ fn codegen_inner(
     module: &mut dyn Module,
     kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
-    oom_strategy: OomStrategy,
 ) {
     let usize_ty = module.target_config().pointer_type();
 
@@ -86,33 +79,18 @@ fn codegen_inner(
         );
     }
 
-    {
+    if tcx.sess.opts.unstable_opts.oom == OomStrategy::Abort {
         let sig = Signature {
             call_conv: module.target_config().default_call_conv,
             params: vec![],
             returns: vec![AbiParam::new(types::I8)],
         };
-        let func_id = module
-            .declare_function(
-                &mangle_internal_symbol(tcx, OomStrategy::SYMBOL),
-                Linkage::Export,
-                &sig,
-            )
-            .unwrap();
-        let mut ctx = Context::new();
-        ctx.func.signature = sig;
-        {
-            let mut func_ctx = FunctionBuilderContext::new();
-            let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
-
-            let block = bcx.create_block();
-            bcx.switch_to_block(block);
-            let value = bcx.ins().iconst(types::I8, oom_strategy.should_panic() as i64);
-            bcx.ins().return_(&[value]);
-            bcx.seal_all_blocks();
-            bcx.finalize();
-        }
-        module.define_function(func_id, &mut ctx).unwrap();
+        crate::common::create_wrapper_function(
+            module,
+            sig,
+            &mangle_internal_symbol(tcx, &global_fn_name(OomStrategy::SYMBOL)),
+            &mangle_internal_symbol(tcx, &default_fn_name(OomStrategy::SYMBOL)),
+        );
     }
 
     {
diff --git a/compiler/rustc_codegen_gcc/src/allocator.rs b/compiler/rustc_codegen_gcc/src/allocator.rs
@@ -1,6 +1,6 @@
 #[cfg(feature = "master")]
 use gccjit::FnAttribute;
-use gccjit::{Context, FunctionType, RValue, ToRValue, Type};
+use gccjit::{Context, FunctionType, ToRValue, Type};
 use rustc_ast::expand::allocator::{
     ALLOC_ERROR_HANDLER, ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
     default_fn_name, global_fn_name,
@@ -73,13 +73,16 @@ pub(crate) unsafe fn codegen(
         );
     }
 
-    create_const_value_function(
-        tcx,
-        context,
-        &mangle_internal_symbol(tcx, OomStrategy::SYMBOL),
-        i8,
-        context.new_rvalue_from_int(i8, tcx.sess.opts.unstable_opts.oom.should_panic() as i32),
-    );
+    if tcx.sess.opts.unstable_opts.oom == OomStrategy::Abort {
+        create_wrapper_function(
+            tcx,
+            context,
+            &mangle_internal_symbol(tcx, &global_fn_name(OomStrategy::SYMBOL)),
+            Some(&mangle_internal_symbol(tcx, &default_fn_name(OomStrategy::SYMBOL))),
+            &[],
+            Some(i8),
+        );
+    }
 
     create_wrapper_function(
         tcx,
@@ -91,34 +94,6 @@ pub(crate) unsafe fn codegen(
     );
 }
 
-fn create_const_value_function(
-    tcx: TyCtxt<'_>,
-    context: &Context<'_>,
-    name: &str,
-    output: Type<'_>,
-    value: RValue<'_>,
-) {
-    let func = context.new_function(None, FunctionType::Exported, output, &[], name, false);
-
-    #[cfg(feature = "master")]
-    {
-        func.add_attribute(FnAttribute::Visibility(symbol_visibility_to_gcc(
-            tcx.sess.default_visibility(),
-        )));
-
-        // FIXME(antoyo): cg_llvm sets AlwaysInline, but AlwaysInline is different in GCC and using
-        // it here will causes linking errors when using LTO.
-        func.add_attribute(FnAttribute::Inline);
-    }
-
-    if tcx.sess.must_emit_unwind_tables() {
-        // TODO(antoyo): emit unwind tables.
-    }
-
-    let block = func.new_block("entry");
-    block.end_with_return(None, value);
-}
-
 fn create_wrapper_function(
     tcx: TyCtxt<'_>,
     context: &Context<'_>,
diff --git a/compiler/rustc_codegen_llvm/src/allocator.rs b/compiler/rustc_codegen_llvm/src/allocator.rs
@@ -14,7 +14,7 @@ use rustc_symbol_mangling::mangle_internal_symbol;
 use crate::attributes::llfn_attrs_from_instance;
 use crate::builder::SBuilder;
 use crate::declare::declare_simple_fn;
-use crate::llvm::{self, FALSE, FromGeneric, TRUE, Type, Value};
+use crate::llvm::{self, FromGeneric, TRUE, Type};
 use crate::{SimpleCx, attributes, debuginfo};
 
 pub(crate) unsafe fn codegen(
@@ -97,16 +97,19 @@ pub(crate) unsafe fn codegen(
         );
     }
 
-    // __rust_alloc_error_handler_should_panic_v2
-    create_const_value_function(
-        tcx,
-        &cx,
-        &mangle_internal_symbol(tcx, OomStrategy::SYMBOL),
-        &i8,
-        unsafe {
-            llvm::LLVMConstInt(i8, tcx.sess.opts.unstable_opts.oom.should_panic() as u64, FALSE)
-        },
-    );
+    if tcx.sess.opts.unstable_opts.oom == OomStrategy::Abort {
+        // __rust_alloc_error_handler_should_panic_v2
+        create_wrapper_function(
+            tcx,
+            &cx,
+            &mangle_internal_symbol(tcx, &global_fn_name(OomStrategy::SYMBOL)),
+            Some(&mangle_internal_symbol(tcx, &default_fn_name(OomStrategy::SYMBOL))),
+            &[],
+            Some(i8),
+            false,
+            &CodegenFnAttrs::new(),
+        );
+    }
 
     // __rust_no_alloc_shim_is_unstable_v2
     create_wrapper_function(
@@ -127,34 +130,6 @@ pub(crate) unsafe fn codegen(
     }
 }
 
-fn create_const_value_function(
-    tcx: TyCtxt<'_>,
-    cx: &SimpleCx<'_>,
-    name: &str,
-    output: &Type,
-    value: &Value,
-) {
-    let ty = cx.type_func(&[], output);
-    let llfn = declare_simple_fn(
-        &cx,
-        name,
-        llvm::CallConv::CCallConv,
-        llvm::UnnamedAddr::Global,
-        llvm::Visibility::from_generic(tcx.sess.default_visibility()),
-        ty,
-    );
-
-    attributes::apply_to_llfn(
-        llfn,
-        llvm::AttributePlace::Function,
-        &[llvm::AttributeKind::AlwaysInline.create_attr(cx.llcx)],
-    );
-
-    let llbb = unsafe { llvm::LLVMAppendBasicBlockInContext(cx.llcx, llfn, c"entry".as_ptr()) };
-    let mut bx = SBuilder::build(&cx, llbb);
-    bx.ret(value);
-}
-
 fn create_wrapper_function(
     tcx: TyCtxt<'_>,
     cx: &SimpleCx<'_>,
diff --git a/compiler/rustc_codegen_ssa/src/back/symbol_export.rs b/compiler/rustc_codegen_ssa/src/back/symbol_export.rs
@@ -501,7 +501,7 @@ pub(crate) fn allocator_shim_symbols(
         .map(move |method| mangle_internal_symbol(tcx, global_fn_name(method.name).as_str()))
         .chain([
             mangle_internal_symbol(tcx, global_fn_name(ALLOC_ERROR_HANDLER).as_str()),
-            mangle_internal_symbol(tcx, OomStrategy::SYMBOL),
+            mangle_internal_symbol(tcx, global_fn_name(OomStrategy::SYMBOL).as_str()),
             mangle_internal_symbol(tcx, NO_ALLOC_SHIM_IS_UNSTABLE),
         ])
         .map(move |symbol_name| {
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
@@ -255,6 +255,8 @@ fn configure_and_expand(
         rustc_builtin_macros::test_harness::inject(&mut krate, sess, features, resolver)
     });
 
+    rustc_builtin_macros::oom_strategy::inject(&mut krate, sess, features, resolver);
+
     let has_proc_macro_decls = sess.time("AST_validation", || {
         rustc_ast_passes::ast_validation::check_crate(
             sess,
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
@@ -3465,14 +3465,7 @@ pub enum OomStrategy {
 }
 
 impl OomStrategy {
-    pub const SYMBOL: &'static str = "__rust_alloc_error_handler_should_panic_v2";
-
-    pub fn should_panic(self) -> u8 {
-        match self {
-            OomStrategy::Panic => 1,
-            OomStrategy::Abort => 0,
-        }
-    }
+    pub const SYMBOL: Symbol = sym::alloc_error_handler_should_panic_v2;
 }
 
 /// How to run proc-macro code when building this crate
diff --git a/compiler/rustc_span/src/hygiene.rs b/compiler/rustc_span/src/hygiene.rs
@@ -1204,6 +1204,7 @@ pub enum AstPass {
     StdImports,
     TestHarness,
     ProcMacroHarness,
+    OomStrategy,
 }
 
 impl AstPass {
@@ -1212,6 +1213,7 @@ impl AstPass {
             AstPass::StdImports => "standard library imports",
             AstPass::TestHarness => "test harness",
             AstPass::ProcMacroHarness => "proc macro harness",
+            AstPass::OomStrategy => "oom strategy definition",
         }
     }
 }
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -437,6 +437,7 @@ symbols! {
         all,
         alloc,
         alloc_error_handler,
+        alloc_error_handler_should_panic_v2,
         alloc_layout,
         alloc_zeroed,
         allocator,
diff --git a/library/alloc/src/alloc.rs b/library/alloc/src/alloc.rs
@@ -422,6 +422,11 @@ pub const fn handle_alloc_error(layout: Layout) -> ! {
 #[allow(unused_attributes)]
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub mod __alloc_error_handler {
+    #[rustc_std_internal_symbol]
+    fn __rdl_alloc_error_handler_should_panic_v2() -> u8 {
+        0
+    }
+
     // called via generated `__rust_alloc_error_handler` if there is no
     // `#[alloc_error_handler]`.
     #[rustc_std_internal_symbol]
diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs
@@ -11,6 +11,7 @@ use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::AllocInit;
 use rustc_middle::ty::{Instance, Ty};
 use rustc_middle::{mir, ty};
+use rustc_session::config::OomStrategy;
 use rustc_span::Symbol;
 use rustc_target::callconv::FnAbi;
 
@@ -69,6 +70,19 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     return interp_ok(Some(handler));
                 }
             }
+            name if name
+                == this.mangle_internal_symbol("__rust_alloc_error_handler_should_panic_v2") =>
+            {
+                // Gets the value of the `oom` option.
+                let [] = this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
+                if this.tcx.sess.opts.unstable_opts.oom == OomStrategy::Abort {
+                    let name = Symbol::intern(
+                        this.mangle_internal_symbol("__rdl_alloc_error_handler_should_panic_v2"),
+                    );
+                    let handler = this.lookup_exported_symbol(name)?.expect("missing oom symbol");
+                    return interp_ok(Some(handler));
+                }
+            }
             _ => {}
         }
 
@@ -681,14 +695,6 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 // This is a no-op shim that only exists to prevent making the allocator shims instantly stable.
                 let [] = this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
             }
-            name if name
-                == this.mangle_internal_symbol("__rust_alloc_error_handler_should_panic_v2") =>
-            {
-                // Gets the value of the `oom` option.
-                let [] = this.check_shim_sig_lenient(abi, CanonAbi::Rust, link_name, args)?;
-                let val = this.tcx.sess.opts.unstable_opts.oom.should_panic();
-                this.write_int(val, dest)?;
-            }
 
             // C memory handling functions
             "memcmp" => {

Original file line number	Diff line number	Diff line change
`@@ -3465,14 +3465,7 @@ pub enum OomStrategy {`
`3465`	`3465`	`}`
`3466`	`3466`
`3467`	`3467`	`impl OomStrategy {`
`3468`		`- pub const SYMBOL: &'static str = "__rust_alloc_error_handler_should_panic_v2";`
`3469`		`-`
`3470`		`- pub fn should_panic(self) -> u8 {`
`3471`		`- match self {`
`3472`		`- OomStrategy::Panic => 1,`
`3473`		`- OomStrategy::Abort => 0,`
`3474`		`- }`
`3475`		`- }`
	`3468`	`+ pub const SYMBOL: Symbol = sym::alloc_error_handler_should_panic_v2;`
`3476`	`3469`	`}`
`3477`	`3470`
`3478`	`3471`	`/// How to run proc-macro code when building this crate`
Original file line number	Diff line number	Diff line change
`@@ -1204,6 +1204,7 @@ pub enum AstPass {`
`1204`	`1204`	`StdImports,`
`1205`	`1205`	`TestHarness,`
`1206`	`1206`	`ProcMacroHarness,`
	`1207`	`+ OomStrategy,`
`1207`	`1208`	`}`
`1208`	`1209`
`1209`	`1210`	`impl AstPass {`
`@@ -1212,6 +1213,7 @@ impl AstPass {`
`1212`	`1213`	`AstPass::StdImports => "standard library imports",`
`1213`	`1214`	`AstPass::TestHarness => "test harness",`
`1214`	`1215`	`AstPass::ProcMacroHarness => "proc macro harness",`
	`1216`	`+ AstPass::OomStrategy => "oom strategy definition",`
`1215`	`1217`	`}`
`1216`	`1218`	`}`
`1217`	`1219`	`}`