Potential fix for the stackrestore problem

src/jump.ll

@@ -1,55 +1,55 @@
-declare void @llvm.eh.sjlj.longjmp(i8*)
-declare i32 @llvm.eh.sjlj.setjmp(i8*)
-declare void @llvm.stackrestore(i8*)
-declare i8* @llvm.frameaddress(i32)
-declare i8* @llvm.stacksave()
+declare void @llvm.eh.sjlj.longjmp(i8*) noreturn nounwind
+declare i32 @llvm.eh.sjlj.setjmp(i8*) nounwind
+declare void @llvm.stackrestore(i8*) nounwind
+declare i8* @llvm.frameaddress(i32) nounwind readnone
+declare i8* @llvm.stacksave() nounwind
 
-; This function is essentially what __builtin_setjmp() emits in C.
 ; We put it here and mark it as alwaysinline for code-reuse.
 ; This function is internal only.
 define private i32
-@jump_save(i8** nonnull %ctx)
-alwaysinline nounwind naked
+@jump_save([5 x i8*]* nonnull %ctx)
+alwaysinline nounwind naked returns_twice
 {
   ; Store the frame address.
   %frame = call i8* @llvm.frameaddress(i32 0)
-  %foff = getelementptr inbounds i8*, i8** %ctx, i32 0
+  %foff = getelementptr inbounds [5 x i8*], [5 x i8*]* %ctx, i64 0, i32 0
   store i8* %frame, i8** %foff
 
   ; Store the stack address.
   %stack = call i8* @llvm.stacksave()
-  %soff = getelementptr inbounds i8*, i8** %ctx, i32 2
+  %soff = getelementptr inbounds [5 x i8*], [5 x i8*]* %ctx, i64 0, i32 2
   store i8* %stack, i8** %soff
 
   ; The rest are architecture specific and stored by setjmp().
-  %buff = bitcast i8** %ctx to i8*
-  %retv = call i32 @llvm.eh.sjlj.setjmp(i8* %buff)
+  %buff = bitcast [5 x i8*]* %ctx to i8*
+  %retv = call i32 @llvm.eh.sjlj.setjmp(i8* %buff) returns_twice
   ret i32 %retv
 }
 
 ; This function is essentially what __builtin_longjmp() emits in C.
 ; The purpose is to expose this intrinsic to Rust (without requiring nightly).
 define dso_local void
-@jump_into(i8** %into)
+@jump_into([5 x i8*]* nonnull %into)
 noreturn nounwind naked
 {
-  %buff = bitcast i8** %into to i8*
+  %buff = bitcast [5 x i8*]* %into to i8*
   call void @llvm.eh.sjlj.longjmp(i8* %buff) ; Call longjmp()
   unreachable
 }
 
 ; This function performs a bidirectional context switch.
 ; It simply calls setjmp(%from) and then longjmp(%into).
 define dso_local void
