Leak the stack if it's unsafe to drop it

We can't free the stack before the generator has returned, or been
unwound. Without unwinding, the only safe course of action is to leak
it.

src/generator.rs

@@ -82,12 +82,27 @@ pub enum State {
 #[derive(Debug)]
 pub struct Generator<'a, Input: 'a, Output: 'a, Stack: stack::Stack> {
   state:     State,
-  stack:     Stack,
-  stack_id:  debug::StackId,
+  stack:     NoDrop<Stack>,
+  stack_id:  NoDrop<debug::StackId>,
   stack_ptr: arch::StackPointer,
   phantom:   PhantomData<(&'a (), *mut Input, *const Output)>
 }
 
+#[allow(unions_with_drop_fields)]
+union NoDrop<T> {
+  value: T,
+  _empty: ()
+}
+
+impl<T: ::core::fmt::Debug> ::core::fmt::Debug for NoDrop<T> {
+  fn fmt(&self, w: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
+    unsafe {
+      // this is safe because we never invoke formatting on the empty variant
+      self.value.fmt(w)
+    }
+  }
+}
+
 impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
     where Input: 'a, Output: 'a, Stack: stack::Stack {
   /// Creates a new generator.
@@ -131,8 +146,8 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
 
     Generator {
       state:     State::Runnable,
-      stack:     stack,
-      stack_id:  stack_id,
+      stack:     NoDrop { value: stack },
+      stack_id:  NoDrop { value: stack_id },
       stack_ptr: stack_ptr,
       phantom:   PhantomData
     }
@@ -151,7 +166,7 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
 
         // Switch to the generator function, and retrieve the yielded value.
         let val = unsafe {
-          let (data_out, stack_ptr) = arch::swap(&input as *const Input as usize, self.stack_ptr, Some(&self.stack));
+          let (data_out, stack_ptr) = arch::swap(&input as *const Input as usize, self.stack_ptr, Some(&self.stack.value));
           self.stack_ptr = stack_ptr;
           mem::forget(input);
           ptr::read(data_out as *const Option<Output>)
@@ -176,8 +191,31 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
   /// (i.e. `self.state() == State::Runnable`), panics.
   pub fn unwrap(self) -> Stack {
     match self.state {
-      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
-      State::Unavailable => self.stack
+      State::Runnable => {
+        mem::forget(self);
+        panic!("Argh! Bastard! Don't touch that!")
+      }
+      State::Unavailable => unsafe { self.unsafe_unwrap() }
+    }
+  }
+
+  unsafe fn unsafe_unwrap(mut self) -> Stack {
+    ptr::drop_in_place(&mut self.stack_id.value);
+    let stack = ptr::read(&mut self.stack.value);
+    mem::forget(self);
+    stack
+  }
+}
+
+impl<'a, Input, Output, Stack> Drop for Generator<'a, Input, Output, Stack>
+    where Input: 'a, Output: 'a, Stack: stack::Stack {
+  fn drop(&mut self) {
+    unsafe {
+      ptr::drop_in_place(&mut self.stack_id.value);
+      match self.state {
+        State::Runnable    => {} // leak the stack
+        State::Unavailable => ptr::drop_in_place(&mut self.stack.value)
+      }
     }
   }
 }

src/lib.rs

@@ -4,7 +4,7 @@
 // http://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(asm, naked_functions, cfg_target_vendor)]
+#![feature(asm, naked_functions, cfg_target_vendor, untagged_unions)]
 #![cfg_attr(feature = "alloc", feature(alloc, heap_api))]
 #![cfg_attr(test, feature(test))]
 #![no_std]