Skip to content

Commit 68dfaed

Browse files
committed
Make platform-specific OOM handling available to liballoc when libstd is linked
1 parent 3f9ad0a commit 68dfaed

File tree

5 files changed

+99
-16
lines changed

5 files changed

+99
-16
lines changed

src/librustc/session/mod.rs

+8
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ pub struct Session {
118118
/// injected.
119119
pub injected_allocator: Once<Option<CrateNum>>,
120120
pub allocator_kind: Once<Option<AllocatorKind>>,
121+
pub injected_default_alloc_error_hook: Once<InjectedDefaultOomHook>,
121122
pub injected_panic_runtime: Once<Option<CrateNum>>,
122123

123124
/// Map from imported macro spans (which consist of
@@ -175,6 +176,12 @@ pub struct PerfStats {
175176
pub normalize_projection_ty: AtomicUsize,
176177
}
177178

179+
pub enum InjectedDefaultOomHook {
180+
None,
181+
Noop,
182+
Platform,
183+
}
184+
178185
/// Enum to support dispatch of one-time diagnostics (in Session.diag_once)
179186
enum DiagnosticBuilderMethod {
180187
Note,
@@ -1119,6 +1126,7 @@ pub fn build_session_(
11191126
next_node_id: OneThread::new(Cell::new(NodeId::new(1))),
11201127
injected_allocator: Once::new(),
11211128
allocator_kind: Once::new(),
1129+
injected_default_alloc_error_hook: Once::new(),
11221130
injected_panic_runtime: Once::new(),
11231131
imported_macro_spans: OneThread::new(RefCell::new(HashMap::new())),
11241132
incr_comp_session: OneThread::new(RefCell::new(IncrCompSession::NotInitialized)),

src/librustc_allocator/lib.rs

+13
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,19 @@ pub static ALLOCATOR_METHODS: &[AllocatorMethod] = &[
4242
},
4343
];
4444

45+
pub static OOM_HANDLING_METHODS: &[AllocatorMethod] = &[
46+
AllocatorMethod {
47+
name: "default_alloc_error_hook",
48+
inputs: &[AllocatorTy::Layout],
49+
output: AllocatorTy::Unit,
50+
},
51+
AllocatorMethod {
52+
name: "abort_internal",
53+
inputs: &[],
54+
output: AllocatorTy::Unit,
55+
},
56+
];
57+
4558
pub struct AllocatorMethod {
4659
pub name: &'static str,
4760
pub inputs: &'static [AllocatorTy],

src/librustc_codegen_llvm/allocator.rs

+48-13
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,9 @@ use std::ptr;
1414
use attributes;
1515
use libc::c_uint;
1616
use rustc::middle::allocator::AllocatorKind;
17+
use rustc::session::InjectedDefaultOomHook;
1718
use rustc::ty::TyCtxt;
18-
use rustc_allocator::{ALLOCATOR_METHODS, AllocatorTy};
19+
use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS, AllocatorTy};
1920

2021
use ModuleLlvm;
2122
use llvm::{self, False, True};
@@ -33,9 +34,10 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
3334
let i8p = llvm::LLVMPointerType(i8, 0);
3435
let void = llvm::LLVMVoidTypeInContext(llcx);
3536

36-
for method in ALLOCATOR_METHODS {
37+
let build = |name: String, inputs: &[AllocatorTy], output: &AllocatorTy,
38+
callee: Option<String>| {
3739
let mut args = Vec::new();
38-
for ty in method.inputs.iter() {
40+
for ty in inputs.iter() {
3941
match *ty {
4042
AllocatorTy::Layout => {
4143
args.push(usize); // size
@@ -48,7 +50,7 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
4850
AllocatorTy::Unit => panic!("invalid allocator arg"),
4951
}
5052
}
51-
let output = match method.output {
53+
let output = match *output {
5254
AllocatorTy::ResultPtr => Some(i8p),
5355
AllocatorTy::Unit => None,
5456

@@ -60,29 +62,40 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
6062
args.as_ptr(),
6163
args.len() as c_uint,
6264
False);
63-
let name = CString::new(format!("__rust_{}", method.name)).unwrap();
65+
let name = CString::new(name).unwrap();
6466
let llfn = llvm::LLVMRustGetOrInsertFunction(llmod,
6567
name.as_ptr(),
6668
ty);
6769

6870
if tcx.sess.target.target.options.default_hidden_visibility {
6971
llvm::LLVMRustSetVisibility(llfn, llvm::Visibility::Hidden);
7072
}
71-
if tcx.sess.target.target.options.requires_uwtable {
72-
attributes::emit_uwtable(llfn, true);
73-
}
74-
75-
let callee = CString::new(kind.fn_name(method.name)).unwrap();
76-
let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
77-
callee.as_ptr(),
78-
ty);
73+
if tcx.sess.target.target.options.requires_uwtable {
74+
attributes::emit_uwtable(llfn, true);
75+
}
7976

8077
let llbb = llvm::LLVMAppendBasicBlockInContext(llcx,
8178
llfn,
8279
"entry\0".as_ptr() as *const _);
8380

8481
let llbuilder = llvm::LLVMCreateBuilderInContext(llcx);
8582
llvm::LLVMPositionBuilderAtEnd(llbuilder, llbb);
83+
84+
let callee = if let Some(callee) = callee {
85+
callee
86+
} else {
87+
// Generate a no-op function
88+
llvm::LLVMBuildRetVoid(llbuilder);
89+
llvm::LLVMDisposeBuilder(llbuilder);
90+
return
91+
};
92+
93+
// Forward the call to another function
94+
let callee = CString::new(callee).unwrap();
95+
let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
96+
callee.as_ptr(),
97+
ty);
98+
8699
let args = args.iter().enumerate().map(|(i, _)| {
87100
llvm::LLVMGetParam(llfn, i as c_uint)
88101
}).collect::<Vec<_>>();
@@ -98,6 +111,28 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
98111
} else {
99112
llvm::LLVMBuildRetVoid(llbuilder);
100113
}
114+
101115
llvm::LLVMDisposeBuilder(llbuilder);
116+
};
117+
118+
for method in ALLOCATOR_METHODS {
119+
let name = format!("__rust_{}", method.name);
120+
build(name, method.inputs, &method.output, Some(kind.fn_name(method.name)))
121+
}
122+
123+
let has_plaftom_functions = match tcx.sess.injected_default_alloc_error_hook.get() {
124+
InjectedDefaultOomHook::None => return,
125+
InjectedDefaultOomHook::Noop => false,
126+
InjectedDefaultOomHook::Platform => true,
127+
};
128+
129+
for method in OOM_HANDLING_METHODS {
130+
let callee = if has_plaftom_functions {
131+
Some(format!("__rust_{}", method.name))
132+
} else {
133+
None
134+
};
135+
let name = format!("__rust_maybe_{}", method.name);
136+
build(name, method.inputs, &method.output, callee)
102137
}
103138
}

src/librustc_codegen_llvm/back/symbol_export.rs

+15-2
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,12 @@ use rustc::hir::CodegenFnAttrFlags;
1717
use rustc::hir::def_id::{CrateNum, DefId, LOCAL_CRATE, CRATE_DEF_INDEX};
1818
use rustc::ich::Fingerprint;
1919
use rustc::middle::exported_symbols::{SymbolExportLevel, ExportedSymbol, metadata_symbol_name};
20-
use rustc::session::config;
20+
use rustc::session::{config, InjectedDefaultOomHook};
2121
use rustc::ty::{TyCtxt, SymbolName};
2222
use rustc::ty::query::Providers;
2323
use rustc::ty::subst::Substs;
2424
use rustc::util::nodemap::{FxHashMap, DefIdMap};
25-
use rustc_allocator::ALLOCATOR_METHODS;
25+
use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS};
2626
use rustc_data_structures::indexed_vec::IndexVec;
2727
use std::collections::hash_map::Entry::*;
2828

@@ -222,6 +222,19 @@ fn exported_symbols_provider_local<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
222222
}
223223
}
224224

