Make platform-specific OOM handling available to liballoc when libstd is linked

Author: SimonSapin

src/librustc/session/mod.rs

@@ -118,6 +118,7 @@ pub struct Session {
     /// injected.
     pub injected_allocator: Once<Option<CrateNum>>,
     pub allocator_kind: Once<Option<AllocatorKind>>,
+    pub injected_default_alloc_error_hook: Once<InjectedDefaultOomHook>,
     pub injected_panic_runtime: Once<Option<CrateNum>>,
 
     /// Map from imported macro spans (which consist of
@@ -175,6 +176,12 @@ pub struct PerfStats {
     pub normalize_projection_ty: AtomicUsize,
 }
 
+pub enum InjectedDefaultOomHook {
+    None,
+    Noop,
+    Platform,
+}
+
 /// Enum to support dispatch of one-time diagnostics (in Session.diag_once)
 enum DiagnosticBuilderMethod {
     Note,
@@ -1119,6 +1126,7 @@ pub fn build_session_(
         next_node_id: OneThread::new(Cell::new(NodeId::new(1))),
         injected_allocator: Once::new(),
         allocator_kind: Once::new(),
+        injected_default_alloc_error_hook: Once::new(),
         injected_panic_runtime: Once::new(),
         imported_macro_spans: OneThread::new(RefCell::new(HashMap::new())),
         incr_comp_session: OneThread::new(RefCell::new(IncrCompSession::NotInitialized)),

src/librustc_allocator/lib.rs

@@ -42,6 +42,19 @@ pub static ALLOCATOR_METHODS: &[AllocatorMethod] = &[
     },
 ];
 
+pub static OOM_HANDLING_METHODS: &[AllocatorMethod] = &[
+    AllocatorMethod {
+        name: "default_alloc_error_hook",
+        inputs: &[AllocatorTy::Layout],
+        output: AllocatorTy::Unit,
+    },
+    AllocatorMethod {
+        name: "abort_internal",
+        inputs: &[],
+        output: AllocatorTy::Unit,
+    },
+];
+
 pub struct AllocatorMethod {
     pub name: &'static str,
     pub inputs: &'static [AllocatorTy],

src/librustc_codegen_llvm/allocator.rs

@@ -14,8 +14,9 @@ use std::ptr;
 use attributes;
 use libc::c_uint;
 use rustc::middle::allocator::AllocatorKind;
+use rustc::session::InjectedDefaultOomHook;
 use rustc::ty::TyCtxt;
-use rustc_allocator::{ALLOCATOR_METHODS, AllocatorTy};
+use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS, AllocatorTy};
 
 use ModuleLlvm;
 use llvm::{self, False, True};
@@ -33,9 +34,10 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
     let i8p = llvm::LLVMPointerType(i8, 0);
     let void = llvm::LLVMVoidTypeInContext(llcx);
 
-    for method in ALLOCATOR_METHODS {
+    let build = |name: String, inputs: &[AllocatorTy], output: &AllocatorTy,
+                 callee: Option<String>| {
         let mut args = Vec::new();
-        for ty in method.inputs.iter() {
+        for ty in inputs.iter() {
             match *ty {
                 AllocatorTy::Layout => {
                     args.push(usize); // size
@@ -48,7 +50,7 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
                 AllocatorTy::Unit => panic!("invalid allocator arg"),
             }
         }
-        let output = match method.output {
+        let output = match *output {
             AllocatorTy::ResultPtr => Some(i8p),
             AllocatorTy::Unit => None,
 
@@ -60,29 +62,40 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
                                         args.as_ptr(),
                                         args.len() as c_uint,
                                         False);
-        let name = CString::new(format!("__rust_{}", method.name)).unwrap();
+        let name = CString::new(name).unwrap();
         let llfn = llvm::LLVMRustGetOrInsertFunction(llmod,
                                                      name.as_ptr(),
                                                      ty);
 
         if tcx.sess.target.target.options.default_hidden_visibility {
             llvm::LLVMRustSetVisibility(llfn, llvm::Visibility::Hidden);
         }
-       if tcx.sess.target.target.options.requires_uwtable {
-           attributes::emit_uwtable(llfn, true);
-       }
-
-        let callee = CString::new(kind.fn_name(method.name)).unwrap();
-        let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
-                                                       callee.as_ptr(),
-                                                       ty);
+        if tcx.sess.target.target.options.requires_uwtable {
+            attributes::emit_uwtable(llfn, true);
+        }
 
         let llbb = llvm::LLVMAppendBasicBlockInContext(llcx,
                                                        llfn,
                                                        "entry\0".as_ptr() as *const _);
 
         let llbuilder = llvm::LLVMCreateBuilderInContext(llcx);
         llvm::LLVMPositionBuilderAtEnd(llbuilder, llbb);
+
+        let callee = if let Some(callee) = callee {
+            callee
+        } else {
+            // Generate a no-op function
+            llvm::LLVMBuildRetVoid(llbuilder);
+            llvm::LLVMDisposeBuilder(llbuilder);
+            return
+        };
+
+        // Forward the call to another function
+        let callee = CString::new(callee).unwrap();
+        let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
+                                                       callee.as_ptr(),
+                                                       ty);
+
         let args = args.iter().enumerate().map(|(i, _)| {
             llvm::LLVMGetParam(llfn, i as c_uint)
         }).collect::<Vec<_>>();
@@ -98,6 +111,28 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
         } else {
             llvm::LLVMBuildRetVoid(llbuilder);
         }
+
         llvm::LLVMDisposeBuilder(llbuilder);
+    };
+
+    for method in ALLOCATOR_METHODS {
+        let name = format!("__rust_{}", method.name);
+        build(name, method.inputs, &method.output, Some(kind.fn_name(method.name)))
+    }
+
+    let has_plaftom_functions = match tcx.sess.injected_default_alloc_error_hook.get() {
+        InjectedDefaultOomHook::None => return,
+        InjectedDefaultOomHook::Noop => false,
+        InjectedDefaultOomHook::Platform => true,
+    };
+
+    for method in OOM_HANDLING_METHODS {
+        let callee = if has_plaftom_functions {
+            Some(format!("__rust_{}", method.name))
+        } else {
+            None
+        };
+        let name = format!("__rust_maybe_{}", method.name);
+        build(name, method.inputs, &method.output, callee)
     }
 }

src/librustc_codegen_llvm/back/symbol_export.rs

@@ -17,12 +17,12 @@ use rustc::hir::CodegenFnAttrFlags;
 use rustc::hir::def_id::{CrateNum, DefId, LOCAL_CRATE, CRATE_DEF_INDEX};
 use rustc::ich::Fingerprint;
 use rustc::middle::exported_symbols::{SymbolExportLevel, ExportedSymbol, metadata_symbol_name};
-use rustc::session::config;
+use rustc::session::{config, InjectedDefaultOomHook};
 use rustc::ty::{TyCtxt, SymbolName};
 use rustc::ty::query::Providers;
 use rustc::ty::subst::Substs;
 use rustc::util::nodemap::{FxHashMap, DefIdMap};
-use rustc_allocator::ALLOCATOR_METHODS;
+use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS};
 use rustc_data_structures::indexed_vec::IndexVec;
 use std::collections::hash_map::Entry::*;
 
@@ -222,6 +222,19 @@ fn exported_symbols_provider_local<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         }
     }
 
