src/librustc_allocator/expand.rs

   1 use rustc::middle::allocator::AllocatorKind;
   2 use rustc_errors;
   3 use smallvec::SmallVec;
   4 use syntax::{
   5     ast::{
   6         self, Arg, Attribute, Crate, Expr, FnHeader, Generics, Ident, Item, ItemKind,
   7         Mac, Mod, Mutability, Ty, TyKind, Unsafety, VisibilityKind,
   8     },
   9     attr,
  10     source_map::{
  11         respan, ExpnInfo, MacroAttribute,
  12     },
  13     ext::{
  14         base::{ExtCtxt, Resolver},
  15         build::AstBuilder,
  16         expand::ExpansionConfig,
  17         hygiene::{self, Mark, SyntaxContext},
  18     },
  19     fold::{self, Folder},
  20     parse::ParseSess,
  21     ptr::P,
  22     symbol::Symbol
  23 };
  24 use syntax_pos::Span;
  25
  26 use {AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
  27
  28 pub fn modify(
  29     sess: &ParseSess,
  30     resolver: &mut dyn Resolver,
  31     krate: Crate,
  32     crate_name: String,
  33     handler: &rustc_errors::Handler,
  34 ) -> ast::Crate {
  35     ExpandAllocatorDirectives {
  36         handler,
  37         sess,
  38         resolver,
  39         found: false,
  40         crate_name: Some(crate_name),
  41         in_submod: -1, // -1 to account for the "root" module
  42     }.fold_crate(krate)
  43 }
  44
  45 struct ExpandAllocatorDirectives<'a> {
  46     found: bool,
  47     handler: &'a rustc_errors::Handler,
  48     sess: &'a ParseSess,
  49     resolver: &'a mut dyn Resolver,
  50     crate_name: Option<String>,
  51
  52     // For now, we disallow `global_allocator` in submodules because hygiene is hard. Keep track of
  53     // whether we are in a submodule or not. If `in_submod > 0` we are in a submodule.
  54     in_submod: isize,
  55 }
  56
  57 impl<'a> Folder for ExpandAllocatorDirectives<'a> {
  58     fn fold_item(&mut self, item: P<Item>) -> SmallVec<[P<Item>; 1]> {
  59         debug!("in submodule {}", self.in_submod);
  60
  61         let name = if attr::contains_name(&item.attrs, "global_allocator") {
  62             "global_allocator"
  63         } else {
  64             return fold::noop_fold_item(item, self);
  65         };
  66         match item.node {
  67             ItemKind::Static(..) => {}
  68             _ => {
  69                 self.handler
  70                     .span_err(item.span, "allocators must be statics");
  71                 return smallvec![item];
  72             }
  73         }
  74
  75         if self.in_submod > 0 {
  76             self.handler
  77                 .span_err(item.span, "`global_allocator` cannot be used in submodules");
  78             return smallvec![item];
  79         }
  80
  81         if self.found {
  82             self.handler
  83                 .span_err(item.span, "cannot define more than one #[global_allocator]");
  84             return smallvec![item];
  85         }
  86         self.found = true;
  87
  88         // Create a fresh Mark for the new macro expansion we are about to do
  89         let mark = Mark::fresh(Mark::root());
  90         mark.set_expn_info(ExpnInfo {
  91             call_site: item.span, // use the call site of the static
  92             def_site: None,
  93             format: MacroAttribute(Symbol::intern(name)),
  94             allow_internal_unstable: true,
  95             allow_internal_unsafe: false,
  96             local_inner_macros: false,
  97             edition: hygiene::default_edition(),
  98         });
  99
 100         // Tie the span to the macro expansion info we just created
 101         let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(mark));
 102
 103         // Create an expansion config
 104         let ecfg = ExpansionConfig::default(self.crate_name.take().unwrap());
 105
 106         // Generate a bunch of new items using the AllocFnFactory
 107         let mut f = AllocFnFactory {
 108             span,
 109             kind: AllocatorKind::Global,
 110             global: item.ident,
 111             core: Ident::from_str("core"),
 112             cx: ExtCtxt::new(self.sess, ecfg, self.resolver),
 113         };
 114
 115         // We will generate a new submodule. To `use` the static from that module, we need to get
 116         // the `super::...` path.
 117         let super_path = f.cx.path(f.span, vec![Ident::from_str("super"), f.global]);
 118
 119         // Generate the items in the submodule
 120         let mut items = vec![
 121             // import `core` to use allocators
 122             f.cx.item_extern_crate(f.span, f.core),
 123             // `use` the `global_allocator` in `super`
 124             f.cx.item_use_simple(
 125                 f.span,
 126                 respan(f.span.shrink_to_lo(), VisibilityKind::Inherited),
 127                 super_path,
 128             ),
 129         ];
 130
 131         // Add the allocator methods to the submodule
 132         items.extend(
 133             ALLOCATOR_METHODS
 134                 .iter()
 135                 .map(|method| f.allocator_fn(method)),
 136         );
 137
 138         // Generate the submodule itself
 139         let name = f.kind.fn_name("allocator_abi");
 140         let allocator_abi = Ident::with_empty_ctxt(Symbol::gensym(&name));
 141         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 142         let module = f.cx.monotonic_expander().fold_item(module).pop().unwrap();
 143
 144         // Return the item and new submodule
 145         smallvec![item, module]
 146     }
 147
 148     // If we enter a submodule, take note.
 149     fn fold_mod(&mut self, m: Mod) -> Mod {
 150         debug!("enter submodule");
 151         self.in_submod += 1;
 152         let ret = fold::noop_fold_mod(m, self);
 153         self.in_submod -= 1;
 154         debug!("exit submodule");
 155         ret
 156     }
 157
 158     // `fold_mac` is disabled by default. Enable it here.
 159     fn fold_mac(&mut self, mac: Mac) -> Mac {
 160         fold::noop_fold_mac(mac, self)
 161     }
 162 }
 163
 164 struct AllocFnFactory<'a> {
 165     span: Span,
 166     kind: AllocatorKind,
 167     global: Ident,
 168     core: Ident,
 169     cx: ExtCtxt<'a>,
 170 }
 171
 172 impl<'a> AllocFnFactory<'a> {
 173     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 174         let mut abi_args = Vec::new();
 175         let mut i = 0;
 176         let ref mut mk = || {
 177             let name = Ident::from_str(&format!("arg{}", i));
 178             i += 1;
 179             name
 180         };
 181         let args = method
 182             .inputs
 183             .iter()
 184             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
 185             .collect();
 186         let result = self.call_allocator(method.name, args);
 187         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
 188         let kind = ItemKind::Fn(
 189             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 190             FnHeader {
 191                 unsafety: Unsafety::Unsafe,
 192                 ..FnHeader::default()
 193             },
 194             Generics::default(),
 195             self.cx.block_expr(output_expr),
 196         );
 197         self.cx.item(
 198             self.span,
 199             Ident::from_str(&self.kind.fn_name(method.name)),
 200             self.attrs(),
 201             kind,
 202         )
 203     }
 204
 205     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 206         let method = self.cx.path(
 207             self.span,
 208             vec![
 209                 self.core,
 210                 Ident::from_str("alloc"),
 211                 Ident::from_str("GlobalAlloc"),
 212                 Ident::from_str(method),
 213             ],
 214         );
 215         let method = self.cx.expr_path(method);
 216         let allocator = self.cx.path_ident(self.span, self.global);
 217         let allocator = self.cx.expr_path(allocator);
 218         let allocator = self.cx.expr_addr_of(self.span, allocator);
 219         args.insert(0, allocator);
 220
 221         self.cx.expr_call(self.span, method, args)
 222     }
 223
 224     fn attrs(&self) -> Vec<Attribute> {
 225         let special = Symbol::intern("rustc_std_internal_symbol");
 226         let special = self.cx.meta_word(self.span, special);
 227         vec![self.cx.attribute(self.span, special)]
 228     }
 229
 230     fn arg_ty(
 231         &self,
 232         ty: &AllocatorTy,
 233         args: &mut Vec<Arg>,
 234         ident: &mut dyn FnMut() -> Ident,
 235     ) -> P<Expr> {
 236         match *ty {
 237             AllocatorTy::Layout => {
 238                 let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 239                 let ty_usize = self.cx.ty_path(usize);
 240                 let size = ident();
 241                 let align = ident();
 242                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 243                 args.push(self.cx.arg(self.span, align, ty_usize));
 244
 245                 let layout_new = self.cx.path(
 246                     self.span,
 247                     vec![
 248                         self.core,
 249                         Ident::from_str("alloc"),
 250                         Ident::from_str("Layout"),
 251                         Ident::from_str("from_size_align_unchecked"),
 252                     ],
 253                 );
 254                 let layout_new = self.cx.expr_path(layout_new);
 255                 let size = self.cx.expr_ident(self.span, size);
 256                 let align = self.cx.expr_ident(self.span, align);
 257                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 258                 layout
 259             }
 260
 261             AllocatorTy::Ptr => {
 262                 let ident = ident();
 263                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 264                 let arg = self.cx.expr_ident(self.span, ident);
 265                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 266             }
 267
 268             AllocatorTy::Usize => {
 269                 let ident = ident();
 270                 args.push(self.cx.arg(self.span, ident, self.usize()));
 271                 self.cx.expr_ident(self.span, ident)
 272             }
 273
 274             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 275                 panic!("can't convert AllocatorTy to an argument")
 276             }
 277         }
 278     }
 279
 280     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 281         match *ty {
 282             AllocatorTy::ResultPtr => {
 283                 // We're creating:
 284                 //
 285                 //      #expr as *mut u8
 286
 287                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 288                 (self.ptr_u8(), expr)
 289             }
 290
 291             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 292
 293             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 294                 panic!("can't convert AllocatorTy to an output")
 295             }
 296         }
 297     }
 298
 299     fn usize(&self) -> P<Ty> {
 300         let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 301         self.cx.ty_path(usize)
 302     }
 303
 304     fn ptr_u8(&self) -> P<Ty> {
 305         let u8 = self.cx.path_ident(self.span, Ident::from_str("u8"));
 306         let ty_u8 = self.cx.ty_path(u8);
 307         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 308     }
 309 }