fix(jit): preserve owned values across native writes

Author: fffonion

src/vm/jit/native/lower.rs

@@ -27,16 +27,16 @@ use crate::vm::native::{
     non_yielding_i64_host_call_signature, non_yielding_scalar_host_call_entry_address,
     non_yielding_scalar_host_call_signature, pack_shared_signature, regex_match_entry_address,
     regex_match_signature, regex_replace_entry_address, regex_replace_signature,
-    restore_active_sparse_exit_state_entry_address, restore_virtual_frame_entry_address,
-    restore_virtual_frame_signature, shared_array_from_buffer_entry_address,
-    shared_bytes_from_buffer_entry_address, shared_string_from_buffer_entry_address,
-    sparse_restore_exit_signature, string_binary_transform_signature,
-    string_contains_entry_address, string_contains_signature, string_lower_ascii_entry_address,
-    string_replace_literal_entry_address, string_replace_signature,
-    string_split_literal_entry_address, string_unary_transform_signature, to_string_entry_address,
-    type_of_entry_address, value_eq_entry_address, value_eq_signature, value_len_entry_address,
-    value_len_signature, value_slot_signature, write_heap_value_to_slot_entry_address,
-    zero_bytes_entry_address,
+    replace_value_in_slot_entry_address, restore_active_sparse_exit_state_entry_address,
+    restore_virtual_frame_entry_address, restore_virtual_frame_signature,
+    shared_array_from_buffer_entry_address, shared_bytes_from_buffer_entry_address,
+    shared_string_from_buffer_entry_address, sparse_restore_exit_signature,
+    string_binary_transform_signature, string_contains_entry_address, string_contains_signature,
+    string_lower_ascii_entry_address, string_replace_literal_entry_address,
+    string_replace_signature, string_split_literal_entry_address, string_unary_transform_signature,
+    to_string_entry_address, type_of_entry_address, value_eq_entry_address, value_eq_signature,
+    value_len_entry_address, value_len_signature, value_slot_signature,
+    write_heap_value_to_slot_entry_address, zero_bytes_entry_address,
 };
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::immediates::Ieee64;
@@ -667,7 +667,8 @@ fn try_compile_ssa_trace(
         };
         let frame_state_ref = b.import_signature(frame_state_sig);
         let deopt_refs = SsaDeoptHelperRefs {
-            clone_value_ref: b.import_signature(clone_value_sig),
+            clone_value_ref: b.import_signature(clone_value_sig.clone()),
+            replace_value_ref: b.import_signature(clone_value_sig),
             value_eq_ref: b.import_signature(value_eq_sig),
             value_len_ref: b.import_signature(value_len_sig),
             non_yielding_host_call_ref: b.import_signature(non_yielding_host_call_sig),
@@ -693,6 +694,7 @@ fn try_compile_ssa_trace(
         };
         let deopt_addrs = SsaDeoptHelperAddrs {
             clone_value: clone_value_to_slot_entry_address(),
+            replace_value: replace_value_in_slot_entry_address(),
             value_eq: value_eq_entry_address(),
             value_len: value_len_entry_address(),
             non_yielding_host_call: non_yielding_host_call_entry_address(),
@@ -1104,6 +1106,7 @@ enum SsaExitAction {
 #[derive(Clone, Copy)]
 struct SsaDeoptHelperRefs {
     clone_value_ref: cranelift_codegen::ir::SigRef,
+    replace_value_ref: cranelift_codegen::ir::SigRef,
     value_eq_ref: cranelift_codegen::ir::SigRef,
     value_len_ref: cranelift_codegen::ir::SigRef,
     non_yielding_host_call_ref: cranelift_codegen::ir::SigRef,
@@ -1130,6 +1133,7 @@ struct SsaDeoptHelperRefs {
 #[derive(Clone, Copy)]
 struct SsaDeoptHelperAddrs {
     clone_value: usize,
+    replace_value: usize,
     value_eq: usize,
     value_len: usize,
     non_yielding_host_call: usize,
@@ -2080,8 +2084,8 @@ fn lower_ssa_inst(
                 b,
                 exit_block,
                 pointer_type,
-                helper_refs.clone_value_ref,
-                helper_addrs.clone_value,
+                helper_refs.replace_value_ref,
+                helper_addrs.replace_value,
                 &[out, values[input]],
             )?;
             out
@@ -4701,6 +4705,7 @@ fn lower_ssa_exit_block(
             active_local_base,
             layout.value.size,
         );
+        let mut cloned_local_addrs = HashMap::new();
         for (local_index, materialization) in
             exit.locals
                 .iter()
@@ -4709,13 +4714,6 @@ fn lower_ssa_exit_block(
                     exit.dirty_locals[local_index].then_some((local_index, materialization))
                 })
         {
-            let index = b.ins().iconst(
-                pointer_type,
-                i64::try_from(local_index).map_err(|_| {
-                    VmError::JitNative("SSA dirty local index out of range".to_string())
-                })?,
-            );
-            let dst_addr = ssa_value_addr(b, pointer_type, vm_locals_ptr, index, layout.value.size);
             let clone_before_clear = match materialization {
                 SsaMaterialization::BoxHeapPtr { .. } => true,
                 SsaMaterialization::Value(value) => !moved_owned_values.contains(value),
@@ -4725,6 +4723,25 @@ fn lower_ssa_exit_block(
                 let temp_slot = ssa_create_value_stack_slot(b, layout.value.size)?;
                 let temp_addr = b.ins().stack_addr(pointer_type, temp_slot, 0);
                 ssa_materialize_slot(b, materialize_ctx, materialization, temp_addr, "local")?;
+                cloned_local_addrs.insert(local_index, temp_addr);
+            }
+        }
+        for (local_index, materialization) in
+            exit.locals
+                .iter()
+                .enumerate()
+                .filter_map(|(local_index, materialization)| {
+                    exit.dirty_locals[local_index].then_some((local_index, materialization))
+                })
+        {
+            let index = b.ins().iconst(
+                pointer_type,
+                i64::try_from(local_index).map_err(|_| {
+                    VmError::JitNative("SSA dirty local index out of range".to_string())
+                })?,
+            );
+            let dst_addr = ssa_value_addr(b, pointer_type, vm_locals_ptr, index, layout.value.size);
+            if let Some(temp_addr) = cloned_local_addrs.get(&local_index).copied() {
                 clear_owned_value_temp_slot(b, pointer_type, deopt_refs, deopt_addrs, dst_addr)?;
                 ssa_copy_value_bytes(b, temp_addr, dst_addr, layout.value.size);
                 ssa_store_tag(b, layout.value, temp_addr, layout.value.null_tag);

src/vm/native/bridge.rs

@@ -264,6 +264,10 @@ pub(crate) fn clone_value_to_slot_entry_address() -> usize {
     pd_vm_native_clone_value_to_slot as *const () as usize
 }
 
+pub(crate) fn replace_value_in_slot_entry_address() -> usize {
+    pd_vm_native_replace_value_in_slot as *const () as usize
+}
+
 pub(crate) fn init_null_value_slot_entry_address() -> usize {
     pd_vm_native_init_null_value_slot as *const () as usize
 }
@@ -636,6 +640,25 @@ pub(crate) extern "C" fn pd_vm_native_clone_value_to_slot(
     STATUS_CONTINUE
 }
 
+pub(crate) extern "C" fn pd_vm_native_replace_value_in_slot(
+    dst: *mut Value,
+    src: *const Value,
+) -> i32 {
+    if dst.is_null() || src.is_null() {
+        store_bridge_error(VmError::JitNative(
+            "native replace-value helper received null slot pointer".to_string(),
+        ));
+        return STATUS_ERROR;
+    }
+
+    unsafe {
+        let replacement = (*src).clone();
+        let previous = std::ptr::replace(dst, replacement);
+        drop(previous);
+    }
+    STATUS_CONTINUE
+}
+
 pub(crate) extern "C" fn pd_vm_native_init_null_value_slot(dst: *mut Value) -> i32 {
     if dst.is_null() {
         store_bridge_error(VmError::JitNative(
@@ -2421,4 +2444,17 @@ mod tests {
         assert_eq!(result.get(&Value::Int(2)), Some(&Value::Int(20)));
         assert!(!Arc::ptr_eq(&result, &alias));
     }
+
+    #[test]
+    fn replace_value_slot_clones_before_dropping_an_aliasing_destination() {
+        let shared = Arc::new(vec![Value::Int(1)]);
+        let mut slot = Value::Array(shared.clone());
+        let slot_ptr = &mut slot as *mut Value;
+
+        let status = pd_vm_native_replace_value_in_slot(slot_ptr, slot_ptr);
+
+        assert_eq!(status, STATUS_CONTINUE);
+        assert_eq!(Arc::strong_count(&shared), 2);
+        assert!(matches!(slot, Value::Array(_)));
+    }
 }

src/vm/native/mod.rs

@@ -22,15 +22,15 @@ pub(crate) use bridge::{
     map_iter_take_key_entry_address, map_iter_take_value_entry_address, map_set_entry_address,
     non_yielding_host_call_entry_address, non_yielding_i64_host_call_entry_address,
     non_yielding_scalar_host_call_entry_address, regex_match_entry_address,
-    regex_replace_entry_address, restore_active_exit_state_entry_address,
-    restore_active_sparse_exit_state_entry_address, restore_exit_state_entry_address,
-    restore_sparse_exit_state_entry_address, restore_virtual_frame_entry_address,
-    shared_array_from_buffer_entry_address, shared_bytes_from_buffer_entry_address,
-    shared_string_from_buffer_entry_address, store_bridge_error, string_contains_entry_address,
-    string_lower_ascii_entry_address, string_replace_literal_entry_address,
-    string_split_literal_entry_address, take_bridge_error, to_string_entry_address,
-    type_of_entry_address, value_eq_entry_address, value_len_entry_address,
-    write_heap_value_to_slot_entry_address, zero_bytes_entry_address,
+    regex_replace_entry_address, replace_value_in_slot_entry_address,
+    restore_active_exit_state_entry_address, restore_active_sparse_exit_state_entry_address,
+    restore_exit_state_entry_address, restore_sparse_exit_state_entry_address,
+    restore_virtual_frame_entry_address, shared_array_from_buffer_entry_address,
+    shared_bytes_from_buffer_entry_address, shared_string_from_buffer_entry_address,
+    store_bridge_error, string_contains_entry_address, string_lower_ascii_entry_address,
+    string_replace_literal_entry_address, string_split_literal_entry_address, take_bridge_error,
+    to_string_entry_address, type_of_entry_address, value_eq_entry_address,
+    value_len_entry_address, write_heap_value_to_slot_entry_address, zero_bytes_entry_address,
 };
 #[cfg(feature = "cranelift-jit")]
 pub(crate) use codegen::{

tests/jit/jit_tests.rs

@@ -3342,6 +3342,63 @@ fn trace_jit_restores_tagged_heap_locals_on_ssa_exit() {
     );
 }
 
+#[test]
+fn trace_jit_snapshots_borrowed_tagged_locals_before_exit_writes() {
+    if !native_jit_supported() {
+        return;
+    }
+
+    let mut bc = BytecodeBuilder::new();
+    bc.ldc(0);
+    bc.stloc(0);
+    bc.ldc(1);
+    bc.stloc(1);
+    bc.ldc(2);
+    bc.stloc(2);
+    let root_ip = bc.position();
+    bc.ldloc(2);
+    bc.ldc(4);
+    bc.clt();
+    let guard_ip = bc.position();
+    bc.brfalse(0);
+    bc.ldloc(0);
+    bc.ldloc(1);
+    bc.stloc(0);
+    bc.stloc(1);
+    bc.ldloc(2);
+    bc.ldc(3);
+    bc.add();
+    bc.stloc(2);
+    bc.br(root_ip);
+    let exit_ip = bc.position();
+    bc.ldloc(1);
+    bc.ret();
+
+    let mut code = bc.finish();
+    patch_branch_target(&mut code, guard_ip, exit_ip);
+    let program = Program::new(
+        vec![
+            Value::string("left"),
+            Value::string("right"),
+            Value::Int(0),
+            Value::Int(1),
+            Value::Int(5),
+        ],
+        code,
+    )
+    .with_local_count(3);
+    let mut vm = Vm::new(program);
+    vm.set_jit_config(JitConfig {
+        enabled: true,
+        hot_loop_threshold: 1,
+        max_trace_len: 512,
+    });
+
+    assert_eq!(vm.run().expect("vm should run"), VmStatus::Halted);
+    assert_eq!(vm.stack(), &[Value::string("left")]);
+    assert!(vm.jit_native_exec_count() > 0, "{}", vm.dump_jit_info());
+}
+
 #[test]
 fn trace_jit_restores_array_and_map_locals_on_ssa_exit() {
     if !native_jit_supported() {