Add control flow information to __rust_probestack (#328)

src/lib.rs

@@ -1,6 +1,7 @@
 #![cfg_attr(feature = "compiler-builtins", compiler_builtins)]
 #![feature(abi_unadjusted)]
 #![feature(asm)]
+#![feature(global_asm)]
 #![feature(cfg_target_has_atomic)]
 #![feature(compiler_builtins)]
 #![feature(core_intrinsics)]

src/probestack.rs

@@ -41,20 +41,68 @@
 //! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
 //! be more than welcome to accept such a change!
 
-#![cfg(not(windows))] // Windows already has builtins to do this
+#![cfg(not(feature = "mangled-names"))]
+// Windows already has builtins to do this.
+#![cfg(not(windows))]
+// We only define stack probing for these architectures today.
+#![cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 
-#[naked]
+extern "C" {
-#[no_mangle]
+    pub fn __rust_probestack();
-#[cfg(all(target_arch = "x86_64", not(feature = "mangled-names")))]
+}
-pub unsafe extern "C" fn __rust_probestack() {
+
-    // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
+// A wrapper for our implementation of __rust_probestack, which allows us to
-    // ensuring that if any pages are unmapped we'll make a page fault.
+// keep the assembly inline while controlling all CFI directives in the assembly
-    //
+// emitted for the function.
-    // The ABI here is that the stack frame size is located in `%eax`. Upon
+//
-    // return we're not supposed to modify `%esp` or `%eax`.
+// This is the ELF version.
-    asm!("
+#[cfg(not(target_vendor = "apple"))]
+macro_rules! define_rust_probestack {
+    ($body: expr) => {
+        concat!(
+            "
+            .pushsection .text.__rust_probestack
+            .globl __rust_probestack
+            .type  __rust_probestack, @function
+        __rust_probestack:
+            ",
+            $body,
+            "
+            .size __rust_probestack, . - __rust_probestack
+            .popsection
+            "
+        )
+    };
+}
+
+// Same as above, but for Mach-O.
+#[cfg(target_vendor = "apple")]
+macro_rules! define_rust_probestack {
+    ($body: expr) => {
+        concat!(
+            "
+            .globl ___rust_probestack
+        ___rust_probestack:
+            ",
+            $body
+        )
+    };
+}
+
+// Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
+// ensuring that if any pages are unmapped we'll make a page fault.
+//
+// The ABI here is that the stack frame size is located in `%rax`. Upon
+// return we're not supposed to modify `%rsp` or `%rax`.
+#[cfg(target_arch = "x86_64")]
+global_asm!(define_rust_probestack!(
+    "
+    .cfi_startproc
     pushq  %rbp
+    .cfi_adjust_cfa_offset 8
+    .cfi_offset %rbp, -16
     movq   %rsp, %rbp
+    .cfi_def_cfa_register %rbp
 
     mov    %rax,%r11        // duplicate %rax as we're clobbering %r11
 
@@ -72,16 +120,16 @@ pub unsafe extern "C" fn __rust_probestack() {
     // Dynamic stack allocation, which is needed to implement unsized
     // rvalues, triggers stackprobe even if %rax < 0x1000.
     // Thus we have to check %r11 first to avoid segfault.
-        cmp    $$0x1000,%r11
+    cmp    $0x1000,%r11
     jna    3f
-    2:
+2:
-        sub    $$0x1000,%rsp
+    sub    $0x1000,%rsp
     test   %rsp,8(%rsp)
-        sub    $$0x1000,%r11
+    sub    $0x1000,%r11
-        cmp    $$0x1000,%r11
+    cmp    $0x1000,%r11
     ja     2b
 
-    3:
+3:
     // Finish up the last remaining stack space requested, getting the last
     // bits out of r11
     sub    %r11,%rsp
@@ -93,43 +141,49 @@ pub unsafe extern "C" fn __rust_probestack() {
     add    %rax,%rsp
 
     leave
+    .cfi_def_cfa_register %rsp
+    .cfi_adjust_cfa_offset -8
     ret
-    " ::: "memory" : "volatile");
+    .cfi_endproc
-    ::core::intrinsics::unreachable();
+    "
-}
+));
 
-#[naked]
+#[cfg(target_arch = "x86")]
-#[no_mangle]
+// This is the same as x86_64 above, only translated for 32-bit sizes. Note
-#[cfg(all(target_arch = "x86", not(feature = "mangled-names")))]
+// that on Unix we're expected to restore everything as it was, this
-pub unsafe extern "C" fn __rust_probestack() {
+// function basically can't tamper with anything.
-    // This is the same as x86_64 above, only translated for 32-bit sizes. Note
+//
-    // that on Unix we're expected to restore everything as it was, this
+// The ABI here is the same as x86_64, except everything is 32-bits large.
-    // function basically can't tamper with anything.
+global_asm!(define_rust_probestack!(
-    //
+    "
-    // The ABI here is the same as x86_64, except everything is 32-bits large.
+    .cfi_startproc
-    asm!("
     push   %ebp
+    .cfi_adjust_cfa_offset 4
+    .cfi_offset %ebp, -8
     mov    %esp, %ebp
+    .cfi_def_cfa_register %ebp
     push   %ecx
     mov    %eax,%ecx
 
-        cmp    $$0x1000,%ecx
+    cmp    $0x1000,%ecx
     jna    3f
-    2:
+2:
-        sub    $$0x1000,%esp
+    sub    $0x1000,%esp
     test   %esp,8(%esp)
-        sub    $$0x1000,%ecx
+    sub    $0x1000,%ecx
-        cmp    $$0x1000,%ecx
+    cmp    $0x1000,%ecx
     ja     2b
 
-    3:
+3:
     sub    %ecx,%esp
     test   %esp,8(%esp)
 
     add    %eax,%esp
     pop    %ecx
     leave
+    .cfi_def_cfa_register %esp
+    .cfi_adjust_cfa_offset -4
     ret
-    " ::: "memory" : "volatile");
+    .cfi_endproc
-    ::core::intrinsics::unreachable();
+    "
-}
+));