Rollup merge of rust-lang#67502 - Mark-Simulacrum:opt-catch, r=alexcrichton

Optimize catch_unwind to match C++ try/catch

This refactors the implementation of catching unwinds to allow LLVM to inline the "try" closure directly into the happy path, avoiding indirection. This means that the catch_unwind implementation is (after this PR) zero-cost unless a panic is thrown.

https://rust.godbolt.org/z/cZcUSB is an example of the current codegen in a simple case. Notably, the codegen is *exactly the same* if `-Cpanic=abort` is passed, which is clearly not great.

This PR, on the other hand, generates the following assembly:

```asm
# -Cpanic=unwind:
	push   rbx
	mov    ebx,0x2a
	call   QWORD PTR [rip+0x1c53c]        # <happy>
	mov    eax,ebx
	pop    rbx
	ret
	mov    rdi,rax
	call   QWORD PTR [rip+0x1c537]        # cleanup function call
	call   QWORD PTR [rip+0x1c539]        # <unfortunate>
	mov    ebx,0xd
	mov    eax,ebx
	pop    rbx
	ret

# -Cpanic=abort:
	push   rax
	call   QWORD PTR [rip+0x20a1]        # <happy>
	mov    eax,0x2a
	pop    rcx
	ret
```

Fixes rust-lang#64224, and resolves rust-lang#64222.

Author: Centril

Cargo.lock

