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, VisibilityKind, Expr};
  17 use syntax::attr;
  18 use syntax::codemap::{dummy_spanned, respan};
  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,
  40         sess,
  41         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                 allow_internal_unsafe: false,
  83             }
  84         });
  85         let span = item.span.with_ctxt(SyntaxContext::empty().apply_mark(mark));
  86         let ecfg = ExpansionConfig::default(name.to_string());
  87         let mut f = AllocFnFactory {
  88             span,
  89             kind: AllocatorKind::Global,
  90             global: item.ident,
  91             alloc: Ident::from_str("alloc"),
  92             cx: ExtCtxt::new(self.sess, ecfg, self.resolver),
  93         };
  94         let super_path = f.cx.path(f.span, vec![
  95             Ident::from_str("super"),
  96             f.global,
  97         ]);
  98         let mut items = vec![
  99             f.cx.item_extern_crate(f.span, f.alloc),
 100             f.cx.item_use_simple(
 101                 f.span,
 102                 respan(f.span.shrink_to_lo(), VisibilityKind::Inherited),
 103                 super_path,
 104             ),
 105         ];
 106         for method in ALLOCATOR_METHODS {
 107             items.push(f.allocator_fn(method));
 108         }
 109         let name = f.kind.fn_name("allocator_abi");
 110         let allocator_abi = Ident::with_empty_ctxt(Symbol::gensym(&name));
 111         let module = f.cx.item_mod(span, span, allocator_abi, Vec::new(), items);
 112         let module = f.cx.monotonic_expander().fold_item(module).pop().unwrap();
 113
 114         let mut ret = SmallVector::new();
 115         ret.push(item);
 116         ret.push(module);
 117         return ret
 118     }
 119
 120     fn fold_mac(&mut self, mac: Mac) -> Mac {
 121         fold::noop_fold_mac(mac, self)
 122     }
 123 }
 124
 125 struct AllocFnFactory<'a> {
 126     span: Span,
 127     kind: AllocatorKind,
 128     global: Ident,
 129     alloc: Ident,
 130     cx: ExtCtxt<'a>,
 131 }
 132
 133 impl<'a> AllocFnFactory<'a> {
 134     fn allocator_fn(&self, method: &AllocatorMethod) -> P<Item> {
 135         let mut abi_args = Vec::new();
 136         let mut i = 0;
 137         let ref mut mk = || {
 138             let name = Ident::from_str(&format!("arg{}", i));
 139             i += 1;
 140             name
 141         };
 142         let args = method.inputs.iter().map(|ty| {
 143             self.arg_ty(ty, &mut abi_args, mk)
 144         }).collect();
 145         let result = self.call_allocator(method.name, args);
 146         let (output_ty, output_expr) =
 147             self.ret_ty(&method.output, &mut abi_args, mk, result);
 148         let kind = ItemKind::Fn(self.cx.fn_decl(abi_args, ast::FunctionRetTy::Ty(output_ty)),
 149                                 Unsafety::Unsafe,
 150                                 dummy_spanned(Constness::NotConst),
 151                                 Abi::Rust,
 152                                 Generics::default(),
 153                                 self.cx.block_expr(output_expr));
 154         self.cx.item(self.span,
 155                      Ident::from_str(&self.kind.fn_name(method.name)),
 156                      self.attrs(),
 157                      kind)
 158     }
 159
 160     fn call_allocator(&self, method: &str, mut args: Vec<P<Expr>>) -> P<Expr> {
 161         let method = self.cx.path(self.span, vec![
 162             self.alloc,
 163             Ident::from_str("heap"),
 164             Ident::from_str("Alloc"),
 165             Ident::from_str(method),
 166         ]);
 167         let method = self.cx.expr_path(method);
 168         let allocator = self.cx.path_ident(self.span, self.global);
 169         let allocator = self.cx.expr_path(allocator);
 170         let allocator = self.cx.expr_addr_of(self.span, allocator);
 171         let allocator = self.cx.expr_mut_addr_of(self.span, allocator);
 172         args.insert(0, allocator);
 173
 174         self.cx.expr_call(self.span, method, args)
 175     }
 176
 177     fn attrs(&self) -> Vec<Attribute> {
 178         let key = Symbol::intern("linkage");
 179         let value = LitKind::Str(Symbol::intern("external"), StrStyle::Cooked);
 180         let linkage = self.cx.meta_name_value(self.span, key, value);
 181
 182         let no_mangle = Symbol::intern("no_mangle");
 183         let no_mangle = self.cx.meta_word(self.span, no_mangle);
 184
 185         let special = Symbol::intern("rustc_std_internal_symbol");
 186         let special = self.cx.meta_word(self.span, special);
 187         vec![
 188             self.cx.attribute(self.span, linkage),
 189             self.cx.attribute(self.span, no_mangle),
 190             self.cx.attribute(self.span, special),
 191         ]
 192     }
 193
 194     fn arg_ty(&self,
 195               ty: &AllocatorTy,
 196               args: &mut Vec<Arg>,
 197               ident: &mut FnMut() -> Ident) -> P<Expr> {
 198         match *ty {
 199             AllocatorTy::Layout => {
 200                 let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 201                 let ty_usize = self.cx.ty_path(usize);
 202                 let size = ident();
 203                 let align = ident();
 204                 args.push(self.cx.arg(self.span, size, ty_usize.clone()));
 205                 args.push(self.cx.arg(self.span, align, ty_usize));
 206
 207                 let layout_new = self.cx.path(self.span, vec![
 208                     self.alloc,
 209                     Ident::from_str("heap"),
 210                     Ident::from_str("Layout"),
 211                     Ident::from_str("from_size_align_unchecked"),
 212                 ]);
 213                 let layout_new = self.cx.expr_path(layout_new);
 214                 let size = self.cx.expr_ident(self.span, size);
 215                 let align = self.cx.expr_ident(self.span, align);
 216                 let layout = self.cx.expr_call(self.span,
 217                                                layout_new,
 218                                                vec![size, align]);
 219                 layout
 220             }
 221
 222             AllocatorTy::LayoutRef => {
 223                 let ident = ident();
 224                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 225
 226                 // Convert our `arg: *const u8` via:
 227                 //
 228                 //      &*(arg as *const Layout)
 229                 let expr = self.cx.expr_ident(self.span, ident);
 230                 let expr = self.cx.expr_cast(self.span, expr, self.layout_ptr());
 231                 let expr = self.cx.expr_deref(self.span, expr);
 232                 self.cx.expr_addr_of(self.span, expr)
 233             }
 234
 235             AllocatorTy::AllocErr => {
 236                 // We're creating:
 237                 //
 238                 //      (*(arg as *const AllocErr)).clone()
 239                 let ident = ident();
 240                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 241                 let expr = self.cx.expr_ident(self.span, ident);
 242                 let expr = self.cx.expr_cast(self.span, expr, self.alloc_err_ptr());
 243                 let expr = self.cx.expr_deref(self.span, expr);
 244                 self.cx.expr_method_call(
 245                     self.span,
 246                     expr,
 247                     Ident::from_str("clone"),
 248                     Vec::new()
 249                 )
 250             }
 251
 252             AllocatorTy::Ptr => {
 253                 let ident = ident();
 254                 args.push(self.cx.arg(self.span, ident, self.ptr_u8()));
 255                 self.cx.expr_ident(self.span, ident)
 256             }
 257
 258             AllocatorTy::ResultPtr |
 259             AllocatorTy::ResultExcess |
 260             AllocatorTy::ResultUnit |
 261             AllocatorTy::Bang |
 262             AllocatorTy::UsizePair |
 263             AllocatorTy::Unit => {
 264                 panic!("can't convert AllocatorTy to an argument")
 265             }
 266         }
 267     }
 268
 269     fn ret_ty(&self,
 270               ty: &AllocatorTy,
 271               args: &mut Vec<Arg>,
 272               ident: &mut FnMut() -> Ident,
 273               expr: P<Expr>) -> (P<Ty>, P<Expr>)
 274     {
 275         match *ty {
 276             AllocatorTy::UsizePair => {
 277                 // We're creating:
 278                 //
 279                 //      let arg = #expr;
 280                 //      *min = arg.0;
 281                 //      *max = arg.1;
 282
 283                 let min = ident();
 284                 let max = ident();
 285
 286                 args.push(self.cx.arg(self.span, min, self.ptr_usize()));
 287                 args.push(self.cx.arg(self.span, max, self.ptr_usize()));
 288
 289                 let ident = ident();
 290                 let stmt = self.cx.stmt_let(self.span, false, ident, expr);
 291                 let min = self.cx.expr_ident(self.span, min);
 292                 let max = self.cx.expr_ident(self.span, max);
 293                 let layout = self.cx.expr_ident(self.span, ident);
 294                 let assign_min = self.cx.expr(self.span, ExprKind::Assign(
 295                     self.cx.expr_deref(self.span, min),
 296                     self.cx.expr_tup_field_access(self.span, layout.clone(), 0),
 297                 ));
 298                 let assign_min = self.cx.stmt_semi(assign_min);
 299                 let assign_max = self.cx.expr(self.span, ExprKind::Assign(
 300                     self.cx.expr_deref(self.span, max),
 301                     self.cx.expr_tup_field_access(self.span, layout.clone(), 1),
 302                 ));
 303                 let assign_max = self.cx.stmt_semi(assign_max);
 304
 305                 let stmts = vec![stmt, assign_min, assign_max];
 306                 let block = self.cx.block(self.span, stmts);
 307                 let ty_unit = self.cx.ty(self.span, TyKind::Tup(Vec::new()));
 308                 (ty_unit, self.cx.expr_block(block))
 309             }
 310
 311             AllocatorTy::ResultExcess => {
 312                 // We're creating:
 313                 //
 314                 //      match #expr {
 315                 //          Ok(ptr) => {
 316                 //              *excess = ptr.1;
 317                 //              ptr.0
 318                 //          }
 319                 //          Err(e) => {
 320                 //              ptr::write(err_ptr, e);
 321                 //              0 as *mut u8
 322                 //          }
 323                 //      }
 324
 325                 let excess_ptr = ident();
 326                 args.push(self.cx.arg(self.span, excess_ptr, self.ptr_usize()));
 327                 let excess_ptr = self.cx.expr_ident(self.span, excess_ptr);
 328
 329                 let err_ptr = ident();
 330                 args.push(self.cx.arg(self.span, err_ptr, self.ptr_u8()));
 331                 let err_ptr = self.cx.expr_ident(self.span, err_ptr);
 332                 let err_ptr = self.cx.expr_cast(self.span,
 333                                                 err_ptr,
 334                                                 self.alloc_err_ptr());
 335
 336                 let name = ident();
 337                 let ok_expr = {
 338                     let ptr = self.cx.expr_ident(self.span, name);
 339                     let write = self.cx.expr(self.span, ExprKind::Assign(
 340                         self.cx.expr_deref(self.span, excess_ptr),
 341                         self.cx.expr_tup_field_access(self.span, ptr.clone(), 1),
 342                     ));
 343                     let write = self.cx.stmt_semi(write);
 344                     let ret = self.cx.expr_tup_field_access(self.span,
 345                                                             ptr.clone(),
 346                                                             0);
 347                     let ret = self.cx.stmt_expr(ret);
 348                     let block = self.cx.block(self.span, vec![write, ret]);
 349                     self.cx.expr_block(block)
 350                 };
 351                 let pat = self.cx.pat_ident(self.span, name);
 352                 let ok = self.cx.path_ident(self.span, Ident::from_str("Ok"));
 353                 let ok = self.cx.pat_tuple_struct(self.span, ok, vec![pat]);
 354                 let ok = self.cx.arm(self.span, vec![ok], ok_expr);
 355
 356                 let name = ident();
 357                 let err_expr = {
 358                     let err = self.cx.expr_ident(self.span, name);
 359                     let write = self.cx.path(self.span, vec![
 360                         self.alloc,
 361                         Ident::from_str("heap"),
 362                         Ident::from_str("__core"),
 363                         Ident::from_str("ptr"),
 364                         Ident::from_str("write"),
 365                     ]);
 366                     let write = self.cx.expr_path(write);
 367                     let write = self.cx.expr_call(self.span, write,
 368                                                   vec![err_ptr, err]);
 369                     let write = self.cx.stmt_semi(write);
 370                     let null = self.cx.expr_usize(self.span, 0);
 371                     let null = self.cx.expr_cast(self.span, null, self.ptr_u8());
 372                     let null = self.cx.stmt_expr(null);
 373                     let block = self.cx.block(self.span, vec![write, null]);
 374                     self.cx.expr_block(block)
 375                 };
 376                 let pat = self.cx.pat_ident(self.span, name);
 377                 let err = self.cx.path_ident(self.span, Ident::from_str("Err"));
 378                 let err = self.cx.pat_tuple_struct(self.span, err, vec![pat]);
 379                 let err = self.cx.arm(self.span, vec![err], err_expr);
 380
 381                 let expr = self.cx.expr_match(self.span, expr, vec![ok, err]);
 382                 (self.ptr_u8(), expr)
 383             }
 384
 385             AllocatorTy::ResultPtr => {
 386                 // We're creating:
 387                 //
 388                 //      match #expr {
 389                 //          Ok(ptr) => ptr,
 390                 //          Err(e) => {
 391                 //              ptr::write(err_ptr, e);
 392                 //              0 as *mut u8
 393                 //          }
 394                 //      }
 395
 396                 let err_ptr = ident();
 397                 args.push(self.cx.arg(self.span, err_ptr, self.ptr_u8()));
 398                 let err_ptr = self.cx.expr_ident(self.span, err_ptr);
 399                 let err_ptr = self.cx.expr_cast(self.span,
 400                                                 err_ptr,
 401                                                 self.alloc_err_ptr());
 402
 403                 let name = ident();
 404                 let ok_expr = self.cx.expr_ident(self.span, name);
 405                 let pat = self.cx.pat_ident(self.span, name);
 406                 let ok = self.cx.path_ident(self.span, Ident::from_str("Ok"));
 407                 let ok = self.cx.pat_tuple_struct(self.span, ok, vec![pat]);
 408                 let ok = self.cx.arm(self.span, vec![ok], ok_expr);
 409
 410                 let name = ident();
 411                 let err_expr = {
 412                     let err = self.cx.expr_ident(self.span, name);
 413                     let write = self.cx.path(self.span, vec![
 414                         self.alloc,
 415                         Ident::from_str("heap"),
 416                         Ident::from_str("__core"),
 417                         Ident::from_str("ptr"),
 418                         Ident::from_str("write"),
 419                     ]);
 420                     let write = self.cx.expr_path(write);
 421                     let write = self.cx.expr_call(self.span, write,
 422                                                   vec![err_ptr, err]);
 423                     let write = self.cx.stmt_semi(write);
 424                     let null = self.cx.expr_usize(self.span, 0);
 425                     let null = self.cx.expr_cast(self.span, null, self.ptr_u8());
 426                     let null = self.cx.stmt_expr(null);
 427                     let block = self.cx.block(self.span, vec![write, null]);
 428                     self.cx.expr_block(block)
 429                 };
 430                 let pat = self.cx.pat_ident(self.span, name);
 431                 let err = self.cx.path_ident(self.span, Ident::from_str("Err"));
 432                 let err = self.cx.pat_tuple_struct(self.span, err, vec![pat]);
 433                 let err = self.cx.arm(self.span, vec![err], err_expr);
 434
 435                 let expr = self.cx.expr_match(self.span, expr, vec![ok, err]);
 436                 (self.ptr_u8(), expr)
 437             }
 438
 439             AllocatorTy::ResultUnit => {
 440                 // We're creating:
 441                 //
 442                 //      #expr.is_ok() as u8
 443
 444                 let cast = self.cx.expr_method_call(
 445                     self.span,
 446                     expr,
 447                     Ident::from_str("is_ok"),
 448                     Vec::new()
 449                 );
 450                 let u8 = self.cx.path_ident(self.span, Ident::from_str("u8"));
 451                 let u8 = self.cx.ty_path(u8);
 452                 let cast = self.cx.expr_cast(self.span, cast, u8.clone());
 453                 (u8, cast)
 454             }
 455
 456             AllocatorTy::Bang => {
 457                 (self.cx.ty(self.span, TyKind::Never), expr)
 458             }
 459
 460             AllocatorTy::Unit => {
 461                 (self.cx.ty(self.span, TyKind::Tup(Vec::new())), expr)
 462             }
 463
 464             AllocatorTy::AllocErr |
 465             AllocatorTy::Layout |
 466             AllocatorTy::LayoutRef |
 467             AllocatorTy::Ptr => {
 468                 panic!("can't convert AllocatorTy to an output")
 469             }
 470         }
 471     }
 472
 473     fn ptr_u8(&self) -> P<Ty> {
 474         let u8 = self.cx.path_ident(self.span, Ident::from_str("u8"));
 475         let ty_u8 = self.cx.ty_path(u8);
 476         self.cx.ty_ptr(self.span, ty_u8, Mutability::Mutable)
 477     }
 478
 479     fn ptr_usize(&self) -> P<Ty> {
 480         let usize = self.cx.path_ident(self.span, Ident::from_str("usize"));
 481         let ty_usize = self.cx.ty_path(usize);
 482         self.cx.ty_ptr(self.span, ty_usize, Mutability::Mutable)
 483     }
 484
 485     fn layout_ptr(&self) -> P<Ty> {
 486         let layout = self.cx.path(self.span, vec![
 487             self.alloc,
 488             Ident::from_str("heap"),
 489             Ident::from_str("Layout"),
 490         ]);
 491         let layout = self.cx.ty_path(layout);
 492         self.cx.ty_ptr(self.span, layout, Mutability::Mutable)
 493     }
 494
 495     fn alloc_err_ptr(&self) -> P<Ty> {
 496         let err = self.cx.path(self.span, vec![
 497             self.alloc,
 498             Ident::from_str("heap"),
 499             Ident::from_str("AllocErr"),
 500         ]);
 501         let err = self.cx.ty_path(err);
 502         self.cx.ty_ptr(self.span, err, Mutability::Mutable)
 503     }
 504 }