src/librustc_allocator/expand.rs

   1 // Copyright 2016 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use rustc::middle::allocator::AllocatorKind;
  12 use rustc_errors;
  13 use syntax::abi::Abi;
  14 use syntax::ast::{Crate, Attribute, LitKind, StrStyle, ExprKind};
  15 use syntax::ast::{Unsafety, Constness, Generics, Mutability, Ty, Mac, Arg};
  16 use syntax::ast::{self, Ident, Item, ItemKind, TyKind, Visibility, Expr};
  17 use syntax::attr;
  18 use syntax::codemap::dummy_spanned;
  19 use syntax::codemap::{ExpnInfo, NameAndSpan, MacroAttribute};
  20 use syntax::ext::base::ExtCtxt;
  21 use syntax::ext::base::Resolver;
  22 use syntax::ext::build::AstBuilder;
  23 use syntax::ext::expand::ExpansionConfig;
  24 use syntax::ext::hygiene::{Mark, SyntaxContext};
  25 use syntax::fold::{self, Folder};
  26 use syntax::parse::ParseSess;
  27 use syntax::ptr::P;
  28 use syntax::symbol::Symbol;
  29 use syntax::util::small_vector::SmallVector;
  30 use syntax_pos::{Span, DUMMY_SP};
  31
  32 use {AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
  33
  34 pub fn modify(sess: &ParseSess,
  35               resolver: &mut Resolver,
  36               krate: Crate,
  37               handler: &rustc_errors::Handler) -> ast::Crate {
  38     ExpandAllocatorDirectives {
  39         handler: handler,
  40         sess: sess,
  41         resolver: resolver,
  42         found: false,
  43     }.fold_crate(krate)
  44 }
  45
  46 struct ExpandAllocatorDirectives<'a> {
  47     found: bool,
  48     handler: &'a rustc_errors::Handler,
  49     sess: &'a ParseSess,
  50     resolver: &'a mut Resolver,
  51 }
  52
  53 impl<'a> Folder for ExpandAllocatorDirectives<'a> {
  54     fn fold_item(&mut self, item: P<Item>) -> SmallVector<P<Item>> {
  55         let name = if attr::contains_name(&item.attrs, "global_allocator") {
  56             "global_allocator"
  57         } else {
  58             return fold::noop_fold_item(item, self)
  59         };
  60         match item.node {
  61             ItemKind::Static(..) => {}
  62             _ => {
  63                 self.handler.span_err(item.span, "allocators must be statics");
  64                 return SmallVector::one(item)
  65             }
  66         }
  67
  68         if self.found {
  69             self.handler.span_err(item.span, "cannot define more than one \
  70                                               #[global_allocator]");
  71             return SmallVector::one(item)
  72         }
  73         self.found = true;
  74
  75         let mark = Mark::fresh(Mark::root());
  76         mark.set_expn_info(ExpnInfo {
  77             call_site: DUMMY_SP,
  78             callee: NameAndSpan {
  79                 format: MacroAttribute(Symbol::intern(name)),
  80                 span: None,
  81                 allow_internal_unstable: true,
  82             }
  83         });
  84         let span = Span {
  85             ctxt: SyntaxContext::empty().apply_mark(mark),
  86             ..item.span
  87         };
  88         let ecfg = ExpansionConfig::default(name.to_string());
  89         let mut f = AllocFnFactory {
  90             span: span,
  91             kind: AllocatorKind::Global,
  92             global: item.ident,
  93             alloc: Ident::from_str("alloc"),
  94             cx: ExtCtxt::new(self.sess, ecfg, self.resolver),
  95         };
  96         let super_path = f.cx.path(f.span, vec![
  97             Ident::from_str("super"),
  98             f.global,
  99         ]);
 100         let mut items = vec![
 101             f.cx.item_extern_crate(f.span, f.alloc),
 102             f.cx.item_use_simple(f.span, Visibility::Inherited, super_path),
 103         ];
 104         for method in ALLOCATOR_METHODS {
 105             items.push(f.allocator_fn(method));
 106         }
 107         let name = f.kind.fn_name("allocator_abi");
 108         let allocator_abi = Ident::with_empty_ctxt(Symbol::gensym(&name));
 109         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 110         let module = f.cx.monotonic_expander().fold_item(module).pop().unwrap();
 111
 112         let mut ret = SmallVector::new();
 113         ret.push(item);
 114         ret.push(module);
 115         return ret
 116     }
 117
 118     fn fold_mac(&mut self, mac: Mac) -> Mac {
 119         fold::noop_fold_mac(mac, self)
 120     }
 121 }
 122
 123 struct AllocFnFactory<'a> {
 124     span: Span,
 125     kind: AllocatorKind,
 126     global: Ident,
 127     alloc: Ident,
 128     cx: ExtCtxt<'a>,
 129 }
 130
 131 impl<'a> AllocFnFactory<'a> {
 132     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 133         let mut abi_args = Vec::new();
 134         let mut i = 0;
 135         let ref mut mk = || {
 136             let name = Ident::from_str(&format!("arg{}", i));
 137             i += 1;
 138             name
 139         };
 140         let args = method.inputs.iter().map(|ty| {
 141             self.arg_ty(ty, &mut abi_args, mk)
 142         }).collect();
 143         let result = self.call_allocator(method.name, args);
 144         let (output_ty, output_expr) =
 145             self.ret_ty(&method.output, &mut abi_args, mk, result);
 146         let kind = ItemKind::Fn(self.cx.fn_decl(abi_args, output_ty),
 147                                 Unsafety::Unsafe,
 148                                 dummy_spanned(Constness::NotConst),
 149                                 Abi::Rust,
 150                                 Generics::default(),
 151                                 self.cx.block_expr(output_expr));
 152         self.cx.item(self.span,
 153                      Ident::from_str(&self.kind.fn_name(method.name)),
 154                      self.attrs(),
 155                      kind)
 156     }
 157
 158     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 159         let method = self.cx.path(self.span, vec![
 160             self.alloc,
 161             Ident::from_str("heap"),
 162             Ident::from_str("Alloc"),
 163             Ident::from_str(method),
 164         ]);
 165         let method = self.cx.expr_path(method);
 166         let allocator = self.cx.path_ident(self.span, self.global);
 167         let allocator = self.cx.expr_path(allocator);
 168         let allocator = self.cx.expr_addr_of(self.span, allocator);
 169         let allocator = self.cx.expr_mut_addr_of(self.span, allocator);
 170         args.insert(0, allocator);
 171
 172         self.cx.expr_call(self.span, method, args)
 173     }
 174
 175     fn attrs(&self) -> Vec<Attribute> {
 176         let key = Symbol::intern("linkage");
 177         let value = LitKind::Str(Symbol::intern("external"), StrStyle::Cooked);
 178         let linkage = self.cx.meta_name_value(self.span, key, value);
 179
 180         let no_mangle = Symbol::intern("no_mangle");
 181         let no_mangle = self.cx.meta_word(self.span, no_mangle);
 182         vec![
 183             self.cx.attribute(self.span, linkage),
 184             self.cx.attribute(self.span, no_mangle),
 185         ]
 186     }
 187
 188     fn arg_ty(&self,
 189               ty: &AllocatorTy,
 190               args: &mut Vec<Arg>,
 191               mut ident: &mut FnMut() -> Ident) -> P<Expr> {
 192         match *ty {
 193             AllocatorTy::Layout => {
 194                 let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 195                 let ty_usize = self.cx.ty_path(usize);
 196                 let size = ident();
 197                 let align = ident();
 198                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 199                 args.push(self.cx.arg(self.span, align, ty_usize));
 200
 201                 let layout_new = self.cx.path(self.span, vec![
 202                     self.alloc,
 203                     Ident::from_str("heap"),
 204                     Ident::from_str("Layout"),
 205                     Ident::from_str("from_size_align_unchecked"),
 206                 ]);
 207                 let layout_new = self.cx.expr_path(layout_new);
 208                 let size = self.cx.expr_ident(self.span, size);
 209                 let align = self.cx.expr_ident(self.span, align);
 210                 let layout = self.cx.expr_call(self.span,
 211                                                layout_new,
 212                                                vec![size, align]);
 213                 layout
 214             }
 215
 216             AllocatorTy::LayoutRef => {
 217                 let ident = ident();
 218                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 219
 220                 // Convert our `arg: *const u8` via:
 221                 //
 222                 //      &*(arg as *const Layout)
 223                 let expr = self.cx.expr_ident(self.span, ident);
 224                 let expr = self.cx.expr_cast(self.span, expr, self.layout_ptr());
 225                 let expr = self.cx.expr_deref(self.span, expr);
 226                 self.cx.expr_addr_of(self.span, expr)
 227             }
 228
 229             AllocatorTy::AllocErr => {
 230                 // We're creating:
 231                 //
 232                 //      (*(arg as *const AllocErr)).clone()
 233                 let ident = ident();
 234                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 235                 let expr = self.cx.expr_ident(self.span, ident);
 236                 let expr = self.cx.expr_cast(self.span, expr, self.alloc_err_ptr());
 237                 let expr = self.cx.expr_deref(self.span, expr);
 238                 self.cx.expr_method_call(
 239                     self.span,
 240                     expr,
 241                     Ident::from_str("clone"),
 242                     Vec::new()
 243                 )
 244             }
 245
 246             AllocatorTy::Ptr => {
 247                 let ident = ident();
 248                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 249                 self.cx.expr_ident(self.span, ident)
 250             }
 251
 252             AllocatorTy::ResultPtr |
 253             AllocatorTy::ResultExcess |
 254             AllocatorTy::ResultUnit |
 255             AllocatorTy::Bang |
 256             AllocatorTy::UsizePair |
 257             AllocatorTy::Unit => {
 258                 panic!("can't convert AllocatorTy to an argument")
 259             }
 260         }
 261     }
 262
 263     fn ret_ty(&self,
 264               ty: &AllocatorTy,
 265               args: &mut Vec<Arg>,
 266               mut ident: &mut FnMut() -> Ident,
 267               expr: P<Expr>) -> (P<Ty>, P<Expr>)
 268     {
 269         match *ty {
 270             AllocatorTy::UsizePair => {
 271                 // We're creating:
 272                 //
 273                 //      let arg = #expr;
 274                 //      *min = arg.0;
 275                 //      *max = arg.1;
 276
 277                 let min = ident();
 278                 let max = ident();
 279
 280                 args.push(self.cx.arg(self.span, min, self.ptr_usize()));
 281                 args.push(self.cx.arg(self.span, max, self.ptr_usize()));
 282
 283                 let ident = ident();
 284                 let stmt = self.cx.stmt_let(self.span, false, ident, expr);
 285                 let min = self.cx.expr_ident(self.span, min);
 286                 let max = self.cx.expr_ident(self.span, max);
 287                 let layout = self.cx.expr_ident(self.span, ident);
 288                 let assign_min = self.cx.expr(self.span, ExprKind::Assign(
 289                     self.cx.expr_deref(self.span, min),
 290                     self.cx.expr_tup_field_access(self.span, layout.clone(), 0),
 291                 ));
 292                 let assign_min = self.cx.stmt_semi(assign_min);
 293                 let assign_max = self.cx.expr(self.span, ExprKind::Assign(
 294                     self.cx.expr_deref(self.span, max),
 295                     self.cx.expr_tup_field_access(self.span, layout.clone(), 1),
 296                 ));
 297                 let assign_max = self.cx.stmt_semi(assign_max);
 298
 299                 let stmts = vec![stmt, assign_min, assign_max];
 300                 let block = self.cx.block(self.span, stmts);
 301                 let ty_unit = self.cx.ty(self.span, TyKind::Tup(Vec::new()));
 302                 (ty_unit, self.cx.expr_block(block))
 303             }
 304
 305             AllocatorTy::ResultExcess => {
 306                 // We're creating:
 307                 //
 308                 //      match #expr {
 309                 //          Ok(ptr) => {
 310                 //              *excess = ptr.1;
 311                 //              ptr.0
 312                 //          }
 313                 //          Err(e) => {
 314                 //              ptr::write(err_ptr, e);
 315                 //              0 as *mut u8
 316                 //          }
 317                 //      }
 318
 319                 let excess_ptr = ident();
 320                 args.push(self.cx.arg(self.span, excess_ptr, self.ptr_usize()));
 321                 let excess_ptr = self.cx.expr_ident(self.span, excess_ptr);
 322
 323                 let err_ptr = ident();
 324                 args.push(self.cx.arg(self.span, err_ptr, self.ptr_u8()));
 325                 let err_ptr = self.cx.expr_ident(self.span, err_ptr);
 326                 let err_ptr = self.cx.expr_cast(self.span,
 327                                                 err_ptr,
 328                                                 self.alloc_err_ptr());
 329
 330                 let name = ident();
 331                 let ok_expr = {
 332                     let ptr = self.cx.expr_ident(self.span, name);
 333                     let write = self.cx.expr(self.span, ExprKind::Assign(
 334                         self.cx.expr_deref(self.span, excess_ptr),
 335                         self.cx.expr_tup_field_access(self.span, ptr.clone(), 1),
 336                     ));
 337                     let write = self.cx.stmt_semi(write);
 338                     let ret = self.cx.expr_tup_field_access(self.span,
 339                                                             ptr.clone(),
 340                                                             0);
 341                     let ret = self.cx.stmt_expr(ret);
 342                     let block = self.cx.block(self.span, vec![write, ret]);
 343                     self.cx.expr_block(block)
 344                 };
 345                 let pat = self.cx.pat_ident(self.span, name);
 346                 let ok = self.cx.path_ident(self.span, Ident::from_str("Ok"));
 347                 let ok = self.cx.pat_tuple_struct(self.span, ok, vec![pat]);
 348                 let ok = self.cx.arm(self.span, vec![ok], ok_expr);
 349
 350                 let name = ident();
 351                 let err_expr = {
 352                     let err = self.cx.expr_ident(self.span, name);
 353                     let write = self.cx.path(self.span, vec![
 354                         self.alloc,
 355                         Ident::from_str("heap"),
 356                         Ident::from_str("__core"),
 357                         Ident::from_str("ptr"),
 358                         Ident::from_str("write"),
 359                     ]);
 360                     let write = self.cx.expr_path(write);
 361                     let write = self.cx.expr_call(self.span, write,
 362                                                   vec![err_ptr, err]);
 363                     let write = self.cx.stmt_semi(write);
 364                     let null = self.cx.expr_usize(self.span, 0);
 365                     let null = self.cx.expr_cast(self.span, null, self.ptr_u8());
 366                     let null = self.cx.stmt_expr(null);
 367                     let block = self.cx.block(self.span, vec![write, null]);
 368                     self.cx.expr_block(block)
 369                 };
 370                 let pat = self.cx.pat_ident(self.span, name);
 371                 let err = self.cx.path_ident(self.span, Ident::from_str("Err"));
 372                 let err = self.cx.pat_tuple_struct(self.span, err, vec![pat]);
 373                 let err = self.cx.arm(self.span, vec![err], err_expr);
 374
 375                 let expr = self.cx.expr_match(self.span, expr, vec![ok, err]);
 376                 (self.ptr_u8(), expr)
 377             }
 378
 379             AllocatorTy::ResultPtr => {
 380                 // We're creating:
 381                 //
 382                 //      match #expr {
 383                 //          Ok(ptr) => ptr,
 384                 //          Err(e) => {
 385                 //              ptr::write(err_ptr, e);
 386                 //              0 as *mut u8
 387                 //          }
 388                 //      }
 389
 390                 let err_ptr = ident();
 391                 args.push(self.cx.arg(self.span, err_ptr, self.ptr_u8()));
 392                 let err_ptr = self.cx.expr_ident(self.span, err_ptr);
 393                 let err_ptr = self.cx.expr_cast(self.span,
 394                                                 err_ptr,
 395                                                 self.alloc_err_ptr());
 396
 397                 let name = ident();
 398                 let ok_expr = self.cx.expr_ident(self.span, name);
 399                 let pat = self.cx.pat_ident(self.span, name);
 400                 let ok = self.cx.path_ident(self.span, Ident::from_str("Ok"));
 401                 let ok = self.cx.pat_tuple_struct(self.span, ok, vec![pat]);
 402                 let ok = self.cx.arm(self.span, vec![ok], ok_expr);
 403
 404                 let name = ident();
 405                 let err_expr = {
 406                     let err = self.cx.expr_ident(self.span, name);
 407                     let write = self.cx.path(self.span, vec![
 408                         self.alloc,
 409                         Ident::from_str("heap"),
 410                         Ident::from_str("__core"),
 411                         Ident::from_str("ptr"),
 412                         Ident::from_str("write"),
 413                     ]);
 414                     let write = self.cx.expr_path(write);
 415                     let write = self.cx.expr_call(self.span, write,
 416                                                   vec![err_ptr, err]);
 417                     let write = self.cx.stmt_semi(write);
 418                     let null = self.cx.expr_usize(self.span, 0);
 419                     let null = self.cx.expr_cast(self.span, null, self.ptr_u8());
 420                     let null = self.cx.stmt_expr(null);
 421                     let block = self.cx.block(self.span, vec![write, null]);
 422                     self.cx.expr_block(block)
 423                 };
 424                 let pat = self.cx.pat_ident(self.span, name);
 425                 let err = self.cx.path_ident(self.span, Ident::from_str("Err"));
 426                 let err = self.cx.pat_tuple_struct(self.span, err, vec![pat]);
 427                 let err = self.cx.arm(self.span, vec![err], err_expr);
 428
 429                 let expr = self.cx.expr_match(self.span, expr, vec![ok, err]);
 430                 (self.ptr_u8(), expr)
 431             }
 432
 433             AllocatorTy::ResultUnit => {
 434                 // We're creating:
 435                 //
 436                 //      #expr.is_ok() as u8
 437
 438                 let cast = self.cx.expr_method_call(
 439                     self.span,
 440                     expr,
 441                     Ident::from_str("is_ok"),
 442                     Vec::new()
 443                 );
 444                 let u8 = self.cx.path_ident(self.span, Ident::from_str("u8"));
 445                 let u8 = self.cx.ty_path(u8);
 446                 let cast = self.cx.expr_cast(self.span, cast, u8.clone());
 447                 (u8, cast)
 448             }
 449
 450             AllocatorTy::Bang => {
 451                 (self.cx.ty(self.span, TyKind::Never), expr)
 452             }
 453
 454             AllocatorTy::Unit => {
 455                 (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr)
 456             }
 457
 458             AllocatorTy::AllocErr |
 459             AllocatorTy::Layout |
 460             AllocatorTy::LayoutRef |
 461             AllocatorTy::Ptr => {
 462                 panic!("can't convert AllocatorTy to an output")
 463             }
 464         }
 465     }
 466
 467     fn ptr_u8(&self) -> P<Ty> {
 468         let u8 = self.cx.path_ident(self.span, Ident::from_str("u8"));
 469         let ty_u8 = self.cx.ty_path(u8);
 470         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 471     }
 472
 473     fn ptr_usize(&self) -> P<Ty> {
 474         let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 475         let ty_usize = self.cx.ty_path(usize);
 476         self.cx.ty_ptr(self.span, ty_usize, Mutability::Mutable)
 477     }
 478
 479     fn layout_ptr(&self) -> P<Ty> {
 480         let layout = self.cx.path(self.span, vec![
 481             self.alloc,
 482             Ident::from_str("heap"),
 483             Ident::from_str("Layout"),
 484         ]);
 485         let layout = self.cx.ty_path(layout);
 486         self.cx.ty_ptr(self.span, layout, Mutability::Mutable)
 487     }
 488
 489     fn alloc_err_ptr(&self) -> P<Ty> {
 490         let err = self.cx.path(self.span, vec![
 491             self.alloc,
 492             Ident::from_str("heap"),
 493             Ident::from_str("AllocErr"),
 494         ]);
 495         let err = self.cx.ty_path(err);
 496         self.cx.ty_ptr(self.span, err, Mutability::Mutable)
 497     }
 498 }