@@ -2301,6 +2301,7 @@ dependencies = [
 name = "panic_abort"
 version = "0.0.0"
 dependencies = [
+ "cfg-if",
  "compiler_builtins",
  "core",
  "libc",

src/libpanic_abort/Cargo.toml

@@ -14,3 +14,4 @@ doc = false
 core = { path = "../libcore" }
 libc = { version = "0.2", default-features = false }
 compiler_builtins = "0.1.0"
+cfg-if = "0.1.8"

src/libpanic_abort/lib.rs

@@ -18,17 +18,14 @@
 #![feature(staged_api)]
 #![feature(rustc_attrs)]
 
-// Rust's "try" function, but if we're aborting on panics we just call the
-// function as there's nothing else we need to do here.
+use core::any::Any;
+
+// We need the definition of TryPayload for __rust_panic_cleanup.
+include!("../libpanic_unwind/payload.rs");
+
 #[rustc_std_internal_symbol]
-pub unsafe extern "C" fn __rust_maybe_catch_panic(
-    f: fn(*mut u8),
-    data: *mut u8,
-    _data_ptr: *mut usize,
-    _vtable_ptr: *mut usize,
-) -> u32 {
-    f(data);
-    0
+pub unsafe extern "C" fn __rust_panic_cleanup(_: TryPayload) -> *mut (dyn Any + Send + 'static) {
+    unreachable!()
 }
 
 // "Leak" the payload and shim to the relevant abort on the platform in

src/libpanic_unwind/dummy.rs

@@ -6,10 +6,6 @@ use alloc::boxed::Box;
 use core::any::Any;
 use core::intrinsics;
 
-pub fn payload() -> *mut u8 {
-    core::ptr::null_mut()
-}
-
 pub unsafe fn cleanup(_ptr: *mut u8) -> Box<dyn Any + Send> {
     intrinsics::abort()
 }

src/libpanic_unwind/emcc.rs

@@ -48,10 +48,6 @@ static EXCEPTION_TYPE_INFO: TypeInfo = TypeInfo {
     name: b"rust_panic\0".as_ptr(),
 };
 
-pub fn payload() -> *mut u8 {
-    ptr::null_mut()
-}
-
 pub unsafe fn cleanup(ptr: *mut u8) -> Box<dyn Any + Send> {
     assert!(!ptr.is_null());
     let adjusted_ptr = __cxa_begin_catch(ptr as *mut libc::c_void);

src/libpanic_unwind/gcc.rs

@@ -48,7 +48,6 @@
 
 use alloc::boxed::Box;
 use core::any::Any;
-use core::ptr;
 
 use crate::dwarf::eh::{self, EHAction, EHContext};
 use libc::{c_int, uintptr_t};
@@ -82,10 +81,6 @@ pub unsafe fn panic(data: Box<dyn Any + Send>) -> u32 {
     }
 }
 
-pub fn payload() -> *mut u8 {
-    ptr::null_mut()
-}
-
 pub unsafe fn cleanup(ptr: *mut u8) -> Box<dyn Any + Send> {
     let my_ep = ptr as *mut Exception;
     let cause = (*my_ep).cause.take();

src/libpanic_unwind/hermit.rs

@@ -6,10 +6,6 @@ use alloc::boxed::Box;
 use core::any::Any;
 use core::ptr;
 
-pub fn payload() -> *mut u8 {
-    ptr::null_mut()
-}
-
 pub unsafe fn cleanup(_ptr: *mut u8) -> Box<dyn Any + Send> {
     extern "C" {
         pub fn __rust_abort() -> !;

src/libpanic_unwind/lib.rs

@@ -22,19 +22,20 @@
 #![feature(libc)]
 #![feature(nll)]
 #![feature(panic_unwind)]
-#![feature(raw)]
 #![feature(staged_api)]
 #![feature(std_internals)]
 #![feature(unwind_attributes)]
+#![feature(rustc_attrs)]
+#![feature(raw)]
 #![panic_runtime]
 #![feature(panic_runtime)]
 
 use alloc::boxed::Box;
-use core::intrinsics;
-use core::mem;
+use core::any::Any;
 use core::panic::BoxMeUp;
-use core::raw;
 
+// If adding to this list, you should also look at the list of TryPayload types
+// defined in payload.rs and likely add to there as well.
 cfg_if::cfg_if! {
     if #[cfg(target_os = "emscripten")] {
         #[path = "emcc.rs"]
@@ -60,30 +61,15 @@ cfg_if::cfg_if! {
     }
 }
 
+include!("payload.rs");
+
 mod dwarf;
 
-// Entry point for catching an exception, implemented using the `try` intrinsic
-// in the compiler.
-//
-// The interaction between the `payload` function and the compiler is pretty
-// hairy and tightly coupled, for more information see the compiler's
-// implementation of this.
 #[no_mangle]
-pub unsafe extern "C" fn __rust_maybe_catch_panic(
-    f: fn(*mut u8),
-    data: *mut u8,
-    data_ptr: *mut usize,
-    vtable_ptr: *mut usize,
-) -> u32 {
-    let mut payload = imp::payload();
-    if intrinsics::r#try(f, data, &mut payload as *mut _ as *mut _) == 0 {
-        0
-    } else {
-        let obj = mem::transmute::<_, raw::TraitObject>(imp::cleanup(payload));
-        *data_ptr = obj.data as usize;
-        *vtable_ptr = obj.vtable as usize;
-        1
-    }
+pub unsafe extern "C" fn __rust_panic_cleanup(
+    payload: TryPayload,
+) -> *mut (dyn Any + Send + 'static) {
+    Box::into_raw(imp::cleanup(payload))
 }
 
 // Entry point for raising an exception, just delegates to the platform-specific

src/libpanic_unwind/payload.rs

@@ -0,0 +1,21 @@
+// Type definition for the payload argument of the try intrinsic.
+//
+// This must be kept in sync with the implementations of the try intrinsic.
+//
+// This file is included by both panic runtimes and libstd. It is part of the
+// panic runtime ABI.
+cfg_if::cfg_if! {
+    if #[cfg(target_os = "emscripten")] {
+        type TryPayload = *mut u8;
+    } else if #[cfg(target_arch = "wasm32")] {
+        type TryPayload = *mut u8;
+    } else if #[cfg(target_os = "hermit")] {
+        type TryPayload = *mut u8;
+    } else if #[cfg(all(target_env = "msvc", target_arch = "aarch64"))] {
+        type TryPayload = *mut u8;
+    } else if #[cfg(target_env = "msvc")] {
+        type TryPayload = [u64; 2];
+    } else {
+        type TryPayload = *mut u8;
+    }
+}

src/libpanic_unwind/seh.rs

@@ -264,10 +264,6 @@ pub unsafe fn panic(data: Box<dyn Any + Send>) -> u32 {
     _CxxThrowException(throw_ptr, &mut THROW_INFO as *mut _ as *mut _);
 }
 
-pub fn payload() -> [u64; 2] {
-    [0; 2]
-}
-
 pub unsafe fn cleanup(payload: [u64; 2]) -> Box<dyn Any + Send> {
     mem::transmute(raw::TraitObject { data: payload[0] as *mut _, vtable: payload[1] as *mut _ })
 }

src/librustc_codegen_llvm/intrinsic.rs

@@ -858,8 +858,10 @@ fn try_intrinsic(
 ) {
     if bx.sess().no_landing_pads() {
         bx.call(func, &[data], None);
-        let ptr_align = bx.tcx().data_layout.pointer_align.abi;
-        bx.store(bx.const_null(bx.type_i8p()), dest, ptr_align);
+        // Return 0 unconditionally from the intrinsic call;
+        // we can never unwind.
+        let ret_align = bx.tcx().data_layout.i32_align.abi;
+        bx.store(bx.const_i32(0), dest, ret_align);
     } else if wants_msvc_seh(bx.sess()) {
         codegen_msvc_try(bx, func, data, local_ptr, dest);
     } else {

src/libstd/panicking.rs

@@ -12,9 +12,8 @@ use core::panic::{BoxMeUp, Location, PanicInfo};
 use crate::any::Any;
 use crate::fmt;
 use crate::intrinsics;
-use crate::mem::{self, ManuallyDrop};
+use crate::mem::{self, ManuallyDrop, MaybeUninit};
 use crate::process;
-use crate::raw;
 use crate::sync::atomic::{AtomicBool, Ordering};
 use crate::sys::stdio::panic_output;
 use crate::sys_common::backtrace::{self, RustBacktrace};
@@ -29,6 +28,9 @@ use crate::io::set_panic;
 #[cfg(test)]
 use realstd::io::set_panic;
 
+// Include the definition of UnwindPayload from libpanic_unwind.
+include!("../libpanic_unwind/payload.rs");
+
 // Binary interface to the panic runtime that the standard library depends on.
 //
 // The standard library is tagged with `#![needs_panic_runtime]` (introduced in
@@ -41,12 +43,9 @@ use realstd::io::set_panic;
 // hook up these functions, but it is not this day!
 #[allow(improper_ctypes)]
 extern "C" {
-    fn __rust_maybe_catch_panic(
-        f: fn(*mut u8),
-        data: *mut u8,
-        data_ptr: *mut usize,
-        vtable_ptr: *mut usize,
-    ) -> u32;
+    /// The payload ptr here is actually the same as the payload ptr for the try
+    /// intrinsic (i.e., is really `*mut [u64; 2]` or `*mut *mut u8`).
+    fn __rust_panic_cleanup(payload: TryPayload) -> *mut (dyn Any + Send + 'static);
 
     /// `payload` is actually a `*mut &mut dyn BoxMeUp` but that would cause FFI warnings.
     /// It cannot be `Box<dyn BoxMeUp>` because the other end of this call does not depend
@@ -241,9 +240,9 @@ pub unsafe fn r#try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<dyn Any + Send>>
     }
 
     // We do some sketchy operations with ownership here for the sake of
-    // performance. We can only  pass pointers down to
-    // `__rust_maybe_catch_panic` (can't pass objects by value), so we do all
-    // the ownership tracking here manually using a union.
+    // performance. We can only pass pointers down to `do_call` (can't pass
+    // objects by value), so we do all the ownership tracking here manually
+    // using a union.
     //
     // We go through a transition where:
     //
@@ -254,7 +253,7 @@ pub unsafe fn r#try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<dyn Any + Send>>
     // * If the closure successfully returns, we write the return value into the
     //   data's return slot. Note that `ptr::write` is used as it's overwriting
     //   uninitialized data.
-    // * Finally, when we come back out of the `__rust_maybe_catch_panic` we're
+    // * Finally, when we come back out of the `try` intrinsic we're
     //   in one of two states:
     //
     //      1. The closure didn't panic, in which case the return value was
@@ -265,28 +264,30 @@ pub unsafe fn r#try<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<dyn Any + Send>>
     //
     // Once we stack all that together we should have the "most efficient'
     // method of calling a catch panic whilst juggling ownership.
-    let mut any_data = 0;
-    let mut any_vtable = 0;
     let mut data = Data { f: ManuallyDrop::new(f) };
 
-    let r = __rust_maybe_catch_panic(
-        do_call::<F, R>,
-        &mut data as *mut _ as *mut u8,
-        &mut any_data,
-        &mut any_vtable,
-    );
+    let mut payload: MaybeUninit<TryPayload> = MaybeUninit::uninit();
 
-    return if r == 0 {
+    let data_ptr = &mut data as *mut _ as *mut u8;
+    let payload_ptr = payload.as_mut_ptr() as *mut _;
+    return if intrinsics::r#try(do_call::<F, R>, data_ptr, payload_ptr) == 0 {
         debug_assert!(update_panic_count(0) == 0);
         Ok(ManuallyDrop::into_inner(data.r))
     } else {
+        Err(cleanup(payload.assume_init()))
+    };
+
+    // We consider unwinding to be rare, so mark this function as cold. However,
+    // do not mark it no-inline -- that decision is best to leave to the
+    // optimizer (in most cases this function is not inlined even as a normal,
+    // non-cold function, though, as of the writing of this comment).
+    #[cold]
+    unsafe fn cleanup(payload: TryPayload) -> Box<dyn Any + Send + 'static> {
+        let obj = Box::from_raw(__rust_panic_cleanup(payload));
         update_panic_count(-1);
         debug_assert!(update_panic_count(0) == 0);
-        Err(mem::transmute(raw::TraitObject {
-            data: any_data as *mut _,
-            vtable: any_vtable as *mut _,
-        }))
-    };
+        obj
+    }
 
     fn do_call<F: FnOnce() -> R, R>(data: *mut u8) {
         unsafe {

src/test/codegen/catch-unwind.rs

@@ -0,0 +1,19 @@
+// compile-flags: -O
+
+#![crate_type = "lib"]
+
+extern "C" {
+    fn bar();
+}
+
+// CHECK-LABEL: @foo
+#[no_mangle]
+pub unsafe fn foo() -> i32 {
+    // CHECK: call void @bar
+    // CHECK: ret i32 0
+    std::panic::catch_unwind(|| {
+        bar();
+        0
+    })
+    .unwrap()
+}

src/test/codegen/try-panic-abort.rs

@@ -0,0 +1,17 @@
+// compile-flags: -C panic=abort -O
+
+#![crate_type = "lib"]
+#![feature(unwind_attributes, core_intrinsics)]
+
+extern "C" {
+    #[unwind(allow)]
+    fn bar(data: *mut u8);
+}
+
+// CHECK-LABEL: @foo
+#[no_mangle]
+pub unsafe fn foo() -> i32 {
+    // CHECK: call void @bar
+    // CHECK: ret i32 0
+    std::intrinsics::r#try(|x| bar(x), 0 as *mut u8, 0 as *mut u8)
+}

src/test/ui/no-landing-pads.rs

@@ -59,6 +59,8 @@ const EXCEPTION_PATHS: &[&str] = &[
     "src/libstd/sys_common/mod.rs",
     "src/libstd/sys_common/net.rs",
     "src/libstd/sys_common/backtrace.rs",
+    // panic_unwind shims
+    "src/libstd/panicking.rs",
     "src/libterm", // Not sure how to make this crate portable, but test crate needs it.
     "src/libtest", // Probably should defer to unstable `std::sys` APIs.
     "src/libstd/sync/mpsc", // some tests are only run on non-emscripten