src/allocator.rs

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