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::{self, Mark, SyntaxContext},
  18     },
  19     mut_visit::{self, MutVisitor},
  20     parse::ParseSess,
  21     ptr::P,
  22     symbol::{keywords, Symbol, sym}
  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: hygiene::default_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 =
 120             f.cx.path(f.span, vec![Ident::with_empty_ctxt(keywords::Super.name()), f.global]);
 121
 122         // Generate the items in the submodule
 123         let mut items = vec![
 124             // import `core` to use allocators
 125             f.cx.item_extern_crate(f.span, f.core),
 126             // `use` the `global_allocator` in `super`
 127             f.cx.item_use_simple(
 128                 f.span,
 129                 respan(f.span.shrink_to_lo(), VisibilityKind::Inherited),
 130                 super_path,
 131             ),
 132         ];
 133
 134         // Add the allocator methods to the submodule
 135         items.extend(
 136             ALLOCATOR_METHODS
 137                 .iter()
 138                 .map(|method| f.allocator_fn(method)),
 139         );
 140
 141         // Generate the submodule itself
 142         let name = f.kind.fn_name("allocator_abi");
 143         let allocator_abi = Ident::with_empty_ctxt(Symbol::gensym(&name));
 144         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 145         let module = f.cx.monotonic_expander().flat_map_item(module).pop().unwrap();
 146
 147         // Return the item and new submodule
 148         smallvec![item, module]
 149     }
 150
 151     // If we enter a submodule, take note.
 152     fn visit_mod(&mut self, m: &mut Mod) {
 153         debug!("enter submodule");
 154         self.in_submod += 1;
 155         mut_visit::noop_visit_mod(m, self);
 156         self.in_submod -= 1;
 157         debug!("exit submodule");
 158     }
 159
 160     // `visit_mac` is disabled by default. Enable it here.
 161     fn visit_mac(&mut self, mac: &mut Mac) {
 162         mut_visit::noop_visit_mac(mac, self)
 163     }
 164 }
 165
 166 struct AllocFnFactory<'a> {
 167     span: Span,
 168     kind: AllocatorKind,
 169     global: Ident,
 170     core: Ident,
 171     cx: ExtCtxt<'a>,
 172 }
 173
 174 impl AllocFnFactory<'_> {
 175     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 176         let mut abi_args = Vec::new();
 177         let mut i = 0;
 178         let ref mut mk = || {
 179             let name = Ident::from_str(&format!("arg{}", i));
 180             i += 1;
 181             name
 182         };
 183         let args = method
 184             .inputs
 185             .iter()
 186             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
 187             .collect();
 188         let result = self.call_allocator(method.name, args);
 189         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
 190         let kind = ItemKind::Fn(
 191             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 192             FnHeader {
 193                 unsafety: Unsafety::Unsafe,
 194                 ..FnHeader::default()
 195             },
 196             Generics::default(),
 197             self.cx.block_expr(output_expr),
 198         );
 199         self.cx.item(
 200             self.span,
 201             Ident::from_str(&self.kind.fn_name(method.name)),
 202             self.attrs(),
 203             kind,
 204         )
 205     }
 206
 207     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 208         let method = self.cx.path(
 209             self.span,
 210             vec![
 211                 self.core,
 212                 Ident::from_str("alloc"),
 213                 Ident::from_str("GlobalAlloc"),
 214                 Ident::from_str(method),
 215             ],
 216         );
 217         let method = self.cx.expr_path(method);
 218         let allocator = self.cx.path_ident(self.span, self.global);
 219         let allocator = self.cx.expr_path(allocator);
 220         let allocator = self.cx.expr_addr_of(self.span, allocator);
 221         args.insert(0, allocator);
 222
 223         self.cx.expr_call(self.span, method, args)
 224     }
 225
 226     fn attrs(&self) -> Vec<Attribute> {
 227         let special = Symbol::intern("rustc_std_internal_symbol");
 228         let special = self.cx.meta_word(self.span, special);
 229         vec![self.cx.attribute(self.span, special)]
 230     }
 231
 232     fn arg_ty(
 233         &self,
 234         ty: &AllocatorTy,
 235         args: &mut Vec<Arg>,
 236         ident: &mut dyn FnMut() -> Ident,
 237     ) -> P<Expr> {
 238         match *ty {
 239             AllocatorTy::Layout => {
 240                 let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 241                 let ty_usize = self.cx.ty_path(usize);
 242                 let size = ident();
 243                 let align = ident();
 244                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 245                 args.push(self.cx.arg(self.span, align, ty_usize));
 246
 247                 let layout_new = self.cx.path(
 248                     self.span,
 249                     vec![
 250                         self.core,
 251                         Ident::from_str("alloc"),
 252                         Ident::from_str("Layout"),
 253                         Ident::from_str("from_size_align_unchecked"),
 254                     ],
 255                 );
 256                 let layout_new = self.cx.expr_path(layout_new);
 257                 let size = self.cx.expr_ident(self.span, size);
 258                 let align = self.cx.expr_ident(self.span, align);
 259                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 260                 layout
 261             }
 262
 263             AllocatorTy::Ptr => {
 264                 let ident = ident();
 265                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 266                 let arg = self.cx.expr_ident(self.span, ident);
 267                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 268             }
 269
 270             AllocatorTy::Usize => {
 271                 let ident = ident();
 272                 args.push(self.cx.arg(self.span, ident, self.usize()));
 273                 self.cx.expr_ident(self.span, ident)
 274             }
 275
 276             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 277                 panic!("can't convert AllocatorTy to an argument")
 278             }
 279         }
 280     }
 281
 282     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 283         match *ty {
 284             AllocatorTy::ResultPtr => {
 285                 // We're creating:
 286                 //
 287                 //      #expr as *mut u8
 288
 289                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 290                 (self.ptr_u8(), expr)
 291             }
 292
 293             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 294
 295             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 296                 panic!("can't convert AllocatorTy to an output")
 297             }
 298         }
 299     }
 300
 301     fn usize(&self) -> P<Ty> {
 302         let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 303         self.cx.ty_path(usize)
 304     }
 305
 306     fn ptr_u8(&self) -> P<Ty> {
 307         let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
 308         let ty_u8 = self.cx.ty_path(u8);
 309         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 310     }
 311 }