WIP unwinding support

src/arch/x86_64.rs

@@ -49,6 +49,7 @@
 use core::mem;
 use stack;
 use arch::StackPointer;
+use unwind;
 
 pub const STACK_ALIGNMENT: usize = 16;
 
@@ -127,12 +128,10 @@ pub unsafe fn init<Stack: stack::Stack>(stack: &Stack, f: unsafe extern "C" fn(u
         # trampoline_2.
         nop
 
-        # Call the provided function.
-        callq   *16(%rsp)
-
-        # Clear the stack pointer. We can't call into this context any more once
-        # the function has returned.
-        xorq    %rsi, %rsi
+        # Call unwind_wrapper with the provided function and the CFA address.
+        movq    %rsp, %rdx
+        movq    16(%rsp), %rcx
+        call    ${0:c}
 
         # Restore the stack pointer of the parent context. No CFI adjustments
         # are needed since we have the same stack frame as trampoline_1.
@@ -143,14 +142,27 @@ pub unsafe fn init<Stack: stack::Stack>(stack: &Stack, f: unsafe extern "C" fn(u
         .cfi_adjust_cfa_offset -8
         .cfi_restore %rbp
 
+        # If the returned value is nonzero, trigger an unwind in the parent
+        # context with the given exception object.
+        movq    %rax, %rdi
+        testq   %rax, %rax
+        jnz     ${1:c}
+
+        # Clear the stack pointer. We can't call into this context any more once
+        # the function has returned.
+        xorq    %rsi, %rsi
+
         # Return into the parent context. Use `pop` and `jmp` instead of a `ret`
         # to avoid return address mispredictions (~8ns per `ret` on Ivy Bridge).
-        popq    %rax
+        popq    %rcx
         .cfi_adjust_cfa_offset -8
-        .cfi_register %rip, %rax
-        jmpq    *%rax
+        .cfi_register %rip, %rcx
+        jmpq    *%rcx
       "#
-      : : : : "volatile")
+      :
+      : "s" (unwind::unwind_wrapper as usize)
+        "s" (unwind::start_unwind as usize)
+      : : "volatile")
   }
 
   // We set up the stack in a somewhat special way so that to the unwinder it
@@ -213,10 +225,10 @@ pub unsafe fn swap_link<Stack: stack::Stack>(arg: usize, new_sp: StackPointer,
 
         # Return into the new context. Use `pop` and `jmp` instead of a `ret`
         # to avoid return address mispredictions (~8ns per `ret` on Ivy Bridge).
-        popq    %rax
+        popq    %rcx
         .cfi_adjust_cfa_offset -8
-        .cfi_register %rip, %rax
-        jmpq    *%rax
+        .cfi_register %rip, %rcx
+        jmpq    *%rcx
       "#
       : : : : "volatile")
   }
@@ -267,10 +279,10 @@ pub unsafe fn swap(arg: usize, new_sp: StackPointer) -> (usize, StackPointer) {
         popq    %rbp
         .cfi_adjust_cfa_offset -8
         .cfi_restore %rbp
-        popq    %rax
+        popq    %rcx
         .cfi_adjust_cfa_offset -8
-        .cfi_register %rip, %rax
-        jmpq    *%rax
+        .cfi_register %rip, %rcx
+        jmpq    *%rcx
       "#
       : : : : "volatile")
   }
@@ -297,3 +309,55 @@ pub unsafe fn swap(arg: usize, new_sp: StackPointer) -> (usize, StackPointer) {
     : "volatile", "alignstack");
   (ret, mem::transmute(ret_sp))
 }
