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![sym::rustc_attrs].into()),
  95             allow_internal_unsafe: false,
  96             local_inner_macros: false,
  97             edition: self.sess.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::with_empty_ctxt(sym::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::with_empty_ctxt(kw::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::from_str(&name).gensym();
 141         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 142         let module = f.cx.monotonic_expander().flat_map_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 visit_mod(&mut self, m: &mut Mod) {
 150         debug!("enter submodule");
 151         self.in_submod += 1;
 152         mut_visit::noop_visit_mod(m, self);
 153         self.in_submod -= 1;
 154         debug!("exit submodule");
 155     }
 156
 157     // `visit_mac` is disabled by default. Enable it here.
 158     fn visit_mac(&mut self, mac: &mut Mac) {
 159         mut_visit::noop_visit_mac(mac, self)
 160     }
 161 }
 162
 163 struct AllocFnFactory<'a> {
 164     span: Span,
 165     kind: AllocatorKind,
 166     global: Ident,
 167     core: Ident,
 168     cx: ExtCtxt<'a>,
 169 }
 170
 171 impl AllocFnFactory<'_> {
 172     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 173         let mut abi_args = Vec::new();
 174         let mut i = 0;
 175         let ref mut mk = || {
 176             let name = Ident::from_str(&format!("arg{}", i));
 177             i += 1;
 178             name
 179         };
 180         let args = method
 181             .inputs
 182             .iter()
 183             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
 184             .collect();
 185         let result = self.call_allocator(method.name, args);
 186         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
 187         let kind = ItemKind::Fn(
 188             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 189             FnHeader {
 190                 unsafety: Unsafety::Unsafe,
 191                 ..FnHeader::default()
 192             },
 193             Generics::default(),
 194             self.cx.block_expr(output_expr),
 195         );
 196         self.cx.item(
 197             self.span,
 198             Ident::from_str(&self.kind.fn_name(method.name)),
 199             self.attrs(),
 200             kind,
 201         )
 202     }
 203
 204     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 205         let method = self.cx.path(
 206             self.span,
 207             vec![
 208                 self.core,
 209                 Ident::from_str("alloc"),
 210                 Ident::from_str("GlobalAlloc"),
 211                 Ident::from_str(method),
 212             ],
 213         );
 214         let method = self.cx.expr_path(method);
 215         let allocator = self.cx.path_ident(self.span, self.global);
 216         let allocator = self.cx.expr_path(allocator);
 217         let allocator = self.cx.expr_addr_of(self.span, allocator);
 218         args.insert(0, allocator);
 219
 220         self.cx.expr_call(self.span, method, args)
 221     }
 222
 223     fn attrs(&self) -> Vec<Attribute> {
 224         let special = sym::rustc_std_internal_symbol;
 225         let special = self.cx.meta_word(self.span, special);
 226         vec![self.cx.attribute(self.span, special)]
 227     }
 228
 229     fn arg_ty(
 230         &self,
 231         ty: &AllocatorTy,
 232         args: &mut Vec<Arg>,
 233         ident: &mut dyn FnMut() -> Ident,
 234     ) -> P<Expr> {
 235         match *ty {
 236             AllocatorTy::Layout => {
 237                 let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 238                 let ty_usize = self.cx.ty_path(usize);
 239                 let size = ident();
 240                 let align = ident();
 241                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 242                 args.push(self.cx.arg(self.span, align, ty_usize));
 243
 244                 let layout_new = self.cx.path(
 245                     self.span,
 246                     vec![
 247                         self.core,
 248                         Ident::from_str("alloc"),
 249                         Ident::from_str("Layout"),
 250                         Ident::from_str("from_size_align_unchecked"),
 251                     ],
 252                 );
 253                 let layout_new = self.cx.expr_path(layout_new);
 254                 let size = self.cx.expr_ident(self.span, size);
 255                 let align = self.cx.expr_ident(self.span, align);
 256                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 257                 layout
 258             }
 259
 260             AllocatorTy::Ptr => {
 261                 let ident = ident();
 262                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 263                 let arg = self.cx.expr_ident(self.span, ident);
 264                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 265             }
 266
 267             AllocatorTy::Usize => {
 268                 let ident = ident();
 269                 args.push(self.cx.arg(self.span, ident, self.usize()));
 270                 self.cx.expr_ident(self.span, ident)
 271             }
 272
 273             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 274                 panic!("can't convert AllocatorTy to an argument")
 275             }
 276         }
 277     }
 278
 279     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 280         match *ty {
 281             AllocatorTy::ResultPtr => {
 282                 // We're creating:
 283                 //
 284                 //      #expr as *mut u8
 285
 286                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 287                 (self.ptr_u8(), expr)
 288             }
 289
 290             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 291
 292             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 293                 panic!("can't convert AllocatorTy to an output")
 294             }
 295         }
 296     }
 297
 298     fn usize(&self) -> P<Ty> {
 299         let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 300         self.cx.ty_path(usize)
 301     }
 302
 303     fn ptr_u8(&self) -> P<Ty> {
 304         let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
 305         let ty_u8 = self.cx.ty_path(u8);
 306         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 307     }
 308 }