src/librustc_allocator/expand.rs

   1 use log::debug;
   2 use rustc::middle::allocator::AllocatorKind;
   3 use smallvec::{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::{Mark, SyntaxContext},
  18     },
  19     mut_visit::{self, MutVisitor},
  20     parse::ParseSess,
  21     ptr::P,
  22     symbol::{kw, sym, Symbol}
  23 };
  24 use syntax_pos::Span;
  25
  26 use crate::{AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
  27
  28 pub fn modify(
  29     sess: &ParseSess,
  30     resolver: &mut dyn Resolver,
  31     krate: &mut Crate,
  32     crate_name: String,
  33     handler: &rustc_errors::Handler,
  34 ) {
  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     }.visit_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 MutVisitor for ExpandAllocatorDirectives<'_> {
  58     fn flat_map_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, sym::global_allocator) {
  62             "global_allocator"
  63         } else {
  64             return mut_visit::noop_flat_map_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: Some(vec![
  95                 Symbol::intern("rustc_attrs"),
  96             ].into()),
  97             allow_internal_unsafe: false,
  98             local_inner_macros: false,
  99             edition: self.sess.edition,
 100         });
 101
 102         // Tie the span to the macro expansion info we just created
 103         let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(mark));
 104
 105         // Create an expansion config
 106         let ecfg = ExpansionConfig::default(self.crate_name.take().unwrap());
 107
 108         // Generate a bunch of new items using the AllocFnFactory
 109         let mut f = AllocFnFactory {
 110             span,
 111             kind: AllocatorKind::Global,
 112             global: item.ident,
 113             core: Ident::with_empty_ctxt(sym::core),
 114             cx: ExtCtxt::new(self.sess, ecfg, self.resolver),
 115         };
 116
 117         // We will generate a new submodule. To `use` the static from that module, we need to get
 118         // the `super::...` path.
 119         let super_path = f.cx.path(f.span, vec![Ident::with_empty_ctxt(kw::Super), f.global]);
 120
 121         // Generate the items in the submodule
 122         let mut items = vec![
 123             // import `core` to use allocators
 124             f.cx.item_extern_crate(f.span, f.core),
 125             // `use` the `global_allocator` in `super`
 126             f.cx.item_use_simple(
 127                 f.span,
 128                 respan(f.span.shrink_to_lo(), VisibilityKind::Inherited),
 129                 super_path,
 130             ),
 131         ];
 132
 133         // Add the allocator methods to the submodule
 134         items.extend(
 135             ALLOCATOR_METHODS
 136                 .iter()
 137                 .map(|method| f.allocator_fn(method)),
 138         );
 139
 140         // Generate the submodule itself
 141         let name = f.kind.fn_name("allocator_abi");
 142         let allocator_abi = Ident::from_str(&name).gensym();
 143         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 144         let module = f.cx.monotonic_expander().flat_map_item(module).pop().unwrap();
 145
 146         // Return the item and new submodule
 147         smallvec![item, module]
 148     }
 149
 150     // If we enter a submodule, take note.
 151     fn visit_mod(&mut self, m: &mut Mod) {
 152         debug!("enter submodule");
 153         self.in_submod += 1;
 154         mut_visit::noop_visit_mod(m, self);
 155         self.in_submod -= 1;
 156         debug!("exit submodule");
 157     }
 158
 159     // `visit_mac` is disabled by default. Enable it here.
 160     fn visit_mac(&mut self, mac: &mut Mac) {
 161         mut_visit::noop_visit_mac(mac, self)
 162     }
 163 }
 164
 165 struct AllocFnFactory<'a> {
 166     span: Span,
 167     kind: AllocatorKind,
 168     global: Ident,
 169     core: Ident,
 170     cx: ExtCtxt<'a>,
 171 }
 172
 173 impl AllocFnFactory<'_> {
 174     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 175         let mut abi_args = Vec::new();
 176         let mut i = 0;
 177         let ref mut mk = || {
 178             let name = Ident::from_str(&format!("arg{}", i));
 179             i += 1;
 180             name
 181         };
 182         let args = method
 183             .inputs
 184             .iter()
 185             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
 186             .collect();
 187         let result = self.call_allocator(method.name, args);
 188         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
 189         let kind = ItemKind::Fn(
 190             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 191             FnHeader {
 192                 unsafety: Unsafety::Unsafe,
 193                 ..FnHeader::default()
 194             },
 195             Generics::default(),
 196             self.cx.block_expr(output_expr),
 197         );
 198         self.cx.item(
 199             self.span,
 200             Ident::from_str(&self.kind.fn_name(method.name)),
 201             self.attrs(),
 202             kind,
 203         )
 204     }
 205
 206     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 207         let method = self.cx.path(
 208             self.span,
 209             vec![
 210                 self.core,
 211                 Ident::from_str("alloc"),
 212                 Ident::from_str("GlobalAlloc"),
 213                 Ident::from_str(method),
 214             ],
 215         );
 216         let method = self.cx.expr_path(method);
 217         let allocator = self.cx.path_ident(self.span, self.global);
 218         let allocator = self.cx.expr_path(allocator);
 219         let allocator = self.cx.expr_addr_of(self.span, allocator);
 220         args.insert(0, allocator);
 221
 222         self.cx.expr_call(self.span, method, args)
 223     }
 224
 225     fn attrs(&self) -> Vec<Attribute> {
 226         let special = Symbol::intern("rustc_std_internal_symbol");
 227         let special = self.cx.meta_word(self.span, special);
 228         vec![self.cx.attribute(self.span, special)]
 229     }
 230
 231     fn arg_ty(
 232         &self,
 233         ty: &AllocatorTy,
 234         args: &mut Vec<Arg>,
 235         ident: &mut dyn FnMut() -> Ident,
 236     ) -> P<Expr> {
 237         match *ty {
 238             AllocatorTy::Layout => {
 239                 let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 240                 let ty_usize = self.cx.ty_path(usize);
 241                 let size = ident();
 242                 let align = ident();
 243                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 244                 args.push(self.cx.arg(self.span, align, ty_usize));
 245
 246                 let layout_new = self.cx.path(
 247                     self.span,
 248                     vec![
 249                         self.core,
 250                         Ident::from_str("alloc"),
 251                         Ident::from_str("Layout"),
 252                         Ident::from_str("from_size_align_unchecked"),
 253                     ],
 254                 );
 255                 let layout_new = self.cx.expr_path(layout_new);
 256                 let size = self.cx.expr_ident(self.span, size);
 257                 let align = self.cx.expr_ident(self.span, align);
 258                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 259                 layout
 260             }
 261
 262             AllocatorTy::Ptr => {
 263                 let ident = ident();
 264                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 265                 let arg = self.cx.expr_ident(self.span, ident);
 266                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 267             }
 268
 269             AllocatorTy::Usize => {
 270                 let ident = ident();
 271                 args.push(self.cx.arg(self.span, ident, self.usize()));
 272                 self.cx.expr_ident(self.span, ident)
 273             }
 274
 275             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 276                 panic!("can't convert AllocatorTy to an argument")
 277             }
 278         }
 279     }
 280
 281     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 282         match *ty {
 283             AllocatorTy::ResultPtr => {
 284                 // We're creating:
 285                 //
 286                 //      #expr as *mut u8
 287
 288                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 289                 (self.ptr_u8(), expr)
 290             }
 291
 292             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 293
 294             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 295                 panic!("can't convert AllocatorTy to an output")
 296             }
 297         }
 298     }
 299
 300     fn usize(&self) -> P<Ty> {
 301         let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 302         self.cx.ty_path(usize)
 303     }
 304
 305     fn ptr_u8(&self) -> P<Ty> {
 306         let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
 307         let ty_u8 = self.cx.ty_path(u8);
 308         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 309     }
 310 }