+    match tcx.sess.injected_default_alloc_error_hook.get() {
+        InjectedDefaultOomHook::None => {}
+        InjectedDefaultOomHook::Noop |
+        InjectedDefaultOomHook::Platform => {
+            for method in OOM_HANDLING_METHODS {
+                let symbol_name = format!("__rust_{}", method.name);
+                let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(&symbol_name));
+
+                symbols.push((exported_symbol, SymbolExportLevel::Rust));
+            }
+        }
+    }
+
     if tcx.sess.opts.debugging_opts.pgo_gen.is_some() {
         // These are weak symbols that point to the profile version and the
         // profile name, which need to be treated as exported so LTO doesn't nix

src/librustc_metadata/creader.rs

@@ -20,7 +20,7 @@ use rustc::hir::svh::Svh;
 use rustc::middle::allocator::AllocatorKind;
 use rustc::middle::cstore::DepKind;
 use rustc::mir::interpret::AllocDecodingState;
-use rustc::session::{Session, CrateDisambiguator};
+use rustc::session::{Session, CrateDisambiguator, InjectedDefaultOomHook};
 use rustc::session::config::{Sanitizer, self};
 use rustc_target::spec::{PanicStrategy, TargetTriple};
 use rustc::session::search_paths::PathKind;
@@ -841,6 +841,7 @@ impl<'a> CrateLoader<'a> {
         if !needs_allocator {
             self.sess.injected_allocator.set(None);
             self.sess.allocator_kind.set(None);
+            self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
             return
         }
 
@@ -862,9 +863,22 @@ impl<'a> CrateLoader<'a> {
         if !need_lib_alloc && !need_exe_alloc {
             self.sess.injected_allocator.set(None);
             self.sess.allocator_kind.set(None);
+            self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
             return
         }
 
+        let mut has_default_lib_allocator = attr::contains_name(&krate.attrs,
+                                                                "default_lib_allocator");;
+        self.cstore.iter_crate_data(|_, data| {
+            has_default_lib_allocator = has_default_lib_allocator ||
+                                        data.root.has_default_lib_allocator;
+        });
+        self.sess.injected_default_alloc_error_hook.set(if has_default_lib_allocator {
+            InjectedDefaultOomHook::Platform
+        } else {
+            InjectedDefaultOomHook::Noop
+        });
+
         // Ok, we need an allocator. Not only that but we're actually going to
         // create an artifact that needs one linked in. Let's go find the one
         // that we're going to link in.