]> git.lizzy.rs Git - rust.git/commitdiff
Analyse storage liveness and preserve it during generator transformation
authorJohn Kåre Alsaker <john.kare.alsaker@gmail.com>
Sun, 10 Sep 2017 20:34:56 +0000 (22:34 +0200)
committerJohn Kåre Alsaker <john.kare.alsaker@gmail.com>
Wed, 13 Sep 2017 21:28:10 +0000 (23:28 +0200)
src/librustc_data_structures/bitslice.rs
src/librustc_data_structures/indexed_set.rs
src/librustc_mir/dataflow/impls/mod.rs
src/librustc_mir/dataflow/impls/storage_liveness.rs [new file with mode: 0644]
src/librustc_mir/dataflow/mod.rs
src/librustc_mir/transform/generator.rs
src/test/run-pass/generator/match-bindings.rs [new file with mode: 0644]

index f74af6ee1632e793f8f6fef0922f9262443e2eb1..7665bfd5b1117316c0a758338ad103adca34e2b5 100644 (file)
@@ -132,6 +132,11 @@ pub trait BitwiseOperator {
     fn join(&self, pred1: usize, pred2: usize) -> usize;
 }
 
+pub struct Intersect;
+impl BitwiseOperator for Intersect {
+    #[inline]
+    fn join(&self, a: usize, b: usize) -> usize { a & b }
+}
 pub struct Union;
 impl BitwiseOperator for Union {
     #[inline]
index 47fa21e3bf0b2bf6a1e47f4d6c78b94c1aca9246..c790463e47adb4f2ce3e897c2229e22dc9f7d88f 100644 (file)
@@ -15,7 +15,7 @@
 use std::ops::{Deref, DerefMut, Range};
 use std::slice;
 use bitslice::{BitSlice, Word};
-use bitslice::{bitwise, Union, Subtract};
+use bitslice::{bitwise, Union, Subtract, Intersect};
 use indexed_vec::Idx;
 
 /// Represents a set (or packed family of sets), of some element type
@@ -164,6 +164,10 @@ pub fn subtract(&mut self, other: &IdxSet<T>) -> bool {
         bitwise(self.words_mut(), other.words(), &Subtract)
     }
 
+    pub fn intersect(&mut self, other: &IdxSet<T>) -> bool {
+        bitwise(self.words_mut(), other.words(), &Intersect)
+    }
+
     pub fn iter(&self) -> Iter<T> {
         Iter {
             cur: None,
index a4421b216c311b4d50134c8e4885123d708216e7..19a595622b92edf62b9a4394407a4e5dffa4a1d9 100644 (file)
 use super::drop_flag_effects_for_location;
 use super::on_lookup_result_bits;
 
+mod storage_liveness;
+
+pub use self::storage_liveness::*;
+
 #[allow(dead_code)]
 pub(super) mod borrows;
 
diff --git a/src/librustc_mir/dataflow/impls/storage_liveness.rs b/src/librustc_mir/dataflow/impls/storage_liveness.rs
new file mode 100644 (file)
index 0000000..98615c6
--- /dev/null
@@ -0,0 +1,82 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+pub use super::*;
+
+use rustc::mir::*;
+use dataflow::BitDenotation;
+
+#[derive(Copy, Clone)]
+pub struct MaybeStorageLive<'a, 'tcx: 'a> {
+    mir: &'a Mir<'tcx>,
+}
+
+impl<'a, 'tcx: 'a> MaybeStorageLive<'a, 'tcx> {
+    pub fn new(mir: &'a Mir<'tcx>)
+               -> Self {
+        MaybeStorageLive { mir: mir }
+    }
+
+    pub fn mir(&self) -> &Mir<'tcx> {
+        self.mir
+    }
+}
+
+impl<'a, 'tcx> BitDenotation for MaybeStorageLive<'a, 'tcx> {
+    type Idx = Local;
+    fn name() -> &'static str { "maybe_storage_live" }
+    fn bits_per_block(&self) -> usize {
+        self.mir.local_decls.len()
+    }
+
+    fn start_block_effect(&self, _sets: &mut BlockSets<Local>) {
+        // Nothing is live on function entry
+    }
+
+    fn statement_effect(&self,
+                        sets: &mut BlockSets<Local>,
+                        loc: Location) {
+        let stmt = &self.mir[loc.block].statements[loc.statement_index];
+
+        match stmt.kind {
+            StatementKind::StorageLive(l) => sets.gen(&l),
+            StatementKind::StorageDead(l) => sets.kill(&l),
+            _ => (),
+        }
+    }
+
+    fn terminator_effect(&self,
+                         _sets: &mut BlockSets<Local>,
+                         _loc: Location) {
+        // Terminators have no effect
+    }
+
+    fn propagate_call_return(&self,
+                             _in_out: &mut IdxSet<Local>,
+                             _call_bb: mir::BasicBlock,
+                             _dest_bb: mir::BasicBlock,
+                             _dest_lval: &mir::Lvalue) {
+        // Nothing to do when a call returns successfully
+    }
+}
+
+impl<'a, 'tcx> BitwiseOperator for MaybeStorageLive<'a, 'tcx> {
+    #[inline]
+    fn join(&self, pred1: usize, pred2: usize) -> usize {
+        pred1 | pred2 // "maybe" means we union effects of both preds
+    }
+}
+
+impl<'a, 'tcx> DataflowOperator for MaybeStorageLive<'a, 'tcx> {
+    #[inline]
+    fn bottom_value() -> bool {
+        false // bottom = dead
+    }
+}
index 9f9909a8f57a99f9b1d4cf74e5319d0432fcb254..9fa5691d647b70015204c5c98d93cc11ee61c4c9 100644 (file)
@@ -24,6 +24,7 @@
 use std::path::PathBuf;
 use std::usize;
 
