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, ExpnKind,
  12     },
  13     ext::{
  14         base::{ExtCtxt, MacroKind, Resolver},
  15         build::AstBuilder,
  16         expand::ExpansionConfig,
  17         hygiene::ExpnId,
  18     },
  19     mut_visit::{self, MutVisitor},
  20     parse::ParseSess,
  21     ptr::P,
  22     symbol::{kw, 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         if !attr::contains_name(&item.attrs, sym::global_allocator) {
  62             return mut_visit::noop_flat_map_item(item, self);
  63         }
  64
  65         match item.node {
  66             ItemKind::Static(..) => {}
  67             _ => {
  68                 self.handler
  69                     .span_err(item.span, "allocators must be statics");
  70                 return smallvec![item];
  71             }
  72         }
  73
  74         if self.in_submod > 0 {
  75             self.handler
  76                 .span_err(item.span, "`global_allocator` cannot be used in submodules");
  77             return smallvec![item];
  78         }
  79
  80         if self.found {
  81             self.handler
  82                 .span_err(item.span, "cannot define more than one `#[global_allocator]`");
  83             return smallvec![item];
  84         }
  85         self.found = true;
  86
  87         // Create a new expansion for the generated allocator code.
  88         let span = item.span.fresh_expansion(ExpnId::root(), ExpnInfo::allow_unstable(
  89             ExpnKind::Macro(MacroKind::Attr, sym::global_allocator), item.span, self.sess.edition,
  90             [sym::rustc_attrs][..].into(),
  91         ));
  92
  93         // Create an expansion config
  94         let ecfg = ExpansionConfig::default(self.crate_name.take().unwrap());
  95
  96         // Generate a bunch of new items using the AllocFnFactory
  97         let mut f = AllocFnFactory {
  98             span,
  99             kind: AllocatorKind::Global,
 100             global: item.ident,
 101             core: Ident::with_empty_ctxt(sym::core),
 102             cx: ExtCtxt::new(self.sess, ecfg, self.resolver),
 103         };
 104
 105         // We will generate a new submodule. To `use` the static from that module, we need to get
 106         // the `super::...` path.
 107         let super_path = f.cx.path(f.span, vec![Ident::with_empty_ctxt(kw::Super), f.global]);
 108
 109         // Generate the items in the submodule
 110         let mut items = vec![
 111             // import `core` to use allocators
 112             f.cx.item_extern_crate(f.span, f.core),
 113             // `use` the `global_allocator` in `super`
 114             f.cx.item_use_simple(
 115                 f.span,
 116                 respan(f.span.shrink_to_lo(), VisibilityKind::Inherited),
 117                 super_path,
 118             ),
 119         ];
 120
 121         // Add the allocator methods to the submodule
 122         items.extend(
 123             ALLOCATOR_METHODS
 124                 .iter()
 125                 .map(|method| f.allocator_fn(method)),
 126         );
 127
 128         // Generate the submodule itself
 129         let name = f.kind.fn_name("allocator_abi");
 130         let allocator_abi = Ident::from_str(&name).gensym();
 131         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 132         let module = f.cx.monotonic_expander().flat_map_item(module).pop().unwrap();
 133
 134         // Return the item and new submodule
 135         smallvec![item, module]
 136     }
 137
 138     // If we enter a submodule, take note.
 139     fn visit_mod(&mut self, m: &mut Mod) {
 140         debug!("enter submodule");
 141         self.in_submod += 1;
 142         mut_visit::noop_visit_mod(m, self);
 143         self.in_submod -= 1;
 144         debug!("exit submodule");
 145     }
 146
 147     // `visit_mac` is disabled by default. Enable it here.
 148     fn visit_mac(&mut self, mac: &mut Mac) {
 149         mut_visit::noop_visit_mac(mac, self)
 150     }
 151 }
 152
 153 struct AllocFnFactory<'a> {
 154     span: Span,
 155     kind: AllocatorKind,
 156     global: Ident,
 157     core: Ident,
 158     cx: ExtCtxt<'a>,
 159 }
 160
 161 impl AllocFnFactory<'_> {
 162     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 163         let mut abi_args = Vec::new();
 164         let mut i = 0;
 165         let ref mut mk = || {
 166             let name = Ident::from_str(&format!("arg{}", i));
 167             i += 1;
 168             name
 169         };
 170         let args = method
 171             .inputs
 172             .iter()
 173             .map(|ty| self.arg_ty(ty, &mut abi_args, mk))
 174             .collect();
 175         let result = self.call_allocator(method.name, args);
 176         let (output_ty, output_expr) = self.ret_ty(&method.output, result);
 177         let kind = ItemKind::Fn(
 178             self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 179             FnHeader {
 180                 unsafety: Unsafety::Unsafe,
 181                 ..FnHeader::default()
 182             },
 183             Generics::default(),
 184             self.cx.block_expr(output_expr),
 185         );
 186         self.cx.item(
 187             self.span,
 188             Ident::from_str(&self.kind.fn_name(method.name)),
 189             self.attrs(),
 190             kind,
 191         )
 192     }
 193
 194     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 195         let method = self.cx.path(
 196             self.span,
 197             vec![
 198                 self.core,
 199                 Ident::from_str("alloc"),
 200                 Ident::from_str("GlobalAlloc"),
 201                 Ident::from_str(method),
 202             ],
 203         );
 204         let method = self.cx.expr_path(method);
 205         let allocator = self.cx.path_ident(self.span, self.global);
 206         let allocator = self.cx.expr_path(allocator);
 207         let allocator = self.cx.expr_addr_of(self.span, allocator);
 208         args.insert(0, allocator);
 209
 210         self.cx.expr_call(self.span, method, args)
 211     }
 212
 213     fn attrs(&self) -> Vec<Attribute> {
 214         let special = sym::rustc_std_internal_symbol;
 215         let special = self.cx.meta_word(self.span, special);
 216         vec![self.cx.attribute(self.span, special)]
 217     }
 218
 219     fn arg_ty(
 220         &self,
 221         ty: &AllocatorTy,
 222         args: &mut Vec<Arg>,
 223         ident: &mut dyn FnMut() -> Ident,
 224     ) -> P<Expr> {
 225         match *ty {
 226             AllocatorTy::Layout => {
 227                 let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 228                 let ty_usize = self.cx.ty_path(usize);
 229                 let size = ident();
 230                 let align = ident();
 231                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 232                 args.push(self.cx.arg(self.span, align, ty_usize));
 233
 234                 let layout_new = self.cx.path(
 235                     self.span,
 236                     vec![
 237                         self.core,
 238                         Ident::from_str("alloc"),
 239                         Ident::from_str("Layout"),
 240                         Ident::from_str("from_size_align_unchecked"),
 241                     ],
 242                 );
 243                 let layout_new = self.cx.expr_path(layout_new);
 244                 let size = self.cx.expr_ident(self.span, size);
 245                 let align = self.cx.expr_ident(self.span, align);
 246                 let layout = self.cx.expr_call(self.span, layout_new, vec![size, align]);
 247                 layout
 248             }
 249
 250             AllocatorTy::Ptr => {
 251                 let ident = ident();
 252                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 253                 let arg = self.cx.expr_ident(self.span, ident);
 254                 self.cx.expr_cast(self.span, arg, self.ptr_u8())
 255             }
 256
 257             AllocatorTy::Usize => {
 258                 let ident = ident();
 259                 args.push(self.cx.arg(self.span, ident, self.usize()));
 260                 self.cx.expr_ident(self.span, ident)
 261             }
 262
 263             AllocatorTy::ResultPtr | AllocatorTy::Unit => {
 264                 panic!("can't convert AllocatorTy to an argument")
 265             }
 266         }
 267     }
 268
 269     fn ret_ty(&self, ty: &AllocatorTy, expr: P<Expr>) -> (P<Ty>, P<Expr>) {
 270         match *ty {
 271             AllocatorTy::ResultPtr => {
 272                 // We're creating:
 273                 //
 274                 //      #expr as *mut u8
 275
 276                 let expr = self.cx.expr_cast(self.span, expr, self.ptr_u8());
 277                 (self.ptr_u8(), expr)
 278             }
 279
 280             AllocatorTy::Unit => (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr),
 281
 282             AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
 283                 panic!("can't convert `AllocatorTy` to an output")
 284             }
 285         }
 286     }
 287
 288     fn usize(&self) -> P<Ty> {
 289         let usize = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::usize));
 290         self.cx.ty_path(usize)
 291     }
 292
 293     fn ptr_u8(&self) -> P<Ty> {
 294         let u8 = self.cx.path_ident(self.span, Ident::with_empty_ctxt(sym::u8));
 295         let ty_u8 = self.cx.ty_path(u8);
 296         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 297     }
 298 }