Pack generator inputs and outputs into usize if possible

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<T> {
+  inner: T,
+}
+
+// Try to pack a value into a usize if it fits, otherwise pass its address in
+// the usize.
+unsafe fn encode<T>(val: &NoDrop<T>) -> usize {
+  if mem::size_of::<T>() <= mem::size_of::<usize>() &&
+     mem::align_of::<T>() <= mem::align_of::<usize>() {
+    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<T>(val: usize) -> T {
+  if mem::size_of::<T>() <= mem::size_of::<usize>() &&
+     mem::align_of::<T>() <= mem::align_of::<usize>() {
+    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, Output>, Input) {
       // Retrieve our environment from the callee and return control to it.
-      let f = ptr::read(env as *const F);
+      let f = decode::<F>(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::<Input>(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::<Input, Output, Stack, F>);
 
     // 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::<Output>(data_out)),
           None => None,
         }
       }
@@ -203,10 +234,10 @@ impl<Input, Output> Yielder<Input, Output>
   #[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::<Input>(data)
     }
   }
 }

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, 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]