Only load setjmp/longjmp from Zig on aarch64

compiler/builtins/bitcode/src/main.zig

@@ -1,4 +1,5 @@
 const std = @import("std");
+const builtin = @import("builtin");
 const math = std.math;
 const utils = @import("utils.zig");
 const expect = @import("expect.zig");
@@ -161,6 +162,11 @@ comptime {
     exportExpectFn(expect.deinitFailuresC, "deinit_failures");
 
     @export(utils.panic, .{ .name = "roc_builtins.utils." ++ "panic", .linkage = .Weak });
+
+    if (builtin.target.cpu.arch == .aarch64) {
+        @export(__roc_force_setjmp, .{ .name = "__roc_force_setjmp", .linkage = .Weak });
+        @export(__roc_force_longjmp, .{ .name = "__roc_force_longjmp", .linkage = .Weak });
+    }
 }
 
 // Utils continued - SJLJ
@@ -177,10 +183,10 @@ pub extern fn sigsetjmp([*c]c_int, c_int) c_int;
 pub extern fn siglongjmp([*c]c_int, c_int) noreturn;
 pub extern fn longjmperror() void;
 // Zig won't expose the externs (and hence link correctly) unless we force them to be used.
-pub export fn __roc_force_setjmp(it: [*c]c_int) c_int {
+fn __roc_force_setjmp(it: [*c]c_int) callconv(.C) c_int {
     return setjmp(it);
 }
-pub export fn __roc_force_longjmp(a0: [*c]c_int, a1: c_int) noreturn {
+fn __roc_force_longjmp(a0: [*c]c_int, a1: c_int) callconv(.C) noreturn {
     longjmp(a0, a1);
 }
 
@@ -214,7 +220,6 @@ fn exportExpectFn(comptime func: anytype, comptime func_name: []const u8) void {
 
 // Custom panic function, as builtin Zig version errors during LLVM verification
 pub fn panic(message: []const u8, stacktrace: ?*std.builtin.StackTrace) noreturn {
-    const builtin = @import("builtin");
     if (builtin.is_test) {
         std.debug.print("{s}: {?}", .{ message, stacktrace });
     } else {

compiler/gen_llvm/src/llvm/build.rs

@@ -64,7 +64,7 @@ use roc_mono::ir::{
     ModifyRc, OptLevel, ProcLayout,
 };
 use roc_mono::layout::{Builtin, LambdaSet, Layout, LayoutIds, TagIdIntType, UnionLayout};
-use roc_target::TargetInfo;
+use roc_target::{PtrWidth, TargetInfo};
 use target_lexicon::{Architecture, OperatingSystem, Triple};
 
 /// This is for Inkwell's FunctionValue::verify - we want to know the verification
@@ -516,11 +516,36 @@ fn add_intrinsics<'ctx>(ctx: &'ctx Context, module: &Module<'ctx>) {
     //
     // https://releases.llvm.org/10.0.0/docs/LangRef.html#standard-c-library-intrinsics
     let i1_type = ctx.bool_type();
+    let i8_type = ctx.i8_type();
+    let i8_ptr_type = i8_type.ptr_type(AddressSpace::Generic);
+    let i32_type = ctx.i32_type();
+    let void_type = ctx.void_type();
 
     if let Some(func) = module.get_function("__muloti4") {
         func.set_linkage(Linkage::WeakAny);
     }
 
+    add_intrinsic(
+        ctx,
+        module,
+        LLVM_SETJMP,
+        i32_type.fn_type(&[i8_ptr_type.into()], false),
+    );
+
+    add_intrinsic(
+        ctx,
+        module,
+        LLVM_LONGJMP,
+        void_type.fn_type(&[i8_ptr_type.into()], false),
+    );
+
+    add_intrinsic(
+        ctx,
+        module,
+        LLVM_FRAME_ADDRESS,
+        i8_ptr_type.fn_type(&[i32_type.into()], false),
+    );
+
     add_float_intrinsic(ctx, module, &LLVM_LOG, |t| t.fn_type(&[t.into()], false));
     add_float_intrinsic(ctx, module, &LLVM_POW, |t| {
         t.fn_type(&[t.into(), t.into()], false)
@@ -573,6 +598,11 @@ static LLVM_FLOOR: IntrinsicName = float_intrinsic!("llvm.floor");
 static LLVM_MEMSET_I64: &str = "llvm.memset.p0i8.i64";
 static LLVM_MEMSET_I32: &str = "llvm.memset.p0i8.i32";
 
+static LLVM_FRAME_ADDRESS: &str = "llvm.frameaddress.p0i8";
+
+static LLVM_SETJMP: &str = "llvm.eh.sjlj.setjmp";
+pub static LLVM_LONGJMP: &str = "llvm.eh.sjlj.longjmp";
+
 const LLVM_ADD_WITH_OVERFLOW: IntrinsicName =
     llvm_int_intrinsic!("llvm.sadd.with.overflow", "llvm.uadd.with.overflow");
 const LLVM_SUB_WITH_OVERFLOW: IntrinsicName =
@@ -3620,14 +3650,21 @@ fn expose_function_to_host_help_c_abi<'a, 'ctx, 'env>(
 
 pub fn get_sjlj_buffer<'a, 'ctx, 'env>(env: &Env<'a, 'ctx, 'env>) -> PointerValue<'ctx> {
     // The size of jump_buf is platform-dependent.
-    // AArch64 needs 3 64-bit words, x86 seemingly needs less than that. Let's just make it 3
+    //   - AArch64 needs 3 machine-sized words
-    // 64-bit words.
+    //   - LLVM says the following about the SJLJ intrinsic:
-    // We can always increase the size, since we hand it off opaquely to both `setjmp` and `longjmp`.
+    //
-    // The unused space won't be touched.
+    //     [It is] a five word buffer in which the calling context is saved.
-    let type_ = env
+    //     The front end places the frame pointer in the first word, and the
-        .context
+    //     target implementation of this intrinsic should place the destination
-        .i32_type()
+    //     address for a llvm.eh.sjlj.longjmp in the second word.
-        .array_type(6 * env.target_info.ptr_width() as u32);
+    //     The following three words are available for use in a target-specific manner.
+    //
+    // So, let's create a 5-word buffer.
+    let word_type = match env.target_info.ptr_width() {
+        PtrWidth::Bytes4 => env.context.i32_type(),
+        PtrWidth::Bytes8 => env.context.i64_type(),
+    };
+    let type_ = word_type.array_type(5);
 
     let global = match env.module.get_global("roc_sjlj_buffer") {
         Some(global) => global,
@@ -3645,6 +3682,44 @@ pub fn get_sjlj_buffer<'a, 'ctx, 'env>(env: &Env<'a, 'ctx, 'env>) -> PointerValu
         .into_pointer_value()
 }
 
+pub fn build_setjmp_call<'a, 'ctx, 'env>(env: &Env<'a, 'ctx, 'env>) -> BasicValueEnum<'ctx> {
+    let jmp_buf = get_sjlj_buffer(env);
+    if cfg!(target_arch = "aarch64") {
+        // Due to https://github.com/rtfeldman/roc/issues/2965, we use a setjmp we linked in from Zig
+        call_bitcode_fn(env, &[jmp_buf.into()], bitcode::UTILS_SETJMP)
+    } else {
+        // Anywhere else, use the LLVM intrinsic.
+        let jmp_buf_i8p = env
+            .builder
+            .build_bitcast(
+                jmp_buf,
+                env.context
+                    .i8_type()
+                    .ptr_type(AddressSpace::Generic)
+                    .array_type(5)
+                    .ptr_type(AddressSpace::Generic),
+                "jmp_buf [5 x i8*]",
+            )
+            .into_pointer_value();
+        // LLVM asks us to please store the frame pointer in the first word.
+        let frame_address = env.call_intrinsic(
+            LLVM_FRAME_ADDRESS,
+            &[env.context.i32_type().const_zero().into()],
+        );
+
+        let zero = env.context.i32_type().const_zero();
+        let fa_index = env.context.i32_type().const_zero();
+        let fa = unsafe {
+            env.builder
+                .build_in_bounds_gep(jmp_buf_i8p, &[zero, fa_index], "frame address index")
+        };
+
+        env.builder.build_store(fa, frame_address);
+
+        env.call_intrinsic(LLVM_SETJMP, &[jmp_buf_i8p.into()])
+    }
+}
+
 /// Pointer to pointer of the panic message.
 pub fn get_panic_msg_ptr<'a, 'ctx, 'env>(env: &Env<'a, 'ctx, 'env>) -> PointerValue<'ctx> {
     let ptr_to_u8_ptr = env.context.i8_type().ptr_type(AddressSpace::Generic);
@@ -3677,9 +3752,7 @@ fn set_jump_and_catch_long_jump<'a, 'ctx, 'env>(
     let catch_block = context.append_basic_block(parent, "catch_block");
     let cont_block = context.append_basic_block(parent, "cont_block");
 
-    let buffer = get_sjlj_buffer(env);
+    let panicked_u32 = build_setjmp_call(env);
-
-    let panicked_u32 = call_bitcode_fn(env, &[buffer.into()], bitcode::UTILS_SETJMP);
     let panicked_bool = env.builder.build_int_compare(
         IntPredicate::NE,
         panicked_u32.into_int_value(),

compiler/gen_llvm/src/llvm/externs.rs

@@ -7,6 +7,8 @@ use inkwell::values::BasicValue;
 use inkwell::AddressSpace;
 use roc_builtins::bitcode;
 
+use super::build::{get_sjlj_buffer, LLVM_LONGJMP};
+
 /// Define functions for roc_alloc, roc_realloc, and roc_dealloc
 /// which use libc implementations (malloc, realloc, and free)
 pub fn add_default_roc_externs(env: &Env<'_, '_, '_>) {
@@ -209,13 +211,10 @@ pub fn add_sjlj_roc_panic(env: &Env<'_, '_, '_>) {
 
         builder.position_at_end(entry);
 
-        let buffer = crate::llvm::build::get_sjlj_buffer(env);
-
         // write our error message pointer
         env.builder.build_store(get_panic_msg_ptr(env), ptr_arg);
 
-        let tag = env.context.i32_type().const_int(1, false);
+        build_longjmp_call(env);
-        let _call = call_void_bitcode_fn(env, &[buffer.into(), tag.into()], bitcode::UTILS_LONGJMP);
 
         builder.build_unreachable();
 
@@ -224,3 +223,21 @@ pub fn add_sjlj_roc_panic(env: &Env<'_, '_, '_>) {
         }
     }
 }
+
+pub fn build_longjmp_call<'a, 'ctx, 'env>(env: &Env<'a, 'ctx, 'env>) {
+    let jmp_buf = get_sjlj_buffer(env);
+    if cfg!(target_arch = "aarch64") {
+        // Call the Zig-linked longjmp: `void longjmp(i32*, i32)`
+        let tag = env.context.i32_type().const_int(1, false);
+        let _call =
+            call_void_bitcode_fn(env, &[jmp_buf.into(), tag.into()], bitcode::UTILS_LONGJMP);
+    } else {
+        // Call the LLVM-intrinsic longjmp: `void @llvm.eh.sjlj.longjmp(i8* %setjmp_buf)`
+        let jmp_buf_i8p = env.builder.build_bitcast(
+            jmp_buf,
+            env.context.i8_type().ptr_type(AddressSpace::Generic),
+            "jmp_buf i8*",
+        );
+        let _call = env.build_intrinsic_call(LLVM_LONGJMP, &[jmp_buf_i8p.into()]);
+    }
+}