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