+
+
+#[inline(always)]
+pub unsafe fn unwind<Stack: stack::Stack>(new_sp: StackPointer, new_stack: &Stack)  {
+  // Address of the topmost CFA stack slot.
+  let new_cfa = StackPointer::stack_base(new_stack).offset(-4);
+
+  // Argument to pass to start_unwind, based on the CFA slot address.
+  let arg = unwind::unwind_arg(new_cfa);
+
+  // This is identical to swap_link, except that it doesn't set %rsi and jumps
+  // (tail call) to start_unwind instead of returning into the target context.
+  #[naked]
+  unsafe extern "C" fn trampoline() {
+    asm!(
+      r#"
+        pushq   %rbp
+        .cfi_adjust_cfa_offset 8
+        .cfi_rel_offset %rbp, 0
+        movq    %rsp, (%rcx)
+        movq    %rdx, %rsp
+        popq    %rbp
+        .cfi_adjust_cfa_offset -8
+        .cfi_restore %rbp
+
+        # Jump to the start_unwind function, which will force a stack unwind in
+        # the target context. This will eventually return to our caller through
+        # the stack link.
+        jmp     ${0:c}
+      "#
+      : : "s" (unwind::start_unwind as usize) : : "volatile")
+  }
+
+  asm!(
+    r#"
+      call    ${0:c}
+    "#
+    :
+    : "s" (trampoline as usize)
+      "{rdi}" (arg)
+      "{rdx}" (new_sp.0)
+      "{rcx}" (new_cfa)
+    : "rax",   "rbx",   "rcx",   "rdx",   "rsi",   "rdi", /*"rbp",   "rsp",*/
+      "r8",    "r9",    "r10",   "r11",   "r12",   "r13",   "r14",   "r15",
+      "mm0",   "mm1",   "mm2",   "mm3",   "mm4",   "mm5",   "mm6",   "mm7",
+      "xmm0",  "xmm1",  "xmm2",  "xmm3",  "xmm4",  "xmm5",  "xmm6",  "xmm7",
+      "xmm8",  "xmm9",  "xmm10", "xmm11", "xmm12", "xmm13", "xmm14", "xmm15",
+      "xmm16", "xmm17", "xmm18", "xmm19", "xmm20", "xmm21", "xmm22", "xmm23",
+      "xmm24", "xmm25", "xmm26", "xmm27", "xmm28", "xmm29", "xmm30", "xmm31",
+      "cc", "dirflag", "fpsr", "flags", "memory"
+    : "volatile", "alignstack");
+}

src/generator.rs

@@ -201,27 +201,39 @@ impl<'a, Input, Output, Stack> Generator<'a, Input, Output, Stack>
   /// Extracts the stack from a generator when the generator function has returned.
   /// If the generator function has not returned
   /// (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
+  pub fn unwrap(mut self) -> Stack {
+    unsafe {
+      self.stack_ptr.map(|stack_ptr| arch::unwind(stack_ptr, &self.stack));
+
+      // We can't just return self.stack since Generator has a Drop impl
+      let stack = ptr::read(&self.stack);
+      ptr::drop_in_place(&mut self.stack_id);
+      mem::forget(self);
+      stack
     }
   }
 }
 
+impl<'a, Input, Output, Stack> Drop for Generator<'a, Input, Output, Stack>
+    where Input: Send + 'a, Output: Send + 'a, Stack: stack::Stack {
+  fn drop(&mut self) {
+    self.stack_ptr.map(|stack_ptr| unsafe { arch::unwind(stack_ptr, &self.stack) });
+  }
+}
+
 /// Yielder is an interface provided to every generator through which it
 /// returns a value.
 #[derive(Debug)]
 pub struct Yielder<Input: Send, Output: Send> {
-  stack_ptr: Cell<StackPointer>,
+  stack_ptr: Cell<Option<StackPointer>>,
   phantom: PhantomData<(*const Input, *mut Output)>
 }
 
 impl<Input, Output> Yielder<Input, Output>
     where Input: Send, Output: Send {
   fn new(stack_ptr: StackPointer) -> Yielder<Input, Output> {
     Yielder {
-      stack_ptr: Cell::new(stack_ptr),
+      stack_ptr: Cell::new(Some(stack_ptr)),
       phantom: PhantomData
     }
   }
@@ -231,9 +243,24 @@ impl<Input, Output> Yielder<Input, Output>
   #[inline(always)]
   pub fn suspend(&self, item: Output) -> Input {
     unsafe {
+      struct PanicGuard<'a>(&'a Cell<Option<StackPointer>>);
+      impl<'a> Drop for PanicGuard<'a> {
+        fn drop(&mut self) {
+          self.0.set(None);
+        }
+      }
+
+      let stack_ptr = self.stack_ptr.get().expect("attempted to yield while unwinding");
       let item2 = NoDrop { inner: item };
-      let (data, stack_ptr) = arch::swap(encode_usize(&item2), self.stack_ptr.get());
-      self.stack_ptr.set(stack_ptr);
+
+      // Use a PanicGuard to set self.stack_ptr to None if unwinding occurs. This
+      // is necessary to guarantee safety in case someone tries to call yield
+      // while we are unwinding since there is nowhere to yield to.
+      let guard = PanicGuard(&self.stack_ptr);
+      let (data, stack_ptr) = arch::swap(encode_usize(&item2), stack_ptr);
+      mem::forget(guard);
+
+      self.stack_ptr.set(Some(stack_ptr));
       decode_usize(data)
     }
   }