-@jump_swap(i8** %from, i8** %into)
+@jump_swap([5 x i8*]* nonnull %from, [5 x i8*]* nonnull %into)
 nounwind
 {
-  %retv = call i32 @jump_save(i8** %from)    ; setjmp(%from)
+  ; setjmp(%from)
+  %retv = call i32 @jump_save([5 x i8*]* %from) returns_twice
   %zero = icmp eq i32 %retv, 0
   br i1 %zero, label %jump, label %done
 
 jump:                                        ; setjmp(%from) returned 0
-  %ibuf = bitcast i8** %into to i8*
+  %ibuf = bitcast [5 x i8*]* %into to i8*
   call void @llvm.eh.sjlj.longjmp(i8* %ibuf) ; longjmp(%into)
   unreachable
 
@@ -62,20 +62,35 @@ done:                                        ; setjmp(%from) returned !0
 ;   2. Set the stack pointer to %addr.
 ;   3. Call %func(%c, %f).
 define dso_local void
-@jump_init(i8* %addr, i8* %c, i8* %f, void (i8**, i8*, i8*)* %func)
+@jump_init(i8* %addr, i8* %c, i8* %f, void ([5 x i8*]*, i8*, i8*)* %func)
 nounwind
 {
+  %buff = alloca [5 x i8*], align 4          ; Allocate setjmp() buffer
 
-  %buff = alloca [5 x i8*]                    ; Allocate setjmp() buffer
-
-  %cast = bitcast [5 x i8*]* %buff to i8**
-  %retv = call i32 @jump_save(i8** %cast)     ; Call setjmp(%buff)
+  ; Call setjmp(%buff)
+  %retv = call i32 @jump_save([5 x i8*]* %buff) returns_twice
   %zero = icmp eq i32 %retv, 0
   br i1 %zero, label %next, label %done
 
 next:                                         ; setjmp(%buff) returned 0
+
+  ; FIXME, the correct solution would be to store those in real registers
+
+  %1 = alloca i8*, align 4
+  %2 = alloca i8*, align 4
+  %3 = alloca void ([5 x i8*]*, i8*, i8*)*, align 4
+
+  store i8* %c, i8** %1
+  store i8* %f, i8** %2
+  store void ([5 x i8*]*, i8*, i8*)* %func, void ([5 x i8*]*, i8*, i8*)** %3
+
   call void @llvm.stackrestore(i8* %addr)     ; Move onto new stack %addr
-  call void %func(i8** %cast, i8* %c, i8* %f) ; Call %func(%buff, %c, %f)
+
+  %gc = load i8*, i8** %1
+  %gf = load i8*, i8** %2
+  %gfunc = load void ([5 x i8*]*, i8*, i8*)*, void ([5 x i8*]*, i8*, i8*)** %3
+
+  call void %gfunc([5 x i8*]* %buff, i8* %gc, i8* %gf) ; Call %func(%buff, %c, %f)
   unreachable
 
 done:                                         ; setjmp(%buff) returned !0

src/lib.rs

@@ -74,24 +74,35 @@ const STACK_ALIGNMENT: usize = 16;
 const STACK_MINIMUM: usize = 4096;
 
 extern "C" {
-    fn jump_into(into: *mut *mut c_void) -> !;
-    fn jump_swap(from: *mut *mut c_void, into: *mut *mut c_void);
+    fn jump_into(into: *mut [*mut c_void; 5]) -> !;
+    fn jump_swap(from: *mut [*mut c_void; 5], into: *mut [*mut c_void; 5]);
     fn jump_init(
         stack: *mut u8,
         ctx: *mut c_void,
         fnc: *mut c_void,
         func: unsafe extern "C" fn(
-            parent: *mut *mut c_void,
-            ctxpp: *mut c_void,
-            fncpp: *mut c_void,
+            parent: *mut [*mut c_void; 5],
+            ctx: *mut c_void,
+            fnc: *mut c_void,
         ) -> !,
     );
 }
 
+#[repr(C, align(16))]
 struct Context<Y, R> {
     parent: [*mut c_void; 5],
     child: [*mut c_void; 5],
-    arg: *mut GeneratorState<Y, R>,
+    arg: MaybeUninit<*mut GeneratorState<Y, R>>,
+}
+
+impl<Y, R> Default for Context<Y, R> {
+    fn default() -> Self {
+        Context {
+            parent: [null_mut(); 5],
+            child: [null_mut(); 5],
+            arg: MaybeUninit::uninit(),
+        }
+    }
 }
 
 #[cfg(not(has_generator_trait))]
@@ -165,39 +176,46 @@ pub struct Canceled(());
 
 pub struct Coroutine<'a, Y, R>(Option<&'a mut Context<Y, R>>);
 
-unsafe extern "C" fn callback<Y, R, F>(p: *mut *mut c_void, c: *mut c_void, f: *mut c_void) -> !
+unsafe extern "C" fn callback<Y, R, F>(
+    p: *mut [*mut c_void; 5],
+    c: *mut c_void,
+    f: *mut c_void,
+) -> !
 where
     F: FnOnce(Control<'_, Y, R>) -> Result<Finished<R>, Canceled>,
 {
     // Allocate a Context and a closure.
-    let mut ctx = MaybeUninit::zeroed().assume_init();
-    let mut fnc = MaybeUninit::uninit().assume_init();
+    let mut ctx = Context::<Y, R>::default();
+    let mut fnc = MaybeUninit::<F>::uninit();
 
     // Cast the incoming pointers to their correct types.
     // See `Coroutine::new()`.
     let c = c as *mut Coroutine<'_, Y, R>;
-    let f = f as *mut &mut F;
+    let f = f as *mut *mut F;
 
     // Pass references to the stack-allocated Context and closure back into
     // Coroutine::new() through the incoming pointers.
     (*c).0 = Some(&mut ctx);
-    *f = &mut fnc;
+    *f = fnc.as_mut_ptr();
 
     // Yield control to the parent. The first call to `Generator::resume()`
     // will resume at this location. The `Coroutine::new()` function is
     // responsible to move the closure into this stack while we are yielded.
-    jump_swap(ctx.child.as_mut_ptr(), p);
+    jump_swap(ctx.child.as_mut_ptr() as _, p as _);
+
+    let fnc = fnc.assume_init();
 
     // Call the closure. If the closure returns, then move the return value
     // into the argument variable in `Generator::resume()`.
     if let Ok(r) = fnc(Control(&mut ctx)) {
-        if !ctx.arg.is_null() {
-            *ctx.arg = GeneratorState::Complete(r.0);
+        let arg = ctx.arg.assume_init();
+        if !arg.is_null() {
+            *arg = GeneratorState::Complete(r.0);
         }
     }
 
     // We cannot be resumed, so jump away forever.
-    jump_into(ctx.parent.as_mut_ptr());
+    jump_into(ctx.parent.as_mut_ptr() as _);
 }
 
 impl<'a, Y, R> Coroutine<'a, Y, R> {
@@ -224,7 +242,7 @@ impl<'a, Y, R> Coroutine<'a, Y, R> {
         // it is going to store references to those instances inside these
         // variables.
         let mut cor = Coroutine(None);
-        let mut fnc: Option<&mut F> = None;
+        let mut fnc = MaybeUninit::<&mut F>::uninit();
 
         assert!(stack.len() >= STACK_MINIMUM);
 
@@ -237,14 +255,14 @@ impl<'a, Y, R> Coroutine<'a, Y, R> {
             jump_init(
                 top,
                 &mut cor as *mut _ as _,
-                &mut fnc as *mut _ as _,
+                fnc.as_mut_ptr() as *mut _ as _,
                 callback::<Y, R, F>,
             );
+            let fnc = fnc.assume_init();
+            // Move the closure onto the coroutine's stack.
+            *fnc = func;
         }
 
-        // Move the closure onto the coroutine's stack.
-        *fnc.unwrap() = func;
-
         cor
     }
 }
@@ -262,22 +280,30 @@ impl<'a, Y, R> Control<'a, Y, R> {
     /// exists.
     pub fn r#yield(self, arg: Y) -> Result<Self, Canceled> {
         unsafe {
+            let ptr_arg = self.0.arg.assume_init();
+
             // The parent `Coroutine` object has been dropped. Resume the child
             // coroutine with the Canceled error. It must clean up and exit.
-            if self.0.arg.is_null() {
+            if ptr_arg.is_null() {
                 return Err(Canceled(()));
             }
 
             // Move the argument value into the argument variable in
             // `Generator::resume()`.
-            *self.0.arg = GeneratorState::Yielded(arg);
+            *ptr_arg = GeneratorState::Yielded(arg);
 
             // Save our current position and yield control to the parent.
-            jump_swap(self.0.child.as_mut_ptr(), self.0.parent.as_mut_ptr());
+            jump_swap(
+                self.0.child.as_mut_ptr() as _,
+                self.0.parent.as_mut_ptr() as _,
+            );
+
+            // Let the compiler re-read *self.0.arg
+            let ptr_arg = self.0.arg.as_mut_ptr().read_volatile();
 
             // The parent `Coroutine` object has been dropped. Resume the child
             // coroutine with the Canceled error. It must clean up and exit.
-            if self.0.arg.is_null() {
+            if ptr_arg.is_null() {
                 return Err(Canceled(()));
             }
         }
@@ -300,40 +326,42 @@ impl<'a, Y, R> Generator for Coroutine<'a, Y, R> {
     fn resume(mut self: Pin<&mut Self>) -> GeneratorState<Y, R> {
         // Allocate an argument variable on the stack. See `Control::r#yield()` and
         // `callback()` for where this is initialized.
-        let mut arg = unsafe { MaybeUninit::uninit().assume_init() };
+        let mut arg = MaybeUninit::<GeneratorState<Y, R>>::uninit();
 
         match self.0 {
             None => panic!("Called Generator::resume() after completion!"),
             Some(ref mut p) => unsafe {
                 // Pass the pointer so that the child can move the argument out.
-                p.arg = &mut arg;
+                p.arg.as_mut_ptr().write_volatile(arg.as_mut_ptr());
 
                 // Jump back into the child.
-                jump_swap(p.parent.as_mut_ptr(), p.child.as_mut_ptr());
+                jump_swap(p.parent.as_mut_ptr() as _, p.child.as_mut_ptr() as _);
 
                 // Clear the pointer as the value is about to become invalid.
-                p.arg = null_mut();
+                p.arg.as_mut_ptr().write_volatile(null_mut());
             },
         }
 
+        let state = unsafe { arg.assume_init() };
+
         // If the child coroutine has completed, we are done. Make it so that
         // we can never resume the coroutine by clearing the reference.
-        if let GeneratorState::Complete(r) = arg {
+        if let GeneratorState::Complete(_) = state {
             self.0 = None;
-            return GeneratorState::Complete(r);
         }
 
-        arg
+        state
     }
 }
 
 impl<'a, Y, R> Drop for Coroutine<'a, Y, R> {
     fn drop(&mut self) {
-        // If we are still able to resume the coroutine, do so. Since we don't
-        // set the argument pointer, `Control::halt()` will return `Canceled`.
+        // If we are still able to resume the coroutine, do so.
         if let Some(x) = self.0.take() {
             unsafe {
-                jump_swap(x.parent.as_mut_ptr(), x.child.as_mut_ptr());
+                // set the argument pointer to null, `Control::r#yield()` will return `Canceled`.
+                x.arg.as_mut_ptr().write_volatile(null_mut());
+                jump_swap(x.parent.as_mut_ptr() as _, x.child.as_mut_ptr() as _);
             }
         }
     }