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_id: NoDrop, stack_ptr: arch::StackPointer, phantom: PhantomData<(&'a (), *mut Input, *const Output)> } +#[allow(unions_with_drop_fields)] +union NoDrop { + value: T, + _empty: () +} + +impl ::core::fmt::Debug for NoDrop { + 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) @@ -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 , 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]