From 4813dd741155a0e1ae063ff777e3eac76979d2e3 Mon Sep 17 00:00:00 2001
From: edef <edef@edef.eu>
Date: Sat, 25 Feb 2017 14:09:35 +0100
Subject: [PATCH] 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 | 52 +++++++++++++++++++++++++++++++++++++++++-------
 src/lib.rs       |  2 +-
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/src/generator.rs b/src/generator.rs
index 0a5728a..2a9ca4b 100644
--- a/src/generator.rs
+++ b/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)
+      }
     }
   }
 }
diff --git a/src/lib.rs b/src/lib.rs
index f9764dd..1a4957b 100644
--- a/src/lib.rs
+++ b/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]