feat(jit): transfer owned side-entry state

Author: fffonion

src/vm/jit/native/lower.rs

@@ -222,6 +222,122 @@ pub(crate) fn compile_system_tail_wrapper(root_entry: *const u8) -> VmResult<Com
     )
 }
 
+fn tail_owned_entry_signature(pointer_type: cranelift_codegen::ir::Type) -> Signature {
+    let mut signature = tail_entry_signature(pointer_type);
+    signature.params.push(AbiParam::new(pointer_type));
+    signature
+}
+
+fn system_owned_entry_signature(
+    pointer_type: cranelift_codegen::ir::Type,
+    call_conv: CallConv,
+) -> Signature {
+    let mut signature = entry_signature(pointer_type, call_conv);
+    signature.params.push(AbiParam::new(pointer_type));
+    signature
+}
+
+pub(crate) fn compile_tail_owned_side_link_body(
+    slot_address: usize,
+    deopt_status: i32,
+) -> VmResult<CompiledTailFunction> {
+    compile_standalone_native_function(
+        "pd_vm_tail_owned_side_link",
+        |pointer_type, _| tail_owned_entry_signature(pointer_type),
+        move |builder, pointer_type, _| {
+            let entry = builder.create_block();
+            let deopt = builder.create_block();
+            let linked = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.switch_to_block(entry);
+            let vm_ptr = builder.block_params(entry)[0];
+            let owned_slot = builder.block_params(entry)[1];
+            let slot_address = iconst_ptr_from_addr(builder, pointer_type, slot_address)?;
+            let target = builder
+                .ins()
+                .load(pointer_type, MemFlags::new(), slot_address, 0);
+            let is_null = builder.ins().icmp_imm(IntCC::Equal, target, 0);
+            builder.ins().brif(is_null, deopt, &[], linked, &[]);
+
+            builder.switch_to_block(deopt);
+            let status = builder.ins().iconst(types::I32, i64::from(deopt_status));
+            builder.ins().return_(&[status]);
+
+            builder.switch_to_block(linked);
+            let signature = builder.import_signature(tail_owned_entry_signature(pointer_type));
+            builder
+                .ins()
+                .return_call_indirect(signature, target, &[vm_ptr, owned_slot]);
+            Ok(())
+        },
+    )
+}
+
+pub(crate) fn compile_tail_owned_clear_body(success_status: i32) -> VmResult<CompiledTailFunction> {
+    compile_standalone_native_function(
+        "pd_vm_tail_owned_clear",
+        |pointer_type, _| tail_owned_entry_signature(pointer_type),
+        move |builder, pointer_type, call_conv| {
+            let entry = builder.create_block();
+            let success = builder.create_block();
+            let failure = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.append_block_param(failure, types::I32);
+            builder.switch_to_block(entry);
+            let owned_slot = builder.block_params(entry)[1];
+            let helper =
+                iconst_ptr_from_addr(builder, pointer_type, clear_value_slot_entry_address())?;
+            let helper_signature =
+                builder.import_signature(value_slot_signature(pointer_type, call_conv));
+            let call = builder
+                .ins()
+                .call_indirect(helper_signature, helper, &[owned_slot]);
+            let helper_status = builder.inst_results(call)[0];
+            let succeeded =
+                builder
+                    .ins()
+                    .icmp_imm(IntCC::Equal, helper_status, i64::from(STATUS_CONTINUE));
+            builder
+                .ins()
+                .brif(succeeded, success, &[], failure, &[helper_status.into()]);
+
+            builder.switch_to_block(success);
+            let status = builder.ins().iconst(types::I32, i64::from(success_status));
+            builder.ins().return_(&[status]);
+
+            builder.switch_to_block(failure);
+            let status = builder.block_params(failure)[0];
+            builder.ins().return_(&[status]);
+            Ok(())
+        },
+    )
+}
+
+pub(crate) fn compile_system_owned_tail_wrapper(
+    root_entry: *const u8,
+) -> VmResult<CompiledTailFunction> {
+    let root_entry = root_entry as usize;
+    compile_standalone_native_function(
+        "pd_vm_tail_owned_wrapper",
+        system_owned_entry_signature,
+        move |builder, pointer_type, _| {
+            let entry = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.switch_to_block(entry);
+            let vm_ptr = builder.block_params(entry)[0];
+            let owned_slot = builder.block_params(entry)[1];
+            let root_entry = iconst_ptr_from_addr(builder, pointer_type, root_entry)?;
+            let signature = builder.import_signature(tail_owned_entry_signature(pointer_type));
+            let call = builder
+                .ins()
+                .call_indirect(signature, root_entry, &[vm_ptr, owned_slot]);
+            let status = builder.inst_results(call)[0];
+            builder.ins().return_(&[status]);
+            Ok(())
+        },
+    )
+}
+
 fn try_compile_ssa_trace(
     trace: &JitTrace,
     ssa: &SsaTrace,

src/vm/jit/native/mod.rs

@@ -188,14 +188,47 @@ pub(super) fn compile_native_region(
 #[cfg(test)]
 mod tests {
     use super::lower::{
-        compile_system_tail_wrapper, compile_tail_side_link_body, compile_tail_status_body,
+        compile_system_owned_tail_wrapper, compile_system_tail_wrapper,
+        compile_tail_owned_clear_body, compile_tail_owned_side_link_body,
+        compile_tail_side_link_body, compile_tail_status_body,
     };
     use super::{
         InheritedStateAbiClass, NativeSideLinkSlot, SideEntryOwnership,
         classify_side_entry_materialization, classify_side_entry_repr, selected_codegen_backend,
     };
-    use crate::ValueType;
     use crate::vm::jit::ir::{SsaMaterialization, SsaValueId, SsaValueRepr};
+    use crate::{Value, ValueType};
+    use std::sync::Arc;
+
+    #[cfg(target_os = "linux")]
+    fn assert_executable_mapping_is_not_writable(entry: *const u8) {
+        let address = entry as usize;
+        let maps = std::fs::read_to_string("/proc/self/maps").expect("process maps should read");
+        let mapping = maps
+            .lines()
+            .find(|line| {
+                let Some((range, _)) = line.split_once(' ') else {
+                    return false;
+                };
+                let Some((start, end)) = range.split_once('-') else {
+                    return false;
+                };
+                let Ok(start) = usize::from_str_radix(start, 16) else {
+                    return false;
+                };
+                let Ok(end) = usize::from_str_radix(end, 16) else {
+                    return false;
+                };
+                start <= address && address < end
+            })
+            .expect("entry mapping should exist");
+        let permissions = mapping
+            .split_whitespace()
+            .nth(1)
+            .expect("mapping permissions should exist");
+        assert!(permissions.contains('x'), "{mapping}");
+        assert!(!permissions.contains('w'), "{mapping}");
+    }
 
     #[test]
     fn side_entry_abi_classifies_scalar_pointer_tagged_borrowed_and_owned_values() {
@@ -252,6 +285,8 @@ mod tests {
             .expect("tail root should compile");
         let wrapper =
             compile_system_tail_wrapper(root.entry()).expect("system wrapper should compile");
+        #[cfg(target_os = "linux")]
+        assert_executable_mapping_is_not_writable(wrapper.entry());
         assert!(child.code_len() > 0);
         assert!(root.code_len() > 0);
         assert!(wrapper.code_len() > 0);
@@ -271,6 +306,53 @@ mod tests {
         assert_eq!(unsafe { entry(std::ptr::null_mut()) }, DEOPT_STATUS);
     }
 
+    #[cfg(feature = "cranelift-jit")]
+    #[test]
+    fn trace_jit_side_entry_transfers_owned_values_once() {
+        if selected_codegen_backend() != "native" {
+            return;
+        }
+        const DEOPT_STATUS: i32 = 29;
+        const CHILD_STATUS: i32 = 31;
+        let slot = Box::new(NativeSideLinkSlot::new());
+        let child =
+            compile_tail_owned_clear_body(CHILD_STATUS).expect("owned child should compile");
+        let root = compile_tail_owned_side_link_body(slot.address() as usize, DEOPT_STATUS)
+            .expect("owned root should compile");
+        let wrapper = compile_system_owned_tail_wrapper(root.entry())
+            .expect("owned system wrapper should compile");
+        let entry = unsafe {
+            std::mem::transmute::<*const u8, unsafe extern "C" fn(*mut crate::Vm, *mut Value) -> i32>(
+                wrapper.entry(),
+            )
+        };
+        let backing = Arc::new("owned".to_string());
+        let mut owned = Value::String(backing.clone());
+        assert_eq!(Arc::strong_count(&backing), 2);
+
+        assert_eq!(
+            unsafe { entry(std::ptr::null_mut(), &mut owned) },
+            DEOPT_STATUS
+        );
+        assert!(matches!(owned, Value::String(_)));
+        assert_eq!(Arc::strong_count(&backing), 2);
+
+        slot.publish(child.entry());
+        assert_eq!(
+            unsafe { entry(std::ptr::null_mut(), &mut owned) },
+            CHILD_STATUS
+        );
+        assert!(matches!(owned, Value::Null));
+        assert_eq!(Arc::strong_count(&backing), 1);
+
+        assert_eq!(
+            unsafe { entry(std::ptr::null_mut(), &mut owned) },
+            CHILD_STATUS
+        );
+        assert!(matches!(owned, Value::Null));
+        assert_eq!(Arc::strong_count(&backing), 1);
+    }
+
     #[test]
     fn selected_backend_is_native() {
         #[cfg(feature = "cranelift-jit")]

src/vm/native/exec.rs

@@ -25,7 +25,10 @@ impl ExecutableBuffer {
                 std::ptr::copy_nonoverlapping(code.as_ptr(), ptr, code.len());
             });
             flush_instruction_cache(ptr, code.len());
-            prepare_for_execution();
+            if let Err(err) = protect_executable_buffer(ptr, code.len()) {
+                free_executable(ptr, code.len());
+                return Err(err);
+            }
         }
 
         Ok(Self {
@@ -67,15 +70,15 @@ impl Drop for ExecutableBuffer {
 #[cfg(windows)]
 unsafe fn alloc_executable(len: usize) -> VmResult<*mut u8> {
     use windows_sys::Win32::System::Memory::{
-        MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE, VirtualAlloc,
+        MEM_COMMIT, MEM_RESERVE, PAGE_READWRITE, VirtualAlloc,
     };
 
     let ptr = unsafe {
         VirtualAlloc(
             std::ptr::null_mut(),
             len,
             MEM_COMMIT | MEM_RESERVE,
-            PAGE_EXECUTE_READWRITE,
+            PAGE_READWRITE,
         )
     } as *mut u8;
     if ptr.is_null() {
@@ -104,11 +107,15 @@ unsafe fn flush_instruction_cache(ptr: *mut u8, len: usize) {
 
 #[cfg(unix)]
 unsafe fn alloc_executable(len: usize) -> VmResult<*mut u8> {
+    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
+    let prot = libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC;
+    #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
+    let prot = libc::PROT_READ | libc::PROT_WRITE;
     let ptr = unsafe {
         libc::mmap(
             std::ptr::null_mut(),
             len,
-            libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
+            prot,
             libc::MAP_PRIVATE | map_anon_flag() | map_jit_flag(),
             -1,
             0,
@@ -173,6 +180,38 @@ const fn map_jit_flag() -> i32 {
     0
 }
 
+#[cfg(all(unix, target_os = "macos", target_arch = "aarch64"))]
+pub(crate) unsafe fn protect_executable_buffer(_ptr: *mut u8, _len: usize) -> VmResult<()> {
+    if unsafe { libc::pthread_jit_write_protect_supported_np() } != 0 {
+        unsafe { libc::pthread_jit_write_protect_np(1) };
+    }
+    Ok(())
+}
+
+#[cfg(all(unix, not(all(target_os = "macos", target_arch = "aarch64"))))]
+pub(crate) unsafe fn protect_executable_buffer(ptr: *mut u8, len: usize) -> VmResult<()> {
+    if unsafe { libc::mprotect(ptr.cast(), len, libc::PROT_READ | libc::PROT_EXEC) } != 0 {
+        return Err(VmError::JitNative(format!(
+            "mprotect failed for executable trace buffer: {}",
+            std::io::Error::last_os_error()
+        )));
+    }
+    Ok(())
+}
+
+#[cfg(windows)]
+pub(crate) unsafe fn protect_executable_buffer(ptr: *mut u8, len: usize) -> VmResult<()> {
+    use windows_sys::Win32::System::Memory::{PAGE_EXECUTE_READ, VirtualProtect};
+
+    let mut previous = 0;
+    if unsafe { VirtualProtect(ptr.cast(), len, PAGE_EXECUTE_READ, &mut previous) } == 0 {
+        return Err(VmError::JitNative(
+            "VirtualProtect failed for executable trace buffer".to_string(),
+        ));
+    }
+    Ok(())
+}
+
 #[cfg(all(unix, target_os = "macos", target_arch = "aarch64"))]
 pub(crate) unsafe fn prepare_for_execution() {
     if unsafe { libc::pthread_jit_write_protect_supported_np() } != 0 {