Avoid multi-return asm and add misopt test

It is undefined to exit an asm block multiple times with only one entry.
And LLVM will derive that a branch ending with a noreturn call can
freely clobber stack variables without having observable effects, which
is not compatible with an all-branch-executed control flow.

See: <https://doc.rust-lang.org/stable/reference/inline-assembly.html#r-asm.rules.only-on-exit>

Discussion: <rust-lang/rfcs#2625 (comment)>

Author: oxalica

src/aarch64.rs

@@ -1,7 +1,5 @@
 use super::NonZero;
 
-pub(crate) type Buf = [*mut (); 4];
-
 #[cfg(not(target_os = "macos"))]
 macro_rules! set_jump_raw_impl {
     ($($tt:tt)*) => {
@@ -27,7 +25,6 @@ macro_rules! set_jump_raw_impl {
             lateout("lr") _,
             // Caller saved registers.
             clobber_abi("C"),
-            options(readonly),
         )
     }
 }
@@ -64,50 +61,58 @@ macro_rules! set_jump_raw_impl {
             lateout("lr") _,
             // Caller saved registers.
             clobber_abi("C"),
-            options(readonly),
         )
     }
 }
 
 macro_rules! set_jump_raw {
-    ($val:expr, $buf_ptr:expr, $lander:block) => {
+    ($val:expr, $f:expr, $data:expr, $lander:block) => {
         set_jump_raw_impl!(
-            "adr x0, {lander}",
-            "mov x2, sp",
-            "stp x0, x2, [x1]",
-            "stp x19, x29, [x1, #16]",
+            "adr x2, {lander}",
+            "stp x2, x2, [sp, #-16]!",
+            ".cfi_adjust_cfa_offset 16",
+            "stp x19, x29, [sp, #-16]!",
+            ".cfi_adjust_cfa_offset 16",
+            "mov x1, sp",
+            "bl {f}",
+            "add sp, sp, 32",
+            ".cfi_adjust_cfa_offset -32",
 
-            in("x1") $buf_ptr, // Restored in long_jump_raw.
-            out("x0") $val,
+            f = sym $f,
             lander = label $lander,
+            inout("x0") $data => $val,
         )
     };
-    ($val:expr, $buf_ptr:expr) => {
+    ($val:expr, $f:expr, $data:expr) => {
         set_jump_raw_impl!(
-            "adr x0, 2f",
-            "mov x2, sp",
-            "stp x0, x2, [x1]",
-            "stp x19, x29, [x1, #16]",
-            "mov x0, 0",
+            "adr x2, 2f",
+            "stp x2, x2, [sp, #-16]!",
+            ".cfi_adjust_cfa_offset 16",
+            "stp x19, x29, [sp, #-16]!",
+            ".cfi_adjust_cfa_offset 16",
+            "mov x1, sp",
+            "bl {f}",
+            "add sp, sp, 32",
+            ".cfi_adjust_cfa_offset -32",
             "2:",
 
-            in("x1") $buf_ptr, // Restored in long_jump_raw.
-            out("x0") $val,
+            f = sym $f,
+            inout("x0") $data => $val,
         )
     };
 }
 
 #[inline]
-pub(crate) unsafe fn long_jump_raw(buf: *mut (), result: NonZero<usize>) -> ! {
+pub(crate) unsafe fn long_jump_raw(jp: *mut (), result: NonZero<usize>) -> ! {
     unsafe {
         core::arch::asm!(
-            "ldp x19, x29, [x1, #16]",
-            "ldp lr, x2, [x1]",
-            "mov sp, x2",
-            "ret",
-            in("x1") buf,
+            "ldp x19, x29, [x1]",
+            "ldr x2, [x1, #16]",
+            "add sp, x1, 32",
+            "br x2",
             in("x0") result.get(),
-            options(noreturn),
+            in("x1") jp,
+            options(noreturn, nostack, readonly),
         )
     }
 }

src/arm.rs

@@ -1,17 +1,9 @@
 use super::NonZero;
 
-pub(crate) type Buf = [*mut (); 4];
-
-macro_rules! set_jump_raw {
-    ($val:expr, $buf_ptr:expr, $lander:block) => {
+macro_rules! set_jump_raw_impl {
+    ($($tt:tt)*) => {
         core::arch::asm!(
-            "adr lr, {lander}",
-            "stm r1, {{r6, r11, sp, lr}}",
-
-            lander = label $lander,
-
-            out("r0") $val,
-            in("r1") $buf_ptr, // Restored in long_jump_raw.
+            $($tt)*
 
             // Callee saved registers.
             lateout("r4") _,
@@ -25,40 +17,42 @@ macro_rules! set_jump_raw {
             lateout("r12") _,
             // lateout("sp") _, // sp
             lateout("lr") _,
-
             // Caller saved registers.
+            // FIXME: inline asm clobber list contains reserved registers: D16-D31.
             clobber_abi("aapcs"),
-            options(readonly),
+        )
+    }
+}
+
+macro_rules! set_jump_raw {
+    ($val:expr, $f:expr, $data:expr, $lander:block) => {
+        set_jump_raw_impl!(
+            "adr lr, {lander}",
+            "push {{r6, r11, sp, lr}}",
+            ".cfi_adjust_cfa_offset 16",
+            "mov r1, sp",
+            "bl {f}",
+            "add sp, sp, 16",
+            ".cfi_adjust_cfa_offset -16",
+
+            f = sym $f,
+            inout("r0") $data => $val,
+            lander = label $lander,
         )
     };
-    ($val:expr, $buf_ptr:expr) => {
-        core::arch::asm!(
+    ($val:expr, $f:expr, $data:expr) => {
+        set_jump_raw_impl!(
             "adr lr, 2f",
-            "stm r1, {{r6, r11, sp, lr}}",
-            "mov r0, 0",
+            "push {{r6, r11, sp, lr}}",
+            ".cfi_adjust_cfa_offset 16",
+            "mov r1, sp",
+            "bl {f}",
+            "add sp, sp, 16",
+            ".cfi_adjust_cfa_offset -16",
             "2:",
 
-            // Caller saved registers.
-            // Workaround: see above.
-            out("r0") $val,
-            in("r1") $buf_ptr, // Restored in long_jump_raw.
-
-            // Callee saved registers.
-            lateout("r4") _,
-            lateout("r5") _,
-            // lateout("r6") _, // LLVM reserved.
-            lateout("r7") _,
-            lateout("r8") _,
-            lateout("r9") _,
-            lateout("r10") _,
-            // lateout("r11") _, // LLVM reserved.
-            lateout("r12") _,
-            // lateout("sp") _, // sp
-            lateout("lr") _,
-
-            // Caller saved registers.
-            clobber_abi("aapcs"),
-            options(readonly),
+            f = sym $f,
+            inout("r0") $data => $val,
         )
     };
 }
@@ -68,9 +62,9 @@ pub(crate) unsafe fn long_jump_raw(buf: *mut (), result: NonZero<usize>) -> ! {
     unsafe {
         core::arch::asm!(
             "ldm r1, {{r6, r11, sp, pc}}",
-            in("r1") buf,
             in("r0") result.get(),
-            options(noreturn),
+            in("r1") buf,
+            options(noreturn, nostack, readonly),
         )
     }
 }

src/lib.rs

@@ -10,8 +10,8 @@
 //!   [fatal interaction with optimizer][misopt].
 //!   This crate does not suffer from the misoptimization (covered in `tests/smoke.rs`).
 //!
-//!   ⚠️  We admit that since it's not yet undefined to force unwind Rust POFs and/or
-//!   longjmp's half execution semantic, `long_jump` is still technically undefined behavior.
+//!   ⚠️  We admit that since it's not yet undefined to force unwind Rust POFs,
+//!   `long_jump` is still technically undefined behavior.
 //!   But this crate is an attempt to make a semantically-correct abstraction free from
 //!   misoptimization, and you accept the risk by using this crate.
 //!   If you find any misoptimization in practice, please open an issue.
@@ -25,7 +25,7 @@
 //!   Single `usize` `JumpPoint`. Let optimizer save only necessary states rather than bulk saving
 //!   all callee-saved registers. Inline-able `set_jump` without procedure call cost.
 //!
-//!   - 2.3ns `set_jump` setup and 3.2ns `long_jump` on a modern x86\_64 CPU.
+//!   - 2.4ns `set_jump` setup and 2.9ns `long_jump` on a modern x86\_64 CPU.
 //!     ~300-490x faster than `catch_unwind`-`panic_any!`.
 //!
 //! - No std.
@@ -90,17 +90,16 @@
 //!
 //! - [`sjlj`](https://crates.io/crates/sjlj)
 //!
-//!   - Uses inline assembly but involving an un-inline-able call instruction.
 //!   - Only x86\_64 is supported.
-//!   - Suffers from [misoptimization][misopt].
+//!   - Suffers from [misoptimization][misopt] due to multi-return.
 //!   - Slower `long_jump` because of more register restoring.
 //!
 //! [misopt]: https://github.com/rust-lang/rfcs/issues/2625
 #![cfg_attr(feature = "unstable-asm-goto", feature(asm_goto))]
 #![cfg_attr(feature = "unstable-asm-goto", feature(asm_goto_with_outputs))]
 #![cfg_attr(not(test), no_std)]
 use core::marker::PhantomData;
-use core::mem::{size_of, MaybeUninit};
+use core::mem::ManuallyDrop;
 use core::num::NonZero;
 
 #[cfg(target_arch = "x86_64")]
@@ -147,8 +146,6 @@ mod imp {
 
     compile_error!("sjlj2: unsupported platform");
 
-    pub type Buf = [usize; 0];
-
     macro_rules! set_jump_raw {
         ($val:tt, $($tt:tt)*) => {
             $val = 0 as _
@@ -206,7 +203,7 @@ impl JumpPoint<'_> {
     }
 
     /// Alias of [`set_jump`].
-    #[inline(always)]
+    #[inline]
     pub fn set_jump<T, F, G>(ordinary: F, lander: G) -> T
     where
         F: FnOnce(JumpPoint<'_>) -> T,
@@ -220,7 +217,7 @@ impl JumpPoint<'_> {
     /// # Safety
     ///
     /// See [`long_jump`].
-    #[inline(always)]
+    #[inline]
     pub unsafe fn long_jump(self, result: NonZero<usize>) -> ! {
         long_jump(self, result)
     }
@@ -237,18 +234,25 @@ impl JumpPoint<'_> {
 /// Since it is implemented with [`core::mem::needs_drop`], it may generates false positive
 /// compile errors.
 ///
-/// This function is even inline-able without procedure-call cost!
+/// This function is inline-able but `ordinary` is called in a C function wrapper because it must
+/// has its own stack frame to avoid [misoptimization][misopt2].
+///
+/// [misopt2]: https://github.com/rust-lang/libc/issues/1596
 ///
 /// # Safety
 ///
 /// Yes, this function is actually safe to use. [`long_jump`] is unsafe, however.
 ///
 /// # Panics
 ///
-/// It is possible to unwind from `ordinary` or `lander` closure.
+/// It is allowed to panic in `ordinary` or `lander` closure. But panics in `ordinary` cannot
+/// unwind across `set_jump`, or it will abort the process.
+/// Panics in `lander` can unwind across `set_jump` though.
 #[doc(alias = "setjmp")]
-#[inline(always)]
-pub fn set_jump<T, F, G>(mut ordinary: F, lander: G) -> T
+#[inline]
+// For better logical order.
+#[allow(clippy::items_after_statements)]
+pub fn set_jump<T, F, G>(ordinary: F, lander: G) -> T
 where
     F: FnOnce(JumpPoint<'_>) -> T,
     G: FnOnce(NonZero<usize>) -> T,
@@ -266,40 +270,51 @@ where
         // unwind between its return from F or G and the return from `set_jump`.
     }
 
-    let mut buf = <MaybeUninit<imp::Buf>>::uninit();
-    let ptr = buf.as_mut_ptr();
-    let mut val: usize;
+    union Data<T, F> {
+        func: ManuallyDrop<F>,
+        ret: ManuallyDrop<T>,
+    }
 
-    // Show the optimizer that `ordinary` may be called in both ordinary and landing paths.
-    // So it cannot assume that variables that are captured by `ordinary` are unchanged in the
-    // lander path -- they must be reloaded.
-    // This fixes <https://github.com/rust-lang/rfcs/issues/2625>.
-    if size_of::<F>() == 0 {
-        // F must not mutate local states because it captures nothing.
-    } else {
-        unsafe {
-            core::arch::asm!(
-                "/*{}*/",
-                in(reg) core::ptr::addr_of_mut!(ordinary),
-                // May access memory.
-            );
+    struct AbortGuard;
+    impl Drop for AbortGuard {
+        fn drop(&mut self) {
+            // Manually trigger a double-panic abort.
+            // Workaround: <https://github.com/rust-lang/rust/issues/67952>
+            panic!("panic cannot unwind across set_jump");
         }
     }
 
+    unsafe extern "C" fn wrap<T, F: FnOnce(JumpPoint<'_>) -> T>(
+        data: &mut Data<T, F>,
+        jp: *mut (),
+    ) -> usize {
+        let guard = AbortGuard;
+        let f = unsafe { ManuallyDrop::take(&mut data.func) };
+        let ret = f(unsafe { JumpPoint::from_raw(jp) });
+        data.ret = ManuallyDrop::new(ret);
+        core::mem::forget(guard);
+        0
+    }
+
+    let mut data = Data::<T, F> {
+        func: ManuallyDrop::new(ordinary),
+    };
+    let val: usize;
+
     #[cfg(feature = "unstable-asm-goto")]
     unsafe {
-        set_jump_raw!(val, ptr, {
+        set_jump_raw!(val, wrap::<T, F>, &mut data, {
             unsafe { return lander(NonZero::new_unchecked(val)) }
         });
-        ordinary(JumpPoint::from_raw(ptr.cast()))
+        ManuallyDrop::take(&mut data.ret)
     }
 
     #[cfg(not(feature = "unstable-asm-goto"))]
-    unsafe {
-        set_jump_raw!(val, ptr);
+    {
+        unsafe { set_jump_raw!(val, wrap::<T, F>, &mut data) };
         match NonZero::new(val) {
-            None => ordinary(JumpPoint::from_raw(ptr.cast())),
-            Some(val) => lander(val),
+            None => unsafe { ManuallyDrop::take(&mut data.ret) },
+            Some(v) => lander(v),
         }
     }
 }