+pub use self::impls::{MaybeStorageLive};
 pub use self::impls::{MaybeInitializedLvals, MaybeUninitializedLvals};
 pub use self::impls::{DefinitelyInitializedLvals};
 pub use self::impls::borrows::{Borrows, BorrowData, BorrowIndex};
@@ -351,6 +352,29 @@ fn apply_local_effect(&mut self,
                           flow_state: &mut Self::FlowState);
 }
 
+pub fn state_for_location<T: BitDenotation>(loc: Location,
+                                            analysis: &T,
+                                            result: &DataflowResults<T>)
+    -> IdxSetBuf<T::Idx> {
+    let mut entry = result.sets().on_entry_set_for(loc.block.index()).to_owned();
+
+    {
+        let mut sets = BlockSets {
+            on_entry: &mut entry.clone(),
+            kill_set: &mut entry.clone(),
+            gen_set: &mut entry,
+        };
+
+        for stmt in 0..loc.statement_index {
+            let mut stmt_loc = loc;
+            stmt_loc.statement_index = stmt;
+            analysis.statement_effect(&mut sets, stmt_loc);
+        }
+    }
+
+    entry
+}
+
 pub struct DataflowAnalysis<'a, 'tcx: 'a, O> where O: BitDenotation
 {
     flow_state: DataflowState<O>,
index a52656becd74534190b114501df7e4bff8d6be39..3be21b72731e96b5e753c674285e7f7394e15b6f 100644 (file)
@@ -38,7 +38,8 @@
 //! This pass computes the meaning of the state field and the MIR locals which are live
 //! across a suspension point. There are however two hardcoded generator states:
 //!     0 - Generator have not been resumed yet
-//!     1 - Generator has been poisoned
+//!     1 - Generator has returned / is completed
+//!     2 - Generator has been poisoned
 //!
 //! It also rewrites `return x` and `yield y` as setting a new generator state and returning
 //! GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
 //! the action to take.
 //!
 //! One of them is the implementation of Generator::resume.
-//! For generators which have already returned it panics.
 //! For generators with state 0 (unresumed) it starts the execution of the generator.
-//! For generators with state 1 (poisoned) it panics.
+//! For generators with state 1 (returned) and state 2 (poisoned) it panics.
 //! Otherwise it continues the execution from the last suspension point.
 //!
 //! The other function is the drop glue for the generator.
-//! For generators which have already returned it does nothing.
 //! For generators with state 0 (unresumed) it drops the upvars of the generator.
-//! For generators with state 1 (poisoned) it does nothing.
+//! For generators with state 1 (returned) and state 2 (poisoned) it does nothing.
 //! Otherwise it drops all the values in scope at the last suspension point.
 
 use rustc::hir;
 use rustc::middle::const_val::ConstVal;
 use rustc::mir::*;
 use rustc::mir::transform::{MirPass, MirSource};
-use rustc::mir::visit::{LvalueContext, MutVisitor};
+use rustc::mir::visit::{LvalueContext, Visitor, MutVisitor};
 use rustc::ty::{self, TyCtxt, AdtDef, Ty, GeneratorInterior};
 use rustc::ty::subst::{Kind, Substs};
 use util::dump_mir;
 use util::liveness;
 use rustc_const_math::ConstInt;
 use rustc_data_structures::indexed_vec::Idx;
+use rustc_data_structures::indexed_set::IdxSetBuf;
 use std::collections::HashMap;
 use std::borrow::Cow;
 use std::iter::once;
 use std::mem;
 use transform::simplify;
 use transform::no_landing_pads::no_landing_pads;
+use dataflow::{self, MaybeStorageLive, state_for_location};
 
 pub struct StateTransform;
 
@@ -126,6 +127,14 @@ fn self_arg() -> Local {
     Local::new(1)
 }
 
+struct SuspensionPoint {
+    state: u32,
+    resume: BasicBlock,
+    drop: Option<BasicBlock>,
+    storage_liveness: liveness::LocalSet,
+    storage_live: Option<BasicBlock>,
+}
+
 struct TransformVisitor<'a, 'tcx: 'a> {
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
     state_adt_ref: &'tcx AdtDef,
@@ -137,18 +146,16 @@ struct TransformVisitor<'a, 'tcx: 'a> {
     // Mapping from Local to (type of local, generator struct index)
     remap: HashMap<Local, (Ty<'tcx>, usize)>,
 
-    // The number of generator states. 0 is unresumed, 1 is poisoned. So this is initialized to 2
-    bb_target_count: u32,
+    mir_local_count: usize,
+
+    // A map from a suspension point in a block to the locals which have live storage at that point
+    storage_liveness: HashMap<BasicBlock, liveness::LocalSet>,
 
-    // Map from a (which block to resume execution at, which block to use to drop the generator)
-    // to a generator state
-    bb_targets: HashMap<(BasicBlock, Option<BasicBlock>), u32>,
+    // A list of suspension points, generated during the transform
+    suspension_points: Vec<SuspensionPoint>,
 
     // The original RETURN_POINTER local
     new_ret_local: Local,
-
-    // The block to resume execution when for Return
-    return_block: BasicBlock,
 }
 
 impl<'a, 'tcx> TransformVisitor<'a, 'tcx> {
@@ -225,28 +232,46 @@ fn visit_basic_block_data(&mut self,
 
         let ret_val = match data.terminator().kind {
             TerminatorKind::Return => Some((1,
-                self.return_block,
+                None,
                 Operand::Consume(Lvalue::Local(self.new_ret_local)),
                 None)),
             TerminatorKind::Yield { ref value, resume, drop } => Some((0,
-                resume,
+                Some(resume),
                 value.clone(),
                 drop)),
             _ => None
         };
 
         if let Some((state_idx, resume, v, drop)) = ret_val {
-            let bb_idx = {
-                let bb_targets = &mut self.bb_targets;
-                let bb_target = &mut self.bb_target_count;
-                *bb_targets.entry((resume, drop)).or_insert_with(|| {
-                    let target = *bb_target;
-                    *bb_target = target.checked_add(1).unwrap();
-                    target
-                })
-            };
             let source_info = data.terminator().source_info;
-            data.statements.push(self.set_state(bb_idx, source_info));
+            let state = if let Some(resume) = resume { // Yield
+                let state = 3 + self.suspension_points.len() as u32;
+
+                let liveness = self.storage_liveness.get(&block).unwrap();
+
+                for i in 0..(self.mir_local_count) {
+                    let l = Local::new(i);
+                    if liveness.contains(&l) && !self.remap.contains_key(&l) {
+                        data.statements.push(Statement {
+                            source_info,
+                            kind: StatementKind::StorageDead(l),
+                        });
+                    }
+                }
+
+                self.suspension_points.push(SuspensionPoint {
+                    state,
+                    resume,
+                    drop,
+                    storage_liveness: liveness.clone(),
+                    storage_live: None,
+                });
+
+                state
+            } else { // Return
+                 1 // state for returned
+            };
+            data.statements.push(self.set_state(state, source_info));
             data.statements.push(Statement {
                 source_info,
                 kind: StatementKind::Assign(Lvalue::Local(RETURN_POINTER),
@@ -286,16 +311,11 @@ fn make_generator_state_argument_indirect<'a, 'tcx>(
 
 fn replace_result_variable<'tcx>(ret_ty: Ty<'tcx>,
                             mir: &mut Mir<'tcx>) -> Local {
-    let source_info = SourceInfo {
-        span: mir.span,
-        scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
-
     let new_ret = LocalDecl {
         mutability: Mutability::Mut,
         ty: ret_ty,
         name: None,
-        source_info,
+        source_info: source_info(mir),
         internal: false,
         is_user_variable: false,
     };
@@ -311,33 +331,82 @@ fn replace_result_variable<'tcx>(ret_ty: Ty<'tcx>,
     new_ret_local
 }
 
+struct StorageIgnored(liveness::LocalSet);
+
+impl<'tcx> Visitor<'tcx> for StorageIgnored {
+    fn visit_statement(&mut self,
+                       _block: BasicBlock,
+                       statement: &Statement<'tcx>,
+                       _location: Location) {
+        match statement.kind {
+            StatementKind::StorageLive(l) |
+            StatementKind::StorageDead(l) => { self.0.remove(&l); }
+            _ => (),
+        }
+    }
+}
+
 fn locals_live_across_suspend_points<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                                mir: &Mir<'tcx>,
-                                               source: MirSource) -> liveness::LocalSet {
+                                               source: MirSource) ->
+                                               (liveness::LocalSet,
+                                                HashMap<BasicBlock, liveness::LocalSet>) {
+    let dead_unwinds = IdxSetBuf::new_empty(mir.basic_blocks().len());
+    let node_id = source.item_id();
+    let analysis = MaybeStorageLive::new(mir);
+    let storage_live =
+        dataflow::do_dataflow(tcx, mir, node_id, &[], &dead_unwinds, analysis,
+                              |bd, p| &bd.mir().local_decls[p]);
+
+    let mut ignored = StorageIgnored(IdxSetBuf::new_filled(mir.basic_blocks().len()));
+    ignored.visit_mir(mir);
+
     let mut set = liveness::LocalSet::new_empty(mir.local_decls.len());
     let result = liveness::liveness_of_locals(mir);
     liveness::dump_mir(tcx, "generator_liveness", source, mir, &result);
 
+    let mut storage_liveness_map = HashMap::new();
+
     for (block, data) in mir.basic_blocks().iter_enumerated() {
         if let TerminatorKind::Yield { .. } = data.terminator().kind {
-            set.union(&result.outs[block]);
+            let loc = Location {
+                block: block,
+                statement_index: data.statements.len(),
+            };
+
+            let mut storage_liveness = state_for_location(loc, &analysis, &storage_live);
+
+            storage_liveness_map.insert(block, storage_liveness.clone());
+
+            // Mark locals without storage statements as always live
+            storage_liveness.union(&ignored.0);
+
+            // Locals live are live at this point only if they are used across suspension points
+            // and their storage is live
+            storage_liveness.intersect(&result.outs[block]);
+
+            // Add the locals life at this suspension point to the set of locals which live across
+            // any suspension points
+            set.union(&storage_liveness);
         }
     }
 
     // The generator argument is ignored
     set.remove(&self_arg());
 
-    set
+    (set, storage_liveness_map)
 }
 
 fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                             source: MirSource,
                             interior: GeneratorInterior<'tcx>,
                             mir: &mut Mir<'tcx>)
-    -> (HashMap<Local, (Ty<'tcx>, usize)>, GeneratorLayout<'tcx>)
+    -> (HashMap<Local, (Ty<'tcx>, usize)>,
+        GeneratorLayout<'tcx>,
+        HashMap<BasicBlock, liveness::LocalSet>)
 {
     // Use a liveness analysis to compute locals which are live across a suspension point
-    let live_locals = locals_live_across_suspend_points(tcx, mir, source);
+    let (live_locals, storage_liveness) = locals_live_across_suspend_points(tcx, mir, source);
 
     // Erase regions from the types passed in from typeck so we can compare them with
     // MIR types
@@ -381,12 +450,31 @@ fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         fields: vars
     };
 
-    (remap, layout)
+    (remap, layout, storage_liveness)
 }
 
-fn insert_entry_point<'tcx>(mir: &mut Mir<'tcx>,
-                            block: BasicBlockData<'tcx>) {
-    mir.basic_blocks_mut().raw.insert(0, block);
+fn insert_switch<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                           mir: &mut Mir<'tcx>,
+                           cases: Vec<(u32, BasicBlock)>,
+                           transform: &TransformVisitor<'a, 'tcx>) {
+    let return_block = insert_return_block(mir);
+
+    let switch = TerminatorKind::SwitchInt {
+        discr: Operand::Consume(transform.make_field(transform.state_field, tcx.types.u32)),
+        switch_ty: tcx.types.u32,
+        values: Cow::from(cases.iter().map(|&(i, _)| ConstInt::U32(i)).collect::<Vec<_>>()),
+        targets: cases.iter().map(|&(_, d)| d).chain(once(return_block)).collect(),
+    };
+
+    let source_info = source_info(mir);
+    mir.basic_blocks_mut().raw.insert(0, BasicBlockData {
+        statements: Vec::new(),
+        terminator: Some(Terminator {
+            source_info,
+            kind: switch,
+        }),
+        is_cleanup: false,
+    });
 
     let blocks = mir.basic_blocks_mut().iter_mut();
 
@@ -458,46 +546,29 @@ fn create_generator_drop_shim<'a, 'tcx>(
                 drop_clean: BasicBlock) -> Mir<'tcx> {
     let mut mir = mir.clone();
 
-    let source_info = SourceInfo {
-        span: mir.span,
-        scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
+    let source_info = source_info(&mir);
 
-    let return_block = BasicBlock::new(mir.basic_blocks().len());
-    mir.basic_blocks_mut().push(BasicBlockData {
-        statements: Vec::new(),
-        terminator: Some(Terminator {
-            source_info,
-            kind: TerminatorKind::Return,
-        }),
-        is_cleanup: false,
-    });
-
-    let mut cases: Vec<_> = transform.bb_targets.iter().filter_map(|(&(_, u), &s)| {
-        u.map(|d| (s, d))
+    let mut cases: Vec<_> = transform.suspension_points.iter().filter_map(|point| {
+        point.drop.map(|drop| {
+            // Make the point's storage live block goto the drop block
+            let block = point.storage_live.unwrap();
+            let term = Terminator {
+                source_info,
+                kind: TerminatorKind::Goto {
+                    target: drop,
+                },
+            };
+            mir.basic_blocks_mut()[block].terminator = Some(term);
+            (point.state, block)
+        })
     }).collect();
 
     cases.insert(0, (0, drop_clean));
 
-    // The poisoned state 1 falls through to the default case which is just to return
-
-    let switch = TerminatorKind::SwitchInt {
-        discr: Operand::Consume(transform.make_field(transform.state_field, tcx.types.u32)),
-        switch_ty: tcx.types.u32,
-        values: Cow::from(cases.iter().map(|&(i, _)| {
-                ConstInt::U32(i)
-            }).collect::<Vec<_>>()),
-        targets: cases.iter().map(|&(_, d)| d).chain(once(return_block)).collect(),
-    };
+    // The returned state 1 and the  poisoned state 2 falls through to
+    // the default case which is just to return
 
-    insert_entry_point(&mut mir, BasicBlockData {
-        statements: Vec::new(),
-        terminator: Some(Terminator {
-            source_info,
-            kind: switch,
-        }),
-        is_cleanup: false,
-    });
+    insert_switch(tcx, &mut mir, cases, &transform);
 
     for block in mir.basic_blocks_mut() {
         let kind = &mut block.terminator_mut().kind;
@@ -507,11 +578,6 @@ fn create_generator_drop_shim<'a, 'tcx>(
     }
 
     // Replace the return variable
-    let source_info = SourceInfo {
-        span: mir.span,
-        scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
-
     mir.return_ty = tcx.mk_nil();
     mir.local_decls[RETURN_POINTER] = LocalDecl {
         mutability: Mutability::Mut,
@@ -548,8 +614,23 @@ fn create_generator_drop_shim<'a, 'tcx>(
     mir
 }
 
-fn insert_panic_on_resume_after_return<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
-                                        mir: &mut Mir<'tcx>) {
+fn insert_return_block<'tcx>(mir: &mut Mir<'tcx>) -> BasicBlock {
+    let return_block = BasicBlock::new(mir.basic_blocks().len());
+    let source_info = source_info(mir);
+    mir.basic_blocks_mut().push(BasicBlockData {
+        statements: Vec::new(),
+        terminator: Some(Terminator {
+            source_info,
+            kind: TerminatorKind::Return,
+        }),
+        is_cleanup: false,
+    });
+    return_block
+}
+
+fn insert_panic_block<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                                mir: &mut Mir<'tcx>,
+                                message: AssertMessage<'tcx>) -> BasicBlock {
     let assert_block = BasicBlock::new(mir.basic_blocks().len());
     let term = TerminatorKind::Assert {
         cond: Operand::Constant(box Constant {
@@ -563,16 +644,12 @@ fn insert_panic_on_resume_after_return<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             },
         }),
         expected: true,
-        msg: AssertMessage::GeneratorResumedAfterReturn,
+        msg: message,
         target: assert_block,
         cleanup: None,
     };
 
-    let source_info = SourceInfo {
-        span: mir.span,
-        scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
-
+    let source_info = source_info(mir);
     mir.basic_blocks_mut().push(BasicBlockData {
         statements: Vec::new(),
         terminator: Some(Terminator {
@@ -581,11 +658,13 @@ fn insert_panic_on_resume_after_return<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         }),
         is_cleanup: false,
     });
+
+    assert_block
 }
 
 fn create_generator_resume_function<'a, 'tcx>(
         tcx: TyCtxt<'a, 'tcx, 'tcx>,
-        mut transform: TransformVisitor<'a, 'tcx>,
+        transform: TransformVisitor<'a, 'tcx>,
         def_id: DefId,
         source: MirSource,
         mir: &mut Mir<'tcx>) {
@@ -597,61 +676,27 @@ fn create_generator_resume_function<'a, 'tcx>(
         }
     }
 
-    let source_info = SourceInfo {
-        span: mir.span,
-        scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
-
-    let poisoned_block = BasicBlock::new(mir.basic_blocks().len());
-
-    let term = TerminatorKind::Assert {
-        cond: Operand::Constant(box Constant {
-            span: mir.span,
-            ty: tcx.types.bool,
-            literal: Literal::Value {
-                value: tcx.mk_const(ty::Const {
-                    val: ConstVal::Bool(false),
-                    ty: tcx.types.bool
-                }),
+    let mut cases: Vec<_> = transform.suspension_points.iter().map(|point| {
+        // Make the point's storage live block goto the resume block
+        let block = point.storage_live.unwrap();
+        let term = Terminator {
+            source_info: source_info(mir),
+            kind: TerminatorKind::Goto {
+                target: point.resume,
             },
-        }),
-        expected: true,
-        msg: AssertMessage::GeneratorResumedAfterPanic,
-        target: transform.return_block,
-        cleanup: None,
-    };
-
-    mir.basic_blocks_mut().push(BasicBlockData {
-        statements: Vec::new(),
-        terminator: Some(Terminator {
-            source_info,
-            kind: term,
-        }),
-        is_cleanup: false,
-    });
+        };
+        mir.basic_blocks_mut()[block].terminator = Some(term);
+        (point.state, block)
+    }).collect();
 
-    transform.bb_targets.insert((poisoned_block, None), 1);
+    // Jump to the entry point on the 0 state
+    cases.insert(0, (0, BasicBlock::new(0)));
+    // Panic when resumed on the returned (1) state
+    cases.insert(1, (1, insert_panic_block(tcx, mir, AssertMessage::GeneratorResumedAfterReturn)));
+    // Panic when resumed on the poisoned (2) state
+    cases.insert(2, (2, insert_panic_block(tcx, mir, AssertMessage::GeneratorResumedAfterPanic)));
 
-    let switch = TerminatorKind::SwitchInt {
-        discr: Operand::Consume(transform.make_field(transform.state_field, tcx.types.u32)),
-        switch_ty: tcx.types.u32,
-        values: Cow::from(transform.bb_targets.values().map(|&i| {
-                ConstInt::U32(i)
-            }).collect::<Vec<_>>()),
-        targets: transform.bb_targets.keys()
-            .map(|&(k, _)| k)
-            .chain(once(transform.return_block))
-            .collect(),
-    };
-
-    insert_entry_point(mir, BasicBlockData {
-        statements: Vec::new(),
-        terminator: Some(Terminator {
-            source_info,
-            kind: switch,
-        }),
-        is_cleanup: false,
-    });
+    insert_switch(tcx, mir, cases, &transform);
 
     make_generator_state_argument_indirect(tcx, def_id, mir);
 
@@ -664,21 +709,15 @@ fn create_generator_resume_function<'a, 'tcx>(
     dump_mir(tcx, None, "generator_resume", &0, source, mir);
 }
 
-fn insert_clean_drop<'a, 'tcx>(mir: &mut Mir<'tcx>) -> BasicBlock {
-    let source_info = SourceInfo {
+fn source_info<'a, 'tcx>(mir: &Mir<'tcx>) -> SourceInfo {
+    SourceInfo {
         span: mir.span,
         scope: ARGUMENT_VISIBILITY_SCOPE,
-    };
+    }
+}
 
-    let return_block = BasicBlock::new(mir.basic_blocks().len());
-    mir.basic_blocks_mut().push(BasicBlockData {
-        statements: Vec::new(),
-        terminator: Some(Terminator {
-            source_info,
-            kind: TerminatorKind::Return,
-        }),
-        is_cleanup: false,
-    });
+fn insert_clean_drop<'a, 'tcx>(mir: &mut Mir<'tcx>) -> BasicBlock {
+    let return_block = insert_return_block(mir);
 
     // Create a block to destroy an unresumed generators. This can only destroy upvars.
     let drop_clean = BasicBlock::new(mir.basic_blocks().len());
@@ -687,6 +726,7 @@ fn insert_clean_drop<'a, 'tcx>(mir: &mut Mir<'tcx>) -> BasicBlock {
         target: return_block,
         unwind: None,
     };
+    let source_info = source_info(mir);
     mir.basic_blocks_mut().push(BasicBlockData {
         statements: Vec::new(),
         terminator: Some(Terminator {
@@ -736,16 +776,11 @@ fn run_pass<'a, 'tcx>(&self,
 
         // Extract locals which are live across suspension point into `layout`
         // `remap` gives a mapping from local indices onto generator struct indices
-        let (remap, layout) = compute_layout(tcx, source, interior, mir);
+        // `storage_liveness` tells us which locals have live storage at suspension points
+        let (remap, layout, storage_liveness) = compute_layout(tcx, source, interior, mir);
 
         let state_field = mir.upvar_decls.len();
 
-        let mut bb_targets = HashMap::new();
-
-        // If we jump to the entry point, we should go to the initial 0 generator state.
-        // FIXME: Could this result in the need for destruction for state 0?
-        bb_targets.insert((BasicBlock::new(0), None), 0);
-
         // Run the transformation which converts Lvalues from Local to generator struct
         // accesses for locals in `remap`.
         // It also rewrites `return x` and `yield y` as writing a new generator state and returning
@@ -755,14 +790,11 @@ fn run_pass<'a, 'tcx>(&self,
             state_adt_ref,
             state_substs,
             remap,
-            bb_target_count: 2,
-            bb_targets,
+            storage_liveness,
+            mir_local_count: mir.local_decls.len(),
+            suspension_points: Vec::new(),
             new_ret_local,
             state_field,
-
-            // For returns we will resume execution at the next added basic block.
-            // This happens in `insert_panic_on_resume_after_return`
-            return_block: BasicBlock::new(mir.basic_blocks().len()),
         };
         transform.visit_mir(mir);
 
@@ -773,9 +805,6 @@ fn run_pass<'a, 'tcx>(&self,
         mir.spread_arg = None;
         mir.generator_layout = Some(layout);
 
-        // Panic if we resumed after returning
-        insert_panic_on_resume_after_return(tcx, mir);
-
         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
         // the unresumed (0) state.
         // This is expanded to a drop ladder in `elaborate_generator_drops`.
@@ -790,6 +819,30 @@ fn run_pass<'a, 'tcx>(&self,
 
         dump_mir(tcx, None, "generator_post-transform", &0, source, mir);
 
+        // Create StorageLive instruction blocks for suspension points
+        for point in &mut transform.suspension_points {
+            point.storage_live = Some(BasicBlock::new(mir.basic_blocks().len()));
+            let source_info = source_info(mir);
+            let mut statements = Vec::new();
+            for i in 0..(transform.mir_local_count) {
+                let l = Local::new(i);
+                if point.storage_liveness.contains(&l) && !transform.remap.contains_key(&l) {
+                    statements.push(Statement {
+                        source_info,
+                        kind: StatementKind::StorageLive(l),
+                    });
+                }
+            }
+            mir.basic_blocks_mut().push(BasicBlockData {
+                statements,
+                terminator: Some(Terminator {
+                    source_info,
+                    kind: TerminatorKind::Unreachable,
+                }),
+                is_cleanup: false,
+            });
+        }
+
         // Create a copy of our MIR and use it to create the drop shim for the generator
         let drop_shim = create_generator_drop_shim(tcx,
             &transform,
diff --git a/src/test/run-pass/generator/match-bindings.rs b/src/test/run-pass/generator/match-bindings.rs
new file mode 100644 (file)
index 0000000..9c6b057
--- /dev/null
@@ -0,0 +1,30 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(generators)]
+
+enum Enum {
+    A(String),
+    B
+}
+
+fn main() {
+    || {
+        loop {
+            if let true = true {
+                match Enum::A(String::new()) {
+                    Enum::A(_var) => {}
+                    Enum::B => {}
+                }
+            }
+            yield;
+        }
+    };
+}
\ No newline at end of file