diff --git a/src/generator.rs b/src/generator.rs index 37101f3f..201bd8d4 100644 --- a/src/generator.rs +++ b/src/generator.rs @@ -19,6 +19,37 @@ use stack; use debug; use arch::{self, StackPointer}; +// Wrapper to prevent the compiler from automatically dropping a value when it +// goes out of scope. This is particularly useful when dealing with unwinding +// since mem::forget won't be executed when unwinding. +#[allow(unions_with_drop_fields)] +union NoDrop { + inner: T, +} + +// Try to pack a value into a usize if it fits, otherwise pass its address in +// the usize. +unsafe fn encode(val: &NoDrop) -> usize { + if mem::size_of::() <= mem::size_of::() && + mem::align_of::() <= mem::align_of::() { + let mut out = 0; + ptr::copy_nonoverlapping(&val.inner, &mut out as *mut usize as *mut T, 1); + out + } else { + &val.inner as *const T as usize + } +} + +// Unpack a usize produced by encode. +unsafe fn decode(val: usize) -> T { + if mem::size_of::() <= mem::size_of::() && + mem::align_of::() <= mem::align_of::() { + ptr::read(&val as *const usize as *const T) + } else { + ptr::read(val as *const T) + } +} + #[derive(Debug, Clone, Copy)] pub enum State { /// Generator can be resumed. This is the initial state. @@ -114,10 +145,10 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack> where Input: Send, Output: Send, Stack: stack::Stack, F: FnOnce(&Yielder, Input) { // Retrieve our environment from the callee and return control to it. - let f = ptr::read(env as *const F); + let f = decode::(env); let (data, stack_ptr) = arch::swap(0, stack_ptr); // See the second half of Yielder::suspend_bare. - let input = ptr::read(data as *const Input); + let input = decode::(data); // Run the body of the generator. let yielder = Yielder::new(stack_ptr); f(&yielder, input); @@ -127,8 +158,8 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack> let stack_ptr = arch::init(&stack, generator_wrapper::); // Transfer environment to the callee. - let stack_ptr = arch::swap_link(&f as *const F as usize, stack_ptr, &stack).1; - mem::forget(f); + let f2 = NoDrop { inner: f }; + let stack_ptr = arch::swap_link(encode(&f2), stack_ptr, &stack).1; Generator { stack: stack, @@ -151,13 +182,13 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack> // Switch to the generator function, and retrieve the yielded value. unsafe { - let (data_out, stack_ptr) = arch::swap_link(&input as *const Input as usize, stack_ptr, &self.stack); + let input2 = NoDrop { inner: input }; + let (data_out, stack_ptr) = arch::swap_link(encode(&input2), stack_ptr, &self.stack); self.stack_ptr = stack_ptr; - mem::forget(input); // If the generator function has finished, return None. match stack_ptr { - Some(_) => Some(ptr::read(data_out as *const Output)), + Some(_) => Some(decode::(data_out)), None => None, } } @@ -203,10 +234,10 @@ impl Yielder #[inline(always)] pub fn suspend(&self, item: Output) -> Input { unsafe { - let (data, stack_ptr) = arch::swap(&item as *const Output as usize, self.stack_ptr.get()); - mem::forget(item); + let item2 = NoDrop { inner: item }; + let (data, stack_ptr) = arch::swap(encode(&item2), self.stack_ptr.get()); self.stack_ptr.set(stack_ptr); - ptr::read(data as *const Input) + decode::(data) } } } diff --git a/src/lib.rs b/src/lib.rs index 3db5cfd5..75432aae 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, nonzero)] +#![feature(asm, naked_functions, cfg_target_vendor, nonzero, untagged_unions)] #![cfg_attr(feature = "alloc", feature(alloc, heap_api))] #![cfg_attr(test, feature(test))] #![no_std]