src/allocator.rs

   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 use rustc_span::symbol::sym;
   9
  10 /// Returns whether an allocator shim was created
  11 pub(crate) fn codegen(
  12     tcx: TyCtxt<'_>,
  13     module: &mut impl Module,
  14     unwind_context: &mut UnwindContext,
  15 ) -> bool {
  16     let any_dynamic_crate = tcx.dependency_formats(LOCAL_CRATE).iter().any(|(_, list)| {
  17         use rustc_middle::middle::dependency_format::Linkage;
  18         list.iter().any(|&linkage| linkage == Linkage::Dynamic)
  19     });
  20     if any_dynamic_crate {
  21         false
  22     } else if let Some(kind) = tcx.allocator_kind() {
  23         codegen_inner(module, unwind_context, kind);
  24         true
  25     } else {
  26         false
  27     }
  28 }
  29
  30 fn codegen_inner(
  31     module: &mut impl Module,
  32     unwind_context: &mut UnwindContext,
  33     kind: AllocatorKind,
  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         //eprintln!("Codegen allocator shim {} -> {} ({:?} -> {:?})", caller_name, callee_name, sig.params, sig.returns);
  69
  70         let func_id = module.declare_function(&caller_name, Linkage::Export, &sig).unwrap();
  71
  72         let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
  73
  74         let mut ctx = Context::new();
  75         ctx.func = Function::with_name_signature(ExternalName::user(0, 0), sig.clone());
  76         {
  77             let mut func_ctx = FunctionBuilderContext::new();
  78             let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
  79
  80             let block = bcx.create_block();
  81             bcx.switch_to_block(block);
  82             let args = arg_tys
  83                 .into_iter()
  84                 .map(|ty| bcx.append_block_param(block, ty))
  85                 .collect::<Vec<Value>>();
  86
  87             let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
  88             let call_inst = bcx.ins().call(callee_func_ref, &args);
  89             let results = bcx.inst_results(call_inst).to_vec(); // Clone to prevent borrow error
  90
  91             bcx.ins().return_(&results);
  92             bcx.seal_all_blocks();
  93             bcx.finalize();
  94         }
  95         module
  96             .define_function(func_id, &mut ctx, &mut NullTrapSink {}, &mut NullStackMapSink {})
  97             .unwrap();
  98         unwind_context.add_function(func_id, &ctx, module.isa());
  99     }
 100
 101     let sig = Signature {
 102         call_conv: CallConv::triple_default(module.isa().triple()),
 103         params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
 104         returns: vec![],
 105     };
 106
 107     let callee_name = kind.fn_name(sym::oom);
 108     //eprintln!("Codegen allocator shim {} -> {} ({:?} -> {:?})", caller_name, callee_name, sig.params, sig.returns);
 109
 110     let func_id =
 111         module.declare_function("__rust_alloc_error_handler", Linkage::Export, &sig).unwrap();
 112
 113     let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
 114
 115     let mut ctx = Context::new();
 116     ctx.func = Function::with_name_signature(ExternalName::user(0, 0), sig);
 117     {
 118         let mut func_ctx = FunctionBuilderContext::new();
 119         let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
 120
 121         let block = bcx.create_block();
 122         bcx.switch_to_block(block);
 123         let args = (&[usize_ty, usize_ty])
 124             .iter()
 125             .map(|&ty| bcx.append_block_param(block, ty))
 126             .collect::<Vec<Value>>();
 127
 128         let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
 129         bcx.ins().call(callee_func_ref, &args);
 130
 131         bcx.ins().trap(TrapCode::UnreachableCodeReached);
 132         bcx.seal_all_blocks();
 133         bcx.finalize();
 134     }
 135     module
 136         .define_function(func_id, &mut ctx, &mut NullTrapSink {}, &mut NullStackMapSink {})
 137         .unwrap();
 138     unwind_context.add_function(func_id, &ctx, module.isa());
 139 }