Use the correct state for poisoning a generator

Author: matthewjasper

src/librustc_mir/transform/generator.rs

@@ -26,7 +26,7 @@
 //!     }
 //!
 //! 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:
+//! across a suspension point. There are however three hardcoded generator states:
 //!     0 - Generator have not been resumed yet
 //!     1 - Generator has returned / is completed
 //!     2 - Generator has been poisoned
@@ -144,6 +144,13 @@ fn self_arg() -> Local {
     Local::new(1)
 }
 
+/// Generator have not been resumed yet
+const UNRESUMED: u32 = 0;
+/// Generator has returned / is completed
+const RETURNED: u32 = 1;
+/// Generator has been poisoned
+const POISONED: u32 = 2;
+
 struct SuspensionPoint {
     state: u32,
     resume: BasicBlock,
@@ -278,7 +285,7 @@ impl<'a, 'tcx> MutVisitor<'tcx> for TransformVisitor<'a, 'tcx> {
 
                 state
             } else { // Return
-                 1 // state for returned
+                RETURNED // state for returned
             };
             data.statements.push(self.set_state(state, source_info));
             data.terminator.as_mut().unwrap().kind = TerminatorKind::Return;
@@ -643,10 +650,10 @@ fn create_generator_drop_shim<'a, 'tcx>(
 
     let mut cases = create_cases(&mut mir, transform, |point| point.drop);
 
-    cases.insert(0, (0, drop_clean));
+    cases.insert(0, (UNRESUMED, drop_clean));
 
-    // The returned state (1) and the poisoned state (2) falls through to
-    // the default case which is just to return
+    // The returned state and the poisoned state fall through to the default
+    // case which is just to return
 
     insert_switch(tcx, &mut mir, cases, &transform, TerminatorKind::Return);
 
@@ -762,7 +769,7 @@ fn create_generator_resume_function<'a, 'tcx>(
     for block in mir.basic_blocks_mut() {
         let source_info = block.terminator().source_info;
         if let &TerminatorKind::Resume = &block.terminator().kind {
-            block.statements.push(transform.set_state(1, source_info));
+            block.statements.push(transform.set_state(POISONED, source_info));
         }
     }
 
@@ -773,12 +780,12 @@ fn create_generator_resume_function<'a, 'tcx>(
         GeneratorResumedAfterReturn,
     };
 
-    // 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, GeneratorResumedAfterReturn)));
-    // Panic when resumed on the poisoned (2) state
-    cases.insert(2, (2, insert_panic_block(tcx, mir, GeneratorResumedAfterPanic)));
+    // Jump to the entry point on the unresumed
+    cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
+    // Panic when resumed on the returned state
+    cases.insert(1, (RETURNED, insert_panic_block(tcx, mir, GeneratorResumedAfterReturn)));
+    // Panic when resumed on the poisoned state
+    cases.insert(2, (POISONED, insert_panic_block(tcx, mir, GeneratorResumedAfterPanic)));
 
     insert_switch(tcx, mir, cases, &transform, TerminatorKind::Unreachable);
 
@@ -942,7 +949,7 @@ impl MirPass for StateTransform {
         mir.generator_layout = Some(layout);
 
         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
-        // the unresumed (0) state.
+        // the unresumed state.
         // This is expanded to a drop ladder in `elaborate_generator_drops`.
         let drop_clean = insert_clean_drop(mir);
 

src/test/run-fail/generator-resume-after-panic.rs

@@ -0,0 +1,22 @@
+// error-pattern:generator resumed after panicking
+
+// Test that we get the correct message for resuming a panicked generator.
+
+#![feature(generators, generator_trait)]
+
+use std::{
+    ops::Generator,
+    pin::Pin,
+    panic,
+};
+
+fn main() {
+    let mut g = || {
+        panic!();
+        yield;
+    };
+    panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        let x = Pin::new(&mut g).resume();
+    }));
+    Pin::new(&mut g).resume();
+}