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