src/libsyntax_ext/global_allocator.rs

   1 use syntax::ast::{ItemKind, Mutability, Stmt, Ty, TyKind, Unsafety};
   2 use syntax::ast::{self, Arg, Attribute, Expr, FnHeader, Generics, Ident};
   3 use syntax::attr::check_builtin_macro_attribute;
   4 use syntax::ext::allocator::{AllocatorKind, AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
   5 use syntax::ext::base::{Annotatable, ExtCtxt};
   6 use syntax::ext::build::AstBuilder;
   7 use syntax::ext::hygiene::SyntaxContext;
   8 use syntax::ptr::P;
   9 use syntax::symbol::{kw, sym, Symbol};
  10 use syntax_pos::Span;
  11
  12 pub fn expand(
  13     ecx: &mut ExtCtxt<'_>,
  14     _span: Span,
  15     meta_item: &ast::MetaItem,
  16     item: Annotatable,
  17 ) -> Vec<Annotatable> {
  18     check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);
  19
  20     let not_static = |item: Annotatable| {
  21         ecx.parse_sess.span_diagnostic.span_err(item.span(), "allocators must be statics");
  22         vec![item]
  23     };
  24     let item = match item {
  25         Annotatable::Item(item) => match item.node {
  26             ItemKind::Static(..) => item,
  27             _ => return not_static(Annotatable::Item(item)),
  28         }
  29         _ => return not_static(item),
  30     };
  31
  32     // Generate a bunch of new items using the AllocFnFactory
  33     let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(ecx.current_expansion.id));
  34     let f = AllocFnFactory {
  35         span,
  36         kind: AllocatorKind::Global,
  37         global: item.ident,
  38         cx: ecx,
  39     };
  40
  41     // Generate item statements for the allocator methods.
  42     let stmts = ALLOCATOR_METHODS.iter().map(|method| f.allocator_fn(method)).collect();
  43
  44     // Generate anonymous constant serving as container for the allocator methods.
  45     let const_ty = ecx.ty(span, TyKind::Tup(Vec::new()));
  46     let const_body = ecx.expr_block(ecx.block(span, stmts));
  47     let const_item =
  48         ecx.item_const(span, Ident::with_empty_ctxt(kw::Underscore), const_ty, const_body);
  49
  50     // Return the original item and the new methods.
  51     vec![Annotatable::Item(item), Annotatable::Item(const_item)]
  52 }
  53
  54 struct AllocFnFactory<'a, 'b> {
  55     span: Span,
  56     kind: AllocatorKind,
  57     global: Ident,
  58     cx: &'b ExtCtxt<'a>,
  59 }
  60
  61 impl AllocFnFactory<'_, '_> {
  62     fn allocator_fn(&self, method: &AllocatorMethod) -> Stmt {
  63         let mut abi_args = Vec::new();
  64         let mut i = 0;
  65         let ref mut mk = || {
  66             let name = Ident::from_str(&format!("arg{}", i));
  67             i += 1;
  68             name
  69         };
  70         let args = method
  71             .inputs
  72             .iter()
  73             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
  74             .collect();
  75         let result = self.call_allocator(method.name, args);
  76         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
  77         let kind = ItemKind::Fn(
  78             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
  79             FnHeader {
  80                 unsafety: Unsafety::Unsafe,
  81                 ..FnHeader::default()
  82             },
  83             Generics::default(),
  84             self.cx.block_expr(output_expr),
  85         );
  86         let item = self.cx.item(
  87             self.span,
  88             Ident::from_str(&self.kind.fn_name(method.name)),
  89             self.attrs(),
  90             kind,
  91         );
  92         self.cx.stmt_item(self.span, item)
  93     }
  94
  95     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
  96         let method = self.cx.std_path(&[
  97             Symbol::intern("alloc"),
  98             Symbol::intern("GlobalAlloc"),
  99             Symbol::intern(method),
 100         ]);
 101         let method = self.cx.expr_path(self.cx.path(self.span, method));
 102         let allocator = self.cx.path_ident(self.span, self.global);
 103         let allocator = self.cx.expr_path(allocator);
 104         let allocator = self.cx.expr_addr_of(self.span, allocator);
 105         args.insert(0, allocator);
 106
 107         self.cx.expr_call(self.span, method, args)
 108     }
 109
 110     fn attrs(&self) -> Vec<Attribute> {
 111         let special = sym::rustc_std_internal_symbol;
 112         let special = self.cx.meta_word(self.span, special);
 113         vec![self.cx.attribute(special)]
 114     }
 115
 116     fn arg_ty(
 117         &self,
 118         ty: &AllocatorTy,
 119         args: &mut Vec<Arg>,
 120         ident: &mut dyn FnMut() -> Ident,
 121     ) -> P<Expr> {
 122         match *ty {
 123             AllocatorTy::Layout => {
 124                 let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 125                 let ty_usize = self.cx.ty_path(usize);
 126                 let size = ident();
 127                 let align = ident();
 128                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 129                 args.push(self.cx.arg(self.span, align, ty_usize));
 130
 131                 let layout_new = self.cx.std_path(&[
 132                     Symbol::intern("alloc"),
 133                     Symbol::intern("Layout"),
 134                     Symbol::intern("from_size_align_unchecked"),
 135                 ]);
 136                 let layout_new = self.cx.expr_path(self.cx.path(self.span, layout_new));
 137                 let size = self.cx.expr_ident(self.span, size);
 138                 let align = self.cx.expr_ident(self.span, align);
 139                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 140                 layout
 141             }
 142
 143             AllocatorTy::Ptr => {
 144                 let ident = ident();
 145                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 146                 let arg = self.cx.expr_ident(self.span, ident);
 147                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 148             }
 149
 150             AllocatorTy::Usize => {
 151                 let ident = ident();
 152                 args.push(self.cx.arg(self.span, ident, self.usize()));
 153                 self.cx.expr_ident(self.span, ident)
 154             }
 155
 156             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 157                 panic!("can't convert AllocatorTy to an argument")
 158             }
 159         }
 160     }
 161
 162     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 163         match *ty {
 164             AllocatorTy::ResultPtr => {
 165                 // We're creating:
 166                 //
 167                 //      #expr as *mut u8
 168
 169                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 170                 (self.ptr_u8(), expr)
 171             }
 172
 173             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 174
 175             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 176                 panic!("can't convert `AllocatorTy` to an output")
 177             }
 178         }
 179     }
 180
 181     fn usize(&self) -> P<Ty> {
 182         let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 183         self.cx.ty_path(usize)
 184     }
 185
 186     fn ptr_u8(&self) -> P<Ty> {
 187         let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
 188         let ty_u8 = self.cx.ty_path(u8);
 189         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 190     }
 191 }