fix(jit): Add MAP_JIT support for 100% stable JIT on ARM64 macOS

cranelift/jit/src/memory/mod.rs

@@ -76,5 +76,16 @@ pub(crate) fn set_readable_and_executable(
         }
     }
 
+    // ARM64 macOS: Switch to execute mode for W^X compliance.
+    #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
+    {
+        unsafe extern "C" {
+            fn pthread_jit_write_protect_np(enabled: libc::c_int);
+        }
+        unsafe {
+            pthread_jit_write_protect_np(1);
+        }
+    }
+
     Ok(())
 }

cranelift/jit/src/memory/system.rs

@@ -3,7 +3,11 @@ use cranelift_module::{ModuleError, ModuleResult};
 #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
 use memmap2::MmapMut;
 
-#[cfg(not(any(feature = "selinux-fix", windows)))]
+#[cfg(not(any(
+    feature = "selinux-fix",
+    windows,
+    all(target_arch = "aarch64", target_os = "macos")
+)))]
 use std::alloc;
 use std::io;
 use std::mem;
@@ -49,7 +53,45 @@ impl PtrLen {
         })
     }
 
-    #[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))]
+    /// macOS ARM64: Use mmap with MAP_JIT for W^X policy compliance.
+    #[cfg(all(
+        target_arch = "aarch64",
+        target_os = "macos",
+        not(feature = "selinux-fix")
+    ))]
+    fn with_size(size: usize) -> io::Result<Self> {
+        assert_ne!(size, 0);
+        let alloc_size = region::page::ceil(size as *const ()) as usize;
+
+        const MAP_JIT: libc::c_int = 0x0800;
+
+        let ptr = unsafe {
+            libc::mmap(
+                ptr::null_mut(),
+                alloc_size,
+                libc::PROT_READ | libc::PROT_WRITE,
+                libc::MAP_PRIVATE | libc::MAP_ANON | MAP_JIT,
+                -1,
+                0,
+            )
+        };
+
+        if ptr == libc::MAP_FAILED {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(Self {
+                ptr: ptr as *mut u8,
+                len: alloc_size,
+            })
+        }
+    }
+
+    /// Non-macOS ARM64: Use standard allocator
+    #[cfg(all(
+        not(target_os = "windows"),
+        not(feature = "selinux-fix"),
+        not(all(target_arch = "aarch64", target_os = "macos"))
+    ))]
     fn with_size(size: usize) -> io::Result<Self> {
         assert_ne!(size, 0);
         let page_size = region::page::size();
@@ -95,7 +137,30 @@ impl PtrLen {
 }
 
 // `MMapMut` from `cfg(feature = "selinux-fix")` already deallocates properly.
-#[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))]
+
+/// macOS ARM64: Free MAP_JIT memory with munmap.
+#[cfg(all(
+    target_arch = "aarch64",
+    target_os = "macos",
+    not(feature = "selinux-fix")
+))]
+impl Drop for PtrLen {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe {
+                let _ = region::protect(self.ptr, self.len, region::Protection::READ_WRITE);
+                libc::munmap(self.ptr as *mut libc::c_void, self.len);
+            }
+        }
+    }
+}
+
+/// Other Unix platforms: Use standard allocator dealloc.
+#[cfg(all(
+    not(target_os = "windows"),
+    not(feature = "selinux-fix"),
+    not(all(target_arch = "aarch64", target_os = "macos"))
+))]
 impl Drop for PtrLen {
     fn drop(&mut self) {
         if !self.ptr.is_null() {

crates/jit-icache-coherence/src/libc.rs

@@ -141,16 +141,23 @@ fn riscv_flush_icache(start: u64, end: u64) -> Result<()> {
 
 pub(crate) use details::*;
 
+// macOS ARM64: Use sys_icache_invalidate for icache coherence.
+#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
+unsafe extern "C" {
+    fn sys_icache_invalidate(start: *const c_void, len: usize);
+}
+
 /// See docs on [crate::clear_cache] for a description of what this function is trying to do.
 #[inline]
 pub(crate) fn clear_cache(_ptr: *const c_void, _len: usize) -> Result<()> {
-    // TODO: On AArch64 we currently rely on the `mprotect` call that switches the memory from W+R
-    // to R+X to do this for us, however that is an implementation detail and should not be relied
-    // upon.
-    // We should call some implementation of `clear_cache` here.
-    //
-    // See: https://github.com/bytecodealliance/wasmtime/issues/3310
+    // macOS ARM64: Use sys_icache_invalidate for icache coherence.
+    #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
+    unsafe {
+        sys_icache_invalidate(_ptr, _len);
+    }
+
     #[cfg(all(target_arch = "riscv64", target_os = "linux"))]
     riscv_flush_icache(_ptr as u64, (_ptr as u64) + (_len as u64))?;
+
     Ok(())
 }