src/librustc_mir/transform/generator.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 //! This is the implementation of the pass which transforms generators into state machines.
  12 //!
  13 //! MIR generation for generators creates a function which has a self argument which
  14 //! passes by value. This argument is effectively a generator type which only contains upvars and
  15 //! is only used for this argument inside the MIR for the generator.
  16 //! It is passed by value to enable upvars to be moved out of it. Drop elaboration runs on that
  17 //! MIR before this pass and creates drop flags for MIR locals.
  18 //! It will also drop the generator argument (which only consists of upvars) if any of the upvars
  19 //! are moved out of. This pass elaborates the drops of upvars / generator argument in the case
  20 //! that none of the upvars were moved out of. This is because we cannot have any drops of this
  21 //! generator in the MIR, since it is used to create the drop glue for the generator. We'd get
  22 //! infinite recursion otherwise.
  23 //!
  24 //! This pass creates the implementation for the Generator::resume function and the drop shim
  25 //! for the generator based on the MIR input. It converts the generator argument from Self to
  26 //! &mut Self adding derefs in the MIR as needed. It computes the final layout of the generator
  27 //! struct which looks like this:
  28 //!     First upvars are stored
  29 //!     It is followed by the generator state field.
  30 //!     Then finally the MIR locals which are live across a suspension point are stored.
  31 //!
  32 //!     struct Generator {
  33 //!         upvars...,
  34 //!         state: u32,
  35 //!         mir_locals...,
  36 //!     }
  37 //!
  38 //! This pass computes the meaning of the state field and the MIR locals which are live
  39 //! across a suspension point. There are however two hardcoded generator states:
  40 //!     0 - Generator have not been resumed yet
  41 //!     1 - Generator has been poisoned
  42 //!
  43 //! It also rewrites `return x` and `yield y` as setting a new generator state and returning
  44 //! GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
  45 //! MIR locals which are live across a suspension point are moved to the generator struct
  46 //! with references to them being updated with references to the generator struct.
  47 //!
  48 //! The pass creates two functions which have a switch on the generator state giving
  49 //! the action to take.
  50 //!
  51 //! One of them is the implementation of Generator::resume.
  52 //! For generators which have already returned it panics.
  53 //! For generators with state 0 (unresumed) it starts the execution of the generator.
  54 //! For generators with state 1 (poisoned) it panics.
  55 //! Otherwise it continues the execution from the last suspension point.
  56 //!
  57 //! The other function is the drop glue for the generator.
  58 //! For generators which have already returned it does nothing.
  59 //! For generators with state 0 (unresumed) it drops the upvars of the generator.
  60 //! For generators with state 1 (poisoned) it does nothing.
  61 //! Otherwise it drops all the values in scope at the last suspension point.
  62
  63 use rustc::hir;
  64 use rustc::hir::def_id::DefId;
  65 use rustc::middle::const_val::ConstVal;
  66 use rustc::mir::*;
  67 use rustc::mir::transform::{MirPass, MirSource};
  68 use rustc::mir::visit::{LvalueContext, MutVisitor};
  69 use rustc::ty::{self, TyCtxt, AdtDef, Ty, GeneratorInterior};
  70 use rustc::ty::subst::{Kind, Substs};
  71 use util::dump_mir;
  72 use util::liveness;
  73 use rustc_const_math::ConstInt;
  74 use rustc_data_structures::indexed_vec::Idx;
  75 use std::collections::HashMap;
  76 use std::borrow::Cow;
  77 use std::iter::once;
  78 use std::mem;
  79 use transform::simplify;
  80 use transform::no_landing_pads::no_landing_pads;
  81
  82 pub struct StateTransform;
  83
  84 struct RenameLocalVisitor {
  85     from: Local,
  86     to: Local,
  87 }
  88
  89 impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor {
  90     fn visit_local(&mut self,
  91                    local: &mut Local,
  92                    _: LvalueContext<'tcx>,
  93                    _: Location) {
  94         if *local == self.from {
  95             *local = self.to;
  96         }
  97     }
  98 }
  99
 100 struct DerefArgVisitor;
 101
 102 impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor {
 103     fn visit_local(&mut self,
 104                    local: &mut Local,
 105                    _: LvalueContext<'tcx>,
 106                    _: Location) {
 107         assert_ne!(*local, self_arg());
 108     }
 109
 110     fn visit_lvalue(&mut self,
 111                     lvalue: &mut Lvalue<'tcx>,
 112                     context: LvalueContext<'tcx>,
 113                     location: Location) {
 114         if *lvalue == Lvalue::Local(self_arg()) {
 115             *lvalue = Lvalue::Projection(Box::new(Projection {
 116                 base: lvalue.clone(),
 117                 elem: ProjectionElem::Deref,
 118             }));
 119         } else {
 120             self.super_lvalue(lvalue, context, location);
 121         }
 122     }
 123 }
 124
 125 fn self_arg() -> Local {
 126     Local::new(1)
 127 }
 128
 129 struct TransformVisitor<'a, 'tcx: 'a> {
 130     tcx: TyCtxt<'a, 'tcx, 'tcx>,
 131     state_adt_ref: &'tcx AdtDef,
 132     state_substs: &'tcx Substs<'tcx>,
 133
 134     // The index of the generator state in the generator struct
 135     state_field: usize,
 136
 137     // Mapping from Local to (type of local, generator struct index)
 138     remap: HashMap<Local, (Ty<'tcx>, usize)>,
 139
 140     // The number of generator states. 0 is unresumed, 1 is poisoned. So this is initialized to 2
 141     bb_target_count: u32,
 142
 143     // Map from a (which block to resume execution at, which block to use to drop the generator)
 144     // to a generator state
 145     bb_targets: HashMap<(BasicBlock, Option<BasicBlock>), u32>,
 146
 147     // The original RETURN_POINTER local
 148     new_ret_local: Local,
 149
 150     // The block to resume execution when for Return
 151     return_block: BasicBlock,
 152 }
 153
 154 impl<'a, 'tcx> TransformVisitor<'a, 'tcx> {
 155     // Make a GeneratorState rvalue
 156     fn make_state(&self, idx: usize, val: Operand<'tcx>) -> Rvalue<'tcx> {
 157         let adt = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None);
 158         Rvalue::Aggregate(box adt, vec![val])
 159     }
 160
 161     // Create a Lvalue referencing a generator struct field
 162     fn make_field(&self, idx: usize, ty: Ty<'tcx>) -> Lvalue<'tcx> {
 163         let base = Lvalue::Local(self_arg());
 164         let field = Projection {
 165             base: base,
 166             elem: ProjectionElem::Field(Field::new(idx), ty),
 167         };
 168         Lvalue::Projection(Box::new(field))
 169     }
 170
 171     // Create a statement which changes the generator state
 172     fn set_state(&self, state_disc: u32, source_info: SourceInfo) -> Statement<'tcx> {
 173         let state = self.make_field(self.state_field, self.tcx.types.u32);
 174         let val = Operand::Constant(box Constant {
 175             span: source_info.span,
 176             ty: self.tcx.types.u32,
 177             literal: Literal::Value {
 178                 value: self.tcx.mk_const(ty::Const {
 179                     val: ConstVal::Integral(ConstInt::U32(state_disc)),
 180                     ty: self.tcx.types.u32
 181                 }),
 182             },
 183         });
 184         Statement {
 185             source_info,
 186             kind: StatementKind::Assign(state, Rvalue::Use(val)),
 187         }
 188     }
 189 }
 190
 191 impl<'a, 'tcx> MutVisitor<'tcx> for TransformVisitor<'a, 'tcx> {
 192     fn visit_local(&mut self,
 193                    local: &mut Local,
 194                    _: LvalueContext<'tcx>,
 195                    _: Location) {
 196         assert_eq!(self.remap.get(local), None);
 197     }
 198
 199     fn visit_lvalue(&mut self,
 200                     lvalue: &mut Lvalue<'tcx>,
 201                     context: LvalueContext<'tcx>,
 202                     location: Location) {
 203         if let Lvalue::Local(l) = *lvalue {
 204             // Replace an Local in the remap with a generator struct access
 205             if let Some(&(ty, idx)) = self.remap.get(&l) {
 206                 *lvalue = self.make_field(idx, ty);
 207             }
 208         } else {
 209             self.super_lvalue(lvalue, context, location);
 210         }
 211     }
 212
 213     fn visit_basic_block_data(&mut self,
 214                               block: BasicBlock,
 215                               data: &mut BasicBlockData<'tcx>) {
 216         // Remove StorageLive and StorageDead statements for remapped locals
 217         data.retain_statements(|s| {
 218             match s.kind {
 219                 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
 220                     !self.remap.contains_key(&l)
 221                 }
 222                 _ => true
 223             }
 224         });
 225
 226         let ret_val = match data.terminator().kind {
 227             TerminatorKind::Return => Some((1,
 228                 self.return_block,
 229                 Operand::Consume(Lvalue::Local(self.new_ret_local)),
 230                 None)),
 231             TerminatorKind::Yield { ref value, resume, drop } => Some((0,
 232                 resume,
 233                 value.clone(),
 234                 drop)),
 235             _ => None
 236         };
 237
 238         if let Some((state_idx, resume, v, drop)) = ret_val {
 239             let bb_idx = {
 240                 let bb_targets = &mut self.bb_targets;
 241                 let bb_target = &mut self.bb_target_count;
 242                 *bb_targets.entry((resume, drop)).or_insert_with(|| {
 243                     let target = *bb_target;
 244                     *bb_target = target.checked_add(1).unwrap();
 245                     target
 246                 })
 247             };
 248             let source_info = data.terminator().source_info;
 249             data.statements.push(self.set_state(bb_idx, source_info));
 250             data.statements.push(Statement {
 251                 source_info,
 252                 kind: StatementKind::Assign(Lvalue::Local(RETURN_POINTER),
 253                     self.make_state(state_idx, v)),
 254             });
 255             data.terminator.as_mut().unwrap().kind = TerminatorKind::Return;
 256         }
 257
 258         self.super_basic_block_data(block, data);
 259     }
 260 }
 261
 262 fn make_generator_state_argument_indirect<'a, 'tcx>(
 263                 tcx: TyCtxt<'a, 'tcx, 'tcx>,
 264                 def_id: DefId,
 265                 mir: &mut Mir<'tcx>) {
 266     let gen_ty = mir.local_decls.raw[1].ty;
 267
 268     let region = ty::ReFree(ty::FreeRegion {
 269         scope: def_id,
 270         bound_region: ty::BoundRegion::BrEnv,
 271     });
 272
 273     let region = tcx.mk_region(region);
 274
 275     let ref_gen_ty = tcx.mk_ref(region, ty::TypeAndMut {
 276         ty: gen_ty,
 277         mutbl: hir::MutMutable
 278     });
 279
 280     // Replace the by value generator argument
 281     mir.local_decls.raw[1].ty = ref_gen_ty;
 282
 283     // Add a deref to accesses of the generator state
 284     DerefArgVisitor.visit_mir(mir);
 285 }
 286
 287 fn replace_result_variable<'tcx>(ret_ty: Ty<'tcx>,
 288                             mir: &mut Mir<'tcx>) -> Local {
 289     let source_info = SourceInfo {
 290         span: mir.span,
 291         scope: ARGUMENT_VISIBILITY_SCOPE,
 292     };
 293
 294     let new_ret = LocalDecl {
 295         mutability: Mutability::Mut,
 296         ty: ret_ty,
 297         name: None,
 298         source_info,
 299         internal: false,
 300         is_user_variable: false,
 301     };
 302     let new_ret_local = Local::new(mir.local_decls.len());
 303     mir.local_decls.push(new_ret);
 304     mir.local_decls.swap(0, new_ret_local.index());
 305
 306     RenameLocalVisitor {
 307         from: RETURN_POINTER,
 308         to: new_ret_local,
 309     }.visit_mir(mir);
 310
 311     new_ret_local
 312 }
 313
 314 fn locals_live_across_suspend_points<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 315                                                mir: &Mir<'tcx>,
 316                                                source: MirSource) -> liveness::LocalSet {
 317     let mut set = liveness::LocalSet::new_empty(mir.local_decls.len());
 318     let result = liveness::liveness_of_locals(mir);
 319     liveness::dump_mir(tcx, "generator_liveness", source, mir, &result);
 320
 321     for (block, data) in mir.basic_blocks().iter_enumerated() {
 322         if let TerminatorKind::Yield { .. } = data.terminator().kind {
 323             set.union(&result.outs[block]);
 324         }
 325     }
 326
 327     // The generator argument is ignored
 328     set.remove(&self_arg());
 329
 330     set
 331 }
 332
 333 fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 334                             source: MirSource,
 335                             interior: GeneratorInterior<'tcx>,
 336                             mir: &mut Mir<'tcx>)
 337     -> (HashMap<Local, (Ty<'tcx>, usize)>, GeneratorLayout<'tcx>)
 338 {
 339     // Use a liveness analysis to compute locals which are live across a suspension point
 340     let live_locals = locals_live_across_suspend_points(tcx, mir, source);
 341
 342     // Erase regions from the types passed in from typeck so we can compare them with
 343     // MIR types
 344     let allowed = tcx.erase_regions(&interior.as_slice());
 345
 346     for (local, decl) in mir.local_decls.iter_enumerated() {
 347         // Ignore locals which are internal or not live
 348         if !live_locals.contains(&local) || decl.internal {
 349             continue;
 350         }
 351
 352         // Sanity check that typeck knows about the type of locals which are
 353         // live across a suspension point
 354         if !allowed.contains(&decl.ty) {
 355             span_bug!(mir.span,
 356                       "Broken MIR: generator contains type {} in MIR, \
 357                        but typeck only knows about {}",
 358                       decl.ty,
 359                       interior);
 360         }
 361     }
 362
 363     let upvar_len = mir.upvar_decls.len();
 364     let dummy_local = LocalDecl::new_internal(tcx.mk_nil(), mir.span);
 365
 366     // Gather live locals and their indices replacing values in mir.local_decls with a dummy
 367     // to avoid changing local indices
 368     let live_decls = live_locals.iter().map(|local| {
 369         let var = mem::replace(&mut mir.local_decls[local], dummy_local.clone());
 370         (local, var)
 371     });
 372
 373     // Create a map from local indices to generator struct indices.
 374     // These are offset by (upvar_len + 1) because of fields which comes before locals.
 375     // We also create a vector of the LocalDecls of these locals.
 376     let (remap, vars) = live_decls.enumerate().map(|(idx, (local, var))| {
 377         ((local, (var.ty, upvar_len + 1 + idx)), var)
 378     }).unzip();
 379
 380     let layout = GeneratorLayout {
 381         fields: vars
 382     };
 383
 384     (remap, layout)
 385 }
 386
 387 fn insert_entry_point<'tcx>(mir: &mut Mir<'tcx>,
 388                             block: BasicBlockData<'tcx>) {
 389     mir.basic_blocks_mut().raw.insert(0, block);
 390
 391     let blocks = mir.basic_blocks_mut().iter_mut();
 392
 393     for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
 394         *target = BasicBlock::new(target.index() + 1);
 395     }
 396 }
 397
 398 fn elaborate_generator_drops<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 399                                        def_id: DefId,
 400                                        mir: &mut Mir<'tcx>) {
 401     use util::elaborate_drops::{elaborate_drop, Unwind};
 402     use util::patch::MirPatch;
 403     use shim::DropShimElaborator;
 404
 405     // Note that `elaborate_drops` only drops the upvars of a generator, and
 406     // this is ok because `open_drop` can only be reached within that own
 407     // generator's resume function.
 408
 409     let param_env = tcx.param_env(def_id);
 410     let gen = self_arg();
 411
 412     for block in mir.basic_blocks().indices() {
 413         let (target, unwind, source_info) = match mir.basic_blocks()[block].terminator() {
 414             &Terminator {
 415                 source_info,
 416                 kind: TerminatorKind::Drop {
 417                     location: Lvalue::Local(local),
 418                     target,
 419                     unwind
 420                 }
 421             } if local == gen => (target, unwind, source_info),
 422             _ => continue,
 423         };
 424         let unwind = if let Some(unwind) = unwind {
 425             Unwind::To(unwind)
 426         } else {
 427             Unwind::InCleanup
 428         };
 429         let patch = {
 430             let mut elaborator = DropShimElaborator {
 431                 mir: &mir,
 432                 patch: MirPatch::new(mir),
 433                 tcx,
 434                 param_env
 435             };
 436             elaborate_drop(
 437                 &mut elaborator,
 438                 source_info,
 439                 &Lvalue::Local(gen),
 440                 (),
 441                 target,
 442                 unwind,
 443                 block
 444             );
 445             elaborator.patch
 446         };
 447         patch.apply(mir);
 448     }
 449 }
 450
 451 fn create_generator_drop_shim<'a, 'tcx>(
 452                 tcx: TyCtxt<'a, 'tcx, 'tcx>,
 453                 transform: &TransformVisitor<'a, 'tcx>,
 454                 def_id: DefId,
 455                 source: MirSource,
 456                 gen_ty: Ty<'tcx>,
 457                 mir: &Mir<'tcx>,
 458                 drop_clean: BasicBlock) -> Mir<'tcx> {
 459     let mut mir = mir.clone();
 460
 461     let source_info = SourceInfo {
 462         span: mir.span,
 463         scope: ARGUMENT_VISIBILITY_SCOPE,
 464     };
 465
 466     let return_block = BasicBlock::new(mir.basic_blocks().len());
 467     mir.basic_blocks_mut().push(BasicBlockData {
 468         statements: Vec::new(),
 469         terminator: Some(Terminator {
 470             source_info,
 471             kind: TerminatorKind::Return,
 472         }),
 473         is_cleanup: false,
 474     });
 475
 476     let mut cases: Vec<_> = transform.bb_targets.iter().filter_map(|(&(_, u), &s)| {
 477         u.map(|d| (s, d))
 478     }).collect();
 479
 480     cases.insert(0, (0, drop_clean));
 481
 482     // The poisoned state 1 falls through to the default case which is just to return
 483
 484     let switch = TerminatorKind::SwitchInt {
 485         discr: Operand::Consume(transform.make_field(transform.state_field, tcx.types.u32)),
 486         switch_ty: tcx.types.u32,
 487         values: Cow::from(cases.iter().map(|&(i, _)| {
 488                 ConstInt::U32(i)
 489             }).collect::<Vec<_>>()),
 490         targets: cases.iter().map(|&(_, d)| d).chain(once(return_block)).collect(),
 491     };
 492
 493     insert_entry_point(&mut mir, BasicBlockData {
 494         statements: Vec::new(),
 495         terminator: Some(Terminator {
 496             source_info,
 497             kind: switch,
 498         }),
 499         is_cleanup: false,
 500     });
 501
 502     for block in mir.basic_blocks_mut() {
 503         let kind = &mut block.terminator_mut().kind;
 504         if let TerminatorKind::GeneratorDrop = *kind {
 505             *kind = TerminatorKind::Return;
 506         }
 507     }
 508
 509     // Replace the return variable
 510     let source_info = SourceInfo {
 511         span: mir.span,
 512         scope: ARGUMENT_VISIBILITY_SCOPE,
 513     };
 514
 515     mir.return_ty = tcx.mk_nil();
 516     mir.local_decls[RETURN_POINTER] = LocalDecl {
 517         mutability: Mutability::Mut,
 518         ty: tcx.mk_nil(),
 519         name: None,
 520         source_info,
 521         internal: false,
 522         is_user_variable: false,
 523     };
 524
 525     make_generator_state_argument_indirect(tcx, def_id, &mut mir);
 526
 527     // Change the generator argument from &mut to *mut
 528     mir.local_decls[self_arg()] = LocalDecl {
 529         mutability: Mutability::Mut,
 530         ty: tcx.mk_ptr(ty::TypeAndMut {
 531             ty: gen_ty,
 532             mutbl: hir::Mutability::MutMutable,
 533         }),
 534         name: None,
 535         source_info,
 536         internal: false,
 537         is_user_variable: false,
 538     };
 539
 540     no_landing_pads(tcx, &mut mir);
 541
 542     // Make sure we remove dead blocks to remove
 543     // unrelated code from the resume part of the function
 544     simplify::remove_dead_blocks(&mut mir);
 545
 546     dump_mir(tcx, None, "generator_drop", &0, source, &mut mir);
 547
 548     mir
 549 }
 550
 551 fn insert_panic_on_resume_after_return<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 552                                         mir: &mut Mir<'tcx>) {
 553     let assert_block = BasicBlock::new(mir.basic_blocks().len());
 554     let term = TerminatorKind::Assert {
 555         cond: Operand::Constant(box Constant {
 556             span: mir.span,
 557             ty: tcx.types.bool,
 558             literal: Literal::Value {
 559                 value: tcx.mk_const(ty::Const {
 560                     val: ConstVal::Bool(false),
 561                     ty: tcx.types.bool
 562                 }),
 563             },
 564         }),
 565         expected: true,
 566         msg: AssertMessage::GeneratorResumedAfterReturn,
 567         target: assert_block,
 568         cleanup: None,
 569     };
 570
 571     let source_info = SourceInfo {
 572         span: mir.span,
 573         scope: ARGUMENT_VISIBILITY_SCOPE,
 574     };
 575
 576     mir.basic_blocks_mut().push(BasicBlockData {
 577         statements: Vec::new(),
 578         terminator: Some(Terminator {
 579             source_info,
 580             kind: term,
 581         }),
 582         is_cleanup: false,
 583     });
 584 }
 585
 586 fn create_generator_resume_function<'a, 'tcx>(
 587         tcx: TyCtxt<'a, 'tcx, 'tcx>,
 588         mut transform: TransformVisitor<'a, 'tcx>,
 589         def_id: DefId,
 590         source: MirSource,
 591         mir: &mut Mir<'tcx>) {
 592     // Poison the generator when it unwinds
 593     for block in mir.basic_blocks_mut() {
 594         let source_info = block.terminator().source_info;
 595         if let &TerminatorKind::Resume = &block.terminator().kind {
 596             block.statements.push(transform.set_state(1, source_info));
 597         }
 598     }
 599
 600     let source_info = SourceInfo {
 601         span: mir.span,
 602         scope: ARGUMENT_VISIBILITY_SCOPE,
 603     };
 604
 605     let poisoned_block = BasicBlock::new(mir.basic_blocks().len());
 606
 607     let term = TerminatorKind::Assert {
 608         cond: Operand::Constant(box Constant {
 609             span: mir.span,
 610             ty: tcx.types.bool,
 611             literal: Literal::Value {
 612                 value: tcx.mk_const(ty::Const {
 613                     val: ConstVal::Bool(false),
 614                     ty: tcx.types.bool
 615                 }),
 616             },
 617         }),
 618         expected: true,
 619         msg: AssertMessage::GeneratorResumedAfterPanic,
 620         target: transform.return_block,
 621         cleanup: None,
 622     };
 623
 624     mir.basic_blocks_mut().push(BasicBlockData {
 625         statements: Vec::new(),
 626         terminator: Some(Terminator {
 627             source_info,
 628             kind: term,
 629         }),
 630         is_cleanup: false,
 631     });
 632
 633     transform.bb_targets.insert((poisoned_block, None), 1);
 634
 635     let switch = TerminatorKind::SwitchInt {
 636         discr: Operand::Consume(transform.make_field(transform.state_field, tcx.types.u32)),
 637         switch_ty: tcx.types.u32,
 638         values: Cow::from(transform.bb_targets.values().map(|&i| {
 639                 ConstInt::U32(i)
 640             }).collect::<Vec<_>>()),
 641         targets: transform.bb_targets.keys()
 642             .map(|&(k, _)| k)
 643             .chain(once(transform.return_block))
 644             .collect(),
 645     };
 646
 647     insert_entry_point(mir, BasicBlockData {
 648         statements: Vec::new(),
 649         terminator: Some(Terminator {
 650             source_info,
 651             kind: switch,
 652         }),
 653         is_cleanup: false,
 654     });
 655
 656     make_generator_state_argument_indirect(tcx, def_id, mir);
 657
 658     no_landing_pads(tcx, mir);
 659
 660     // Make sure we remove dead blocks to remove
 661     // unrelated code from the drop part of the function
 662     simplify::remove_dead_blocks(mir);
 663
 664     dump_mir(tcx, None, "generator_resume", &0, source, mir);
 665 }
 666
 667 fn insert_clean_drop<'a, 'tcx>(mir: &mut Mir<'tcx>) -> BasicBlock {
 668     let source_info = SourceInfo {
 669         span: mir.span,
 670         scope: ARGUMENT_VISIBILITY_SCOPE,
 671     };
 672
 673     let return_block = BasicBlock::new(mir.basic_blocks().len());
 674     mir.basic_blocks_mut().push(BasicBlockData {
 675         statements: Vec::new(),
 676         terminator: Some(Terminator {
 677             source_info,
 678             kind: TerminatorKind::Return,
 679         }),
 680         is_cleanup: false,
 681     });
 682
 683     // Create a block to destroy an unresumed generators. This can only destroy upvars.
 684     let drop_clean = BasicBlock::new(mir.basic_blocks().len());
 685     let term = TerminatorKind::Drop {
 686         location: Lvalue::Local(self_arg()),
 687         target: return_block,
 688         unwind: None,
 689     };
 690     mir.basic_blocks_mut().push(BasicBlockData {
 691         statements: Vec::new(),
 692         terminator: Some(Terminator {
 693             source_info,
 694             kind: term,
 695         }),
 696         is_cleanup: false,
 697     });
 698
 699     drop_clean
 700 }
 701
 702 impl MirPass for StateTransform {
 703     fn run_pass<'a, 'tcx>(&self,
 704                     tcx: TyCtxt<'a, 'tcx, 'tcx>,
 705                     source: MirSource,
 706                     mir: &mut Mir<'tcx>) {
 707         let yield_ty = if let Some(yield_ty) = mir.yield_ty {
 708             yield_ty
 709         } else {
 710             // This only applies to generators
 711             return
 712         };
 713
 714         assert!(mir.generator_drop.is_none());
 715
 716         let node_id = source.item_id();
 717         let def_id = tcx.hir.local_def_id(source.item_id());
 718         let hir_id = tcx.hir.node_to_hir_id(node_id);
 719
 720         // Get the interior types which typeck computed
 721         let interior = *tcx.typeck_tables_of(def_id).generator_interiors().get(hir_id).unwrap();
 722
 723         // The first argument is the generator type passed by value
 724         let gen_ty = mir.local_decls.raw[1].ty;
 725
 726         // Compute GeneratorState<yield_ty, return_ty>
 727         let state_did = tcx.lang_items().gen_state().unwrap();
 728         let state_adt_ref = tcx.adt_def(state_did);
 729         let state_substs = tcx.mk_substs([Kind::from(yield_ty),
 730             Kind::from(mir.return_ty)].iter());
 731         let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
 732
 733         // We rename RETURN_POINTER which has type mir.return_ty to new_ret_local
 734         // RETURN_POINTER then is a fresh unused local with type ret_ty.
 735         let new_ret_local = replace_result_variable(ret_ty, mir);
 736
 737         // Extract locals which are live across suspension point into `layout`
 738         // `remap` gives a mapping from local indices onto generator struct indices
 739         let (remap, layout) = compute_layout(tcx, source, interior, mir);
 740
 741         let state_field = mir.upvar_decls.len();
 742
 743         let mut bb_targets = HashMap::new();
 744
 745         // If we jump to the entry point, we should go to the initial 0 generator state.
 746         // FIXME: Could this result in the need for destruction for state 0?
 747         bb_targets.insert((BasicBlock::new(0), None), 0);
 748
 749         // Run the transformation which converts Lvalues from Local to generator struct
 750         // accesses for locals in `remap`.
 751         // It also rewrites `return x` and `yield y` as writing a new generator state and returning
 752         // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
 753         let mut transform = TransformVisitor {
 754             tcx,
 755             state_adt_ref,
 756             state_substs,
 757             remap,
 758             bb_target_count: 2,
 759             bb_targets,
 760             new_ret_local,
 761             state_field,
 762
 763             // For returns we will resume execution at the next added basic block.
 764             // This happens in `insert_panic_on_resume_after_return`
 765             return_block: BasicBlock::new(mir.basic_blocks().len()),
 766         };
 767         transform.visit_mir(mir);
 768
 769         // Update our MIR struct to reflect the changed we've made
 770         mir.return_ty = ret_ty;
 771         mir.yield_ty = None;
 772         mir.arg_count = 1;
 773         mir.spread_arg = None;
 774         mir.generator_layout = Some(layout);
 775
 776         // Panic if we resumed after returning
 777         insert_panic_on_resume_after_return(tcx, mir);
 778
 779         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
 780         // the unresumed (0) state.
 781         // This is expanded to a drop ladder in `elaborate_generator_drops`.
 782         let drop_clean = insert_clean_drop(mir);
 783
 784         dump_mir(tcx, None, "generator_pre-elab", &0, source, mir);
 785
 786         // Expand `drop(generator_struct)` to a drop ladder which destroys upvars.
 787         // If any upvars are moved out of, drop elaboration will handle upvar destruction.
 788         // However we need to also elaborate the code generated by `insert_clean_drop`.
 789         elaborate_generator_drops(tcx, def_id, mir);
 790
 791         dump_mir(tcx, None, "generator_post-transform", &0, source, mir);
 792
 793         // Create a copy of our MIR and use it to create the drop shim for the generator
 794         let drop_shim = create_generator_drop_shim(tcx,
 795             &transform,
 796             def_id,
 797             source,
 798             gen_ty,
 799             &mir,
 800             drop_clean);
 801
 802         mir.generator_drop = Some(box drop_shim);
 803
 804         // Create the Generator::resume function
 805         create_generator_resume_function(tcx, transform, def_id, source, mir);
 806     }
 807 }