src/lib.rs

@@ -51,6 +51,8 @@ pub const STACK_ALIGNMENT: usize = arch::STACK_ALIGNMENT;
 
 mod debug;
 
+mod unwind;
+
 mod stack;
 mod slice_stack;
 pub mod generator;

src/os/mod.rs

@@ -28,8 +28,8 @@ impl Stack {
   pub fn new(size: usize) -> Result<Stack, IoError> {
     let page_size = sys::page_size();
 
-    // Stacks have to be at least one page long.
-    let len = if size == 0 { page_size } else { size };
+    // Stacks have to be at least 16KB to support unwinding.
+    let len = if size == 0 { 16384 } else { size };
 
     // Round the length one page size up, using the fact that the page size
     // is a power of two.

src/unwind.rs

@@ -0,0 +1,67 @@
+// This file is part of libfringe, a low-level green threading library.
+// Copyright (c) Amanieu d'Antras <amanieu@gmail.com>,
+// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
+// 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.
+extern crate std;
+
+use self::std::panic;
+use self::std::boxed::Box;
+use core::any::Any;
+use arch::StackPointer;
+
+// Marker object that is passed through the stack unwinding
+struct UnwindMarker {
+  // We use the CFA address as an identifier so that nested generators are handled
+  // correctly. When unwinding, we will want to continue through any number of
+  // nested generators until we reach the one identified with a matching CFA address.
+  cfa_addr: *mut usize,
+}
+unsafe impl Send for UnwindMarker {}
+
+// Whether the current platform support unwinding across multiple stacks.
+fn have_cross_stack_unwind() -> bool {
+  // The only platform which doesn't support cross-stack unwinding is Windows
+  // since it uses SEH for unwinding instead of libunwind.
+  !cfg!(windows)
+}
+
+pub unsafe extern "C" fn unwind_wrapper(arg: usize, sp: StackPointer, cfa_addr: *mut usize,
+                                        f: unsafe extern "C" fn(usize, StackPointer)) -> Option<Box<Box<Any + Send>>> {
+  // Catch any attempts to unwind out of the context.
+  match panic::catch_unwind(move || f(arg, sp)) {
+    Ok(_) => None,
+    Err(err) => {
+      // If the unwinding is due to an UnwindMarker, check whether it is intended
+      // for us by comparing the CFA of the caller with ours. If it is the same
+      // then we can swallow the exception and return to the caller normally.
+      if let Some(marker) = err.downcast_ref::<UnwindMarker>() {
+        if marker.cfa_addr == cfa_addr {
+          return None;
+        }
+      }
+
+      // Otherwise, propagate the panic to the parent context.
+      if have_cross_stack_unwind() {
+        panic::resume_unwind(err)
+      } else {
+        // The assembly code will call start_unwind in the parent context and
+        // pass it this Box as parameter.
+        Some(Box::new(err))
+      }
+    }
+  }
+}
+
+pub unsafe extern "C" fn start_unwind(panic: Box<Box<Any + Send>>) -> ! {
+  // Use resume_unwind instead of panic! to avoid printing a message.
+  panic::resume_unwind(*panic)
+}
+
+pub fn unwind_arg(cfa_addr: *mut usize) -> usize {
+  let marker = UnwindMarker {
+    cfa_addr: cfa_addr
+  };
+  Box::into_raw(Box::new(Box::new(marker) as Box<Any + Send>)) as usize
+}

tests/generator.rs

@@ -68,7 +68,7 @@ fn panic_safety() {
 
 #[test]
 fn with_slice_stack() {
-  let mut memory = [0; 1024];
+  let mut memory = [0; 16384];
   let stack = SliceStack(&mut memory);
   let mut add_one = unsafe { Generator::unsafe_new(stack, add_one_fn) };
   assert_eq!(add_one.resume(1), Some(2));
@@ -77,7 +77,7 @@ fn with_slice_stack() {
 
 #[test]
 fn with_owned_stack() {
-  let stack = OwnedStack::new(1024);
+  let stack = OwnedStack::new(16384);
   let mut add_one = unsafe { Generator::unsafe_new(stack, add_one_fn) };
   assert_eq!(add_one.resume(1), Some(2));
   assert_eq!(add_one.resume(2), Some(3));