1 //! Allocator shim
2 // Adapted from rustc
3
4 use crate::prelude::*;
5
6 use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
7 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
8
9 /// Returns whether an allocator shim was created
codegen( tcx: TyCtxt<'_>, module: &mut impl Module, unwind_context: &mut UnwindContext, ) -> bool10 pub(crate) fn codegen(
11 tcx: TyCtxt<'_>,
12 module: &mut impl Module,
13 unwind_context: &mut UnwindContext,
14 ) -> bool {
15 let any_dynamic_crate = tcx.dependency_formats(()).iter().any(|(_, list)| {
16 use rustc_middle::middle::dependency_format::Linkage;
17 list.iter().any(|&linkage| linkage == Linkage::Dynamic)
18 });
19 if any_dynamic_crate {
20 false
21 } else if let Some(kind) = tcx.allocator_kind(()) {
22 codegen_inner(module, unwind_context, kind, tcx.lang_items().oom().is_some());
23 true
24 } else {
25 false
26 }
27 }
28
codegen_inner( module: &mut impl Module, unwind_context: &mut UnwindContext, kind: AllocatorKind, has_alloc_error_handler: bool, )29 fn codegen_inner(
30 module: &mut impl Module,
31 unwind_context: &mut UnwindContext,
32 kind: AllocatorKind,
33 has_alloc_error_handler: bool,
34 ) {
35 let usize_ty = module.target_config().pointer_type();
36
37 for method in ALLOCATOR_METHODS {
38 let mut arg_tys = Vec::with_capacity(method.inputs.len());
39 for ty in method.inputs.iter() {
40 match *ty {
41 AllocatorTy::Layout => {
42 arg_tys.push(usize_ty); // size
43 arg_tys.push(usize_ty); // align
44 }
45 AllocatorTy::Ptr => arg_tys.push(usize_ty),
46 AllocatorTy::Usize => arg_tys.push(usize_ty),
47
48 AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
49 }
50 }
51 let output = match method.output {
52 AllocatorTy::ResultPtr => Some(usize_ty),
53 AllocatorTy::Unit => None,
54
55 AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
56 panic!("invalid allocator output")
57 }
58 };
59
60 let sig = Signature {
61 call_conv: CallConv::triple_default(module.isa().triple()),
62 params: arg_tys.iter().cloned().map(AbiParam::new).collect(),
63 returns: output.into_iter().map(AbiParam::new).collect(),
64 };
65
66 let caller_name = format!("__rust_{}", method.name);
67 let callee_name = kind.fn_name(method.name);
68
69 let func_id = module.declare_function(&caller_name, Linkage::Export, &sig).unwrap();
70
71 let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
72
73 let mut ctx = Context::new();
74 ctx.func = Function::with_name_signature(ExternalName::user(0, 0), sig.clone());
75 {
76 let mut func_ctx = FunctionBuilderContext::new();
77 let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
78
79 let block = bcx.create_block();
80 bcx.switch_to_block(block);
81 let args = arg_tys
82 .into_iter()
83 .map(|ty| bcx.append_block_param(block, ty))
84 .collect::<Vec<Value>>();
85
86 let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
87 let call_inst = bcx.ins().call(callee_func_ref, &args);
88 let results = bcx.inst_results(call_inst).to_vec(); // Clone to prevent borrow error
89
90 bcx.ins().return_(&results);
91 bcx.seal_all_blocks();
92 bcx.finalize();
93 }
94 module
95 .define_function(func_id, &mut ctx, &mut NullTrapSink {}, &mut NullStackMapSink {})
96 .unwrap();
97 unwind_context.add_function(func_id, &ctx, module.isa());
98 }
99
100 let sig = Signature {
101 call_conv: CallConv::triple_default(module.isa().triple()),
102 params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
103 returns: vec![],
104 };
105
106 let callee_name = if has_alloc_error_handler { "__rg_oom" } else { "__rdl_oom" };
107
108 let func_id =
109 module.declare_function("__rust_alloc_error_handler", Linkage::Export, &sig).unwrap();
110
111 let callee_func_id = module.declare_function(callee_name, Linkage::Import, &sig).unwrap();
112
113 let mut ctx = Context::new();
114 ctx.func = Function::with_name_signature(ExternalName::user(0, 0), sig);
115 {
116 let mut func_ctx = FunctionBuilderContext::new();
117 let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
118
119 let block = bcx.create_block();
120 bcx.switch_to_block(block);
121 let args = (&[usize_ty, usize_ty])
122 .iter()
123 .map(|&ty| bcx.append_block_param(block, ty))
124 .collect::<Vec<Value>>();
125
126 let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
127 bcx.ins().call(callee_func_ref, &args);
128
129 bcx.ins().trap(TrapCode::UnreachableCodeReached);
130 bcx.seal_all_blocks();
131 bcx.finalize();
132 }
133 module
134 .define_function(func_id, &mut ctx, &mut NullTrapSink {}, &mut NullStackMapSink {})
135 .unwrap();
136 unwind_context.add_function(func_id, &ctx, module.isa());
137 }
138