225+
match tcx.sess.injected_default_alloc_error_hook.get() {
226+
InjectedDefaultOomHook::None => {}
227+
InjectedDefaultOomHook::Noop |
228+
InjectedDefaultOomHook::Platform => {
229+
for method in OOM_HANDLING_METHODS {
230+
let symbol_name = format!("__rust_{}", method.name);
231+
let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(&symbol_name));
232+
233+
symbols.push((exported_symbol, SymbolExportLevel::Rust));
234+
}
235+
}
236+
}
237+
225238
if tcx.sess.opts.debugging_opts.pgo_gen.is_some() {
226239
// These are weak symbols that point to the profile version and the
227240
// profile name, which need to be treated as exported so LTO doesn't nix

src/librustc_metadata/creader.rs

+15-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use rustc::hir::svh::Svh;
2020
use rustc::middle::allocator::AllocatorKind;
2121
use rustc::middle::cstore::DepKind;
2222
use rustc::mir::interpret::AllocDecodingState;
23-
use rustc::session::{Session, CrateDisambiguator};
23+
use rustc::session::{Session, CrateDisambiguator, InjectedDefaultOomHook};
2424
use rustc::session::config::{Sanitizer, self};
2525
use rustc_target::spec::{PanicStrategy, TargetTriple};
2626
use rustc::session::search_paths::PathKind;
@@ -841,6 +841,7 @@ impl<'a> CrateLoader<'a> {
841841
if !needs_allocator {
842842
self.sess.injected_allocator.set(None);
843843
self.sess.allocator_kind.set(None);
844+
self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
844845
return
845846
}
846847

@@ -862,9 +863,22 @@ impl<'a> CrateLoader<'a> {
862863
if !need_lib_alloc && !need_exe_alloc {
863864
self.sess.injected_allocator.set(None);
864865
self.sess.allocator_kind.set(None);
866+
self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
865867
return
866868
}
867869

870+
let mut has_default_lib_allocator = attr::contains_name(&krate.attrs,
871+
"default_lib_allocator");;
872+
self.cstore.iter_crate_data(|_, data| {
873+
has_default_lib_allocator = has_default_lib_allocator ||
874+
data.root.has_default_lib_allocator;
875+
});
876+
self.sess.injected_default_alloc_error_hook.set(if has_default_lib_allocator {
877+
InjectedDefaultOomHook::Platform
878+
} else {
879+
InjectedDefaultOomHook::Noop
880+
});
881+
868882
// Ok, we need an allocator. Not only that but we're actually going to
869883
// create an artifact that needs one linked in. Let's go find the one
870884
// that we're going to link in.

0 commit comments

Comments
 (0)