From 917e3040b7c2105eda00f4da39ae86f4e375410a Mon Sep 17 00:00:00 2001 From: fffonion Date: Wed, 15 Jul 2026 23:20:04 +0800 Subject: [PATCH] feat(jit): specialize Set/ArrayPush on delayed-move collection rebinds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Record same-local rebinds (stloc Null then Set/ArrayPush) directly in native traces, removing the call boundary that previously forced interpreter fallback. Key changes: - ValueInfo now carries a source_local provenance tag set exclusively at Ldloc; arithmetic, constants, and call outputs clear it, ensuring the recorder and static loop-header plan agree on which local a value came from. - Recorder selects ArraySet / MapSet / ArrayPush specialization only when the container operand still carries the source local and that symbolic local is Null (dirty + known_type Null). Additional args that also share the moved local are rejected to prevent alias correctness issues. - New native bridge helpers pd_vm_native_collection_set and pd_vm_native_array_push use ptr::replace to take true ownership of the container slot, perform COW/rebind/shared-impl logic, and write the result back to the same owned-temp slot, compatible with native loop-back. - Reusable owned helpers builtin_set_owned / builtin_array_push_owned keep interpreter and JIT paths on the same semantics. - Lowering allocates MutationArgBox temp slots for index/key/value and calls collection_set / array_push native helpers with status checking. - Benchmark harness extended to gate per-sample native execution growth and freeze trace_attempts / recorded_traces / native traces during measurement. Tests (all passing): - trace_jit_specializes_same_local_array_set_through_loop_back - trace_jit_specializes_same_local_map_set_through_loop_back (dual mutation with null-delete) - trace_jit_specializes_same_local_array_push_through_loop_back - trace_jit_array_set_preserves_cow_alias - trace_jit_does_not_consume_non_moved_array_set_container Collection rebind benchmark (release, 50K iters, width=256, 15 samples): array: interpreter 5928 µs → JIT 1181 µs (0.20×, ~5× speedup) map: interpreter 22822 µs → JIT 2325 µs (0.10×, ~10× speedup) WAF benchmark sub-par because its Set calls do not use the rebind pattern; JIT trace statistics confirm the expected LoopBack shape for the micro-bench workloads. --- examples/collection_rebind_bench.rs | 33 ++- src/builtins/runtime/core.rs | 34 +-- src/vm/jit/ir.rs | 30 +++ src/vm/jit/native/lower.rs | 226 ++++++++++++++++- src/vm/jit/recorder.rs | 172 ++++++++++++- src/vm/native/bridge.rs | 64 +++++ src/vm/native/codegen.rs | 14 ++ src/vm/native/mod.rs | 26 +- tests/jit/jit_tests.rs | 374 ++++++++++++++++++++++++++++ 9 files changed, 930 insertions(+), 43 deletions(-) diff --git a/examples/collection_rebind_bench.rs b/examples/collection_rebind_bench.rs index d2e52755..f954e6c2 100644 --- a/examples/collection_rebind_bench.rs +++ b/examples/collection_rebind_bench.rs @@ -40,6 +40,8 @@ impl ExecMode { struct Measurement { samples: Vec, + trace_attempts: usize, + recorded_traces: usize, native_traces: usize, call_boundary_traces: usize, loop_back_traces: usize, @@ -128,12 +130,14 @@ fn run() -> Result<(), String> { .collect::>() .join(","); println!( - "workload={} mode={} width={} iterations={} warmup_runs=1 reused_vm=true median_ns={} native_traces={} call_boundary_traces={} loop_back_traces={} native_execs={} trace_exits={} native_loop_backs={} helper_fallbacks={} generic_builtin_calls={} samples_ns={}", + "workload={} mode={} width={} iterations={} warmup_runs=1 reused_vm=true median_ns={} trace_attempts={} recorded_traces={} native_traces={} call_boundary_traces={} loop_back_traces={} native_execs={} trace_exits={} native_loop_backs={} helper_fallbacks={} generic_builtin_calls={} samples_ns={}", workload.label(), mode.label(), config.width, config.iterations, median.as_nanos(), + measurement.trace_attempts, + measurement.recorded_traces, measurement.native_traces, measurement.call_boundary_traces, measurement.loop_back_traces, @@ -202,7 +206,9 @@ fn measure( let snapshot_before = vm.jit_snapshot(); let metrics_before = snapshot_before.metrics; - let native_traces = snapshot_before.traces.len(); + let trace_attempts = snapshot_before.attempts.len(); + let recorded_traces = snapshot_before.traces.len(); + let native_traces = vm.jit_native_trace_count(); let call_boundary_traces = snapshot_before .traces .iter() @@ -217,30 +223,45 @@ fn measure( let mut generic_builtin_calls = 0u64; for _ in 0..config.samples { vm.reset_for_reuse(); + let native_execs_before_sample = vm.jit_native_exec_count(); let started = Instant::now(); let status = vm .run() .map_err(|err| format!("{} {} run failed: {err}", workload.label(), mode.label()))?; let elapsed = started.elapsed(); verify_result(&vm, status, config)?; + if matches!(mode, ExecMode::Jit) && vm.jit_native_exec_count() <= native_execs_before_sample + { + return Err(format!( + "{} measured run did not execute a warmed native trace:\n{}", + workload.label(), + vm.dump_jit_info() + )); + } generic_builtin_calls = generic_builtin_calls .saturating_add(vm.interpreter_metrics_snapshot().generic_builtin_call_count); black_box(vm.stack()); samples.push(elapsed); } let snapshot_after = vm.jit_snapshot(); - if snapshot_after.traces.len() != native_traces { + let measured_trace_attempts = snapshot_after.attempts.len(); + let measured_recorded_traces = snapshot_after.traces.len(); + let measured_native_traces = vm.jit_native_trace_count(); + if measured_trace_attempts != trace_attempts + || measured_recorded_traces != recorded_traces + || measured_native_traces != native_traces + { return Err(format!( - "{} {} compiled additional traces during measured runs: warmup={} measured={}", + "{} {} changed JIT trace state during measured runs: warmup=attempts:{trace_attempts}/recorded:{recorded_traces}/native:{native_traces} measured=attempts:{measured_trace_attempts}/recorded:{measured_recorded_traces}/native:{measured_native_traces}", workload.label(), mode.label(), - native_traces, - snapshot_after.traces.len() )); } let metrics_after = snapshot_after.metrics; Ok(Measurement { samples, + trace_attempts, + recorded_traces, native_traces, call_boundary_traces, loop_back_traces, diff --git a/src/builtins/runtime/core.rs b/src/builtins/runtime/core.rs index 9372f00f..6993efc1 100644 --- a/src/builtins/runtime/core.rs +++ b/src/builtins/runtime/core.rs @@ -220,15 +220,18 @@ fn builtin_array_push_shared_impl(mut items: SharedArray, value: AnyValue) -> Sh items } -pub(super) fn builtin_array_push(args: &mut [Value]) -> VmResult { - let items = match take_arg(args, 0, "array_push array")? { +pub(crate) fn builtin_array_push_owned(items: Value, value: Value) -> VmResult { + let items = match items { Value::Array(values) => values, _ => return Err(VmError::TypeMismatch("array")), }; + Ok(Value::Array(builtin_array_push_shared_impl(items, value))) +} + +pub(super) fn builtin_array_push(args: &mut [Value]) -> VmResult { + let items = take_arg(args, 0, "array_push array")?; let value = take_arg(args, 1, "array_push value")?; - Ok(return_one(Value::Array(builtin_array_push_shared_impl( - items, value, - )))) + builtin_array_push_owned(items, value).map(return_one) } /// Create an empty map. @@ -684,20 +687,23 @@ fn builtin_set_map_shared_impl( entries } -pub(super) fn builtin_set(args: &mut [Value]) -> VmResult { - let container: Value = take_arg(args, 0, "set container")?; - let key: Value = take_arg(args, 1, "set key")?; - let value: Value = take_arg(args, 2, "set value")?; +pub(crate) fn builtin_set_owned(container: Value, key: Value, value: Value) -> VmResult { match container { - Value::Array(values) => builtin_set_array_shared_impl(values, key.as_int()?, value) - .map(|values| return_one(Value::Array(values))), - Value::Map(entries) => Ok(return_one(Value::Map(builtin_set_map_shared_impl( - entries, key, value, - )))), + Value::Array(values) => { + builtin_set_array_shared_impl(values, key.as_int()?, value).map(Value::Array) + } + Value::Map(entries) => Ok(Value::Map(builtin_set_map_shared_impl(entries, key, value))), _ => Err(VmError::TypeMismatch("array/map")), } } +pub(super) fn builtin_set(args: &mut [Value]) -> VmResult { + let container = take_arg(args, 0, "set container")?; + let key = take_arg(args, 1, "set key")?; + let value = take_arg(args, 2, "set value")?; + builtin_set_owned(container, key, value).map(return_one) +} + /// Return an array of container keys or indices. #[pd_host_function(name = "keys")] pub(super) fn builtin_keys_array_impl(items: VmArrayRef<'_>) -> VmArray { diff --git a/src/vm/jit/ir.rs b/src/vm/jit/ir.rs index 4f0bea06..3689c92e 100644 --- a/src/vm/jit/ir.rs +++ b/src/vm/jit/ir.rs @@ -176,6 +176,15 @@ pub(crate) enum SsaInstKind { array: SsaValueId, index: SsaValueId, }, + ArraySet { + array: SsaValueId, + index: SsaValueId, + value: SsaValueId, + }, + ArrayPush { + array: SsaValueId, + value: SsaValueId, + }, MapLen { map: SsaValueId, }, @@ -187,6 +196,11 @@ pub(crate) enum SsaInstKind { map: SsaValueId, key: SsaValueId, }, + MapSet { + map: SsaValueId, + key: SsaValueId, + value: SsaValueId, + }, IntNeg { input: SsaValueId, }, @@ -363,8 +377,15 @@ impl SsaInstKind { Self::BytesToArrayU8 { bytes } => vec![*bytes], Self::ArrayGet { array, index } => vec![*array, *index], Self::ArrayHas { array, index } => vec![*array, *index], + Self::ArraySet { + array, + index, + value, + } => vec![*array, *index, *value], + Self::ArrayPush { array, value } => vec![*array, *value], Self::MapGet { map, key } => vec![*map, *key], Self::MapHas { map, key } => vec![*map, *key], + Self::MapSet { map, key, value } => vec![*map, *key, *value], Self::IntAdd { lhs, rhs } | Self::IntSub { lhs, rhs } | Self::IntMul { lhs, rhs } @@ -963,9 +984,18 @@ fn render_inst_kind(kind: &SsaInstKind) -> String { SsaInstKind::ArrayLen { array } => format!("array_len {array}"), SsaInstKind::ArrayGet { array, index } => format!("array_get {array}, {index}"), SsaInstKind::ArrayHas { array, index } => format!("array_has {array}, {index}"), + SsaInstKind::ArraySet { + array, + index, + value, + } => format!("array_set {array}, {index}, {value}"), + SsaInstKind::ArrayPush { array, value } => format!("array_push {array}, {value}"), SsaInstKind::MapLen { map } => format!("map_len {map}"), SsaInstKind::MapGet { map, key } => format!("map_get {map}, {key}"), SsaInstKind::MapHas { map, key } => format!("map_has {map}, {key}"), + SsaInstKind::MapSet { map, key, value } => { + format!("map_set {map}, {key}, {value}") + } SsaInstKind::IntNeg { input } => format!("ineg {input}"), SsaInstKind::IntAdd { lhs, rhs } => format!("iadd {lhs}, {rhs}"), SsaInstKind::IntAddImm { lhs, imm } => format!("iadd_imm {lhs}, {imm}"), diff --git a/src/vm/jit/native/lower.rs b/src/vm/jit/native/lower.rs index 48edcbbe..81b829d4 100644 --- a/src/vm/jit/native/lower.rs +++ b/src/vm/jit/native/lower.rs @@ -11,9 +11,10 @@ use crate::vm::native::{ ExecutableBuffer, HeapIntrinsicAddrs, HeapIntrinsicRefs, NativeInterruptMode, NativeInterruptSettings, NativeStackLayout, STATUS_CONTINUE, STATUS_ERROR, STATUS_HALTED, STATUS_OUT_OF_FUEL, STATUS_TRACE_EXIT, alloc_buffer_signature, alloc_byte_buffer_entry_address, - alloc_value_buffer_entry_address, box_heap_value_signature, checked_add_i32, - clear_value_slot_entry_address, clone_value_signature, clone_value_to_slot_entry_address, - collection_get_signature, collection_predicate_signature, copy_bytes_entry_address, + alloc_value_buffer_entry_address, array_push_entry_address, box_heap_value_signature, + checked_add_i32, clear_value_slot_entry_address, clone_value_signature, + clone_value_to_slot_entry_address, collection_get_signature, collection_mutation_signature, + collection_predicate_signature, collection_set_entry_address, copy_bytes_entry_address, copy_bytes_signature, detect_native_stack_layout, entry_signature, free_buffer_signature, init_null_value_slot_entry_address, jump_with_status, map_get_entry_address, map_has_entry_address, pack_shared_signature, restore_sparse_exit_state_entry_address, @@ -102,6 +103,8 @@ fn try_compile_ssa_trace( let copy_bytes_sig = copy_bytes_signature(pointer_type, call_conv); let map_has_sig = collection_predicate_signature(pointer_type, call_conv); let map_get_sig = collection_get_signature(pointer_type, call_conv); + let array_push_sig = collection_get_signature(pointer_type, call_conv); + let collection_set_sig = collection_mutation_signature(pointer_type, call_conv); let sparse_restore_exit_sig = sparse_restore_exit_signature(pointer_type, call_conv); let resume_linked_trace_sig = entry_signature(pointer_type, call_conv); let string_contains_sig = string_contains_signature(pointer_type, call_conv); @@ -157,6 +160,8 @@ fn try_compile_ssa_trace( box_heap_value_ref: b.import_signature(box_heap_value_sig), map_has_ref: b.import_signature(map_has_sig), map_get_ref: b.import_signature(map_get_sig), + array_push_ref: b.import_signature(array_push_sig), + collection_set_ref: b.import_signature(collection_set_sig), sparse_restore_exit_ref: b.import_signature(sparse_restore_exit_sig), resume_linked_trace_ref: b.import_signature(resume_linked_trace_sig), }; @@ -167,6 +172,8 @@ fn try_compile_ssa_trace( box_heap_value: write_heap_value_to_slot_entry_address(), map_has: map_has_entry_address(), map_get: map_get_entry_address(), + array_push: array_push_entry_address(), + collection_set: collection_set_entry_address(), sparse_restore_exit: restore_sparse_exit_state_entry_address(), resume_linked_trace: resume_linked_trace_entry_address(), }; @@ -408,6 +415,8 @@ struct SsaDeoptHelperRefs { box_heap_value_ref: cranelift_codegen::ir::SigRef, map_has_ref: cranelift_codegen::ir::SigRef, map_get_ref: cranelift_codegen::ir::SigRef, + array_push_ref: cranelift_codegen::ir::SigRef, + collection_set_ref: cranelift_codegen::ir::SigRef, sparse_restore_exit_ref: cranelift_codegen::ir::SigRef, resume_linked_trace_ref: cranelift_codegen::ir::SigRef, } @@ -420,6 +429,8 @@ struct SsaDeoptHelperAddrs { box_heap_value: usize, map_has: usize, map_get: usize, + array_push: usize, + collection_set: usize, sparse_restore_exit: usize, resume_linked_trace: usize, } @@ -444,6 +455,7 @@ struct SsaStringHelperAddrs { enum SsaTempValueSlotKey { Output(SsaValueId), MapKeyBox(SsaValueId), + MutationArgBox(SsaValueId, u8), } #[derive(Clone)] @@ -527,9 +539,12 @@ fn ssa_trace_supported(ssa: &SsaTrace) -> bool { | SsaInstKind::ArrayLen { .. } | SsaInstKind::ArrayGet { .. } | SsaInstKind::ArrayHas { .. } + | SsaInstKind::ArraySet { .. } + | SsaInstKind::ArrayPush { .. } | SsaInstKind::MapLen { .. } | SsaInstKind::MapGet { .. } | SsaInstKind::MapHas { .. } + | SsaInstKind::MapSet { .. } | SsaInstKind::IntNeg { .. } | SsaInstKind::IntAdd { .. } | SsaInstKind::IntAddImm { .. } @@ -655,6 +670,18 @@ fn allocate_owned_value_temps( ordered.push(slot); slots.insert(SsaTempValueSlotKey::MapKeyBox(output.id), slot); } + if matches!( + inst.kind, + SsaInstKind::ArraySet { .. } + | SsaInstKind::ArrayPush { .. } + | SsaInstKind::MapSet { .. } + ) { + for arg in 0..2 { + let slot = ssa_create_value_stack_slot(b, value_size)?; + ordered.push(slot); + slots.insert(SsaTempValueSlotKey::MutationArgBox(output.id, arg), slot); + } + } } } Ok(SsaOwnedValueTemps { ordered, slots }) @@ -664,7 +691,10 @@ fn ssa_inst_requires_owned_value_slot(kind: &SsaInstKind) -> bool { matches!( kind, SsaInstKind::ArrayGet { .. } + | SsaInstKind::ArraySet { .. } + | SsaInstKind::ArrayPush { .. } | SsaInstKind::MapGet { .. } + | SsaInstKind::MapSet { .. } | SsaInstKind::StringSlice { .. } | SsaInstKind::BytesSlice { .. } | SsaInstKind::StringGet { .. } @@ -1905,6 +1935,127 @@ fn lower_ssa_inst( ); ssa_index_in_range(b, index, len) } + SsaInstKind::ArraySet { + array, + index, + value, + } => { + let out = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::Output(output.id), + )?; + let index_box = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::MutationArgBox(output.id, 0), + )?; + let index_addr = ssa_ensure_boxed_value_addr( + b, + SsaBoxCtx { + exit_block, + pointer_type, + value_layout: layout.value, + helper_refs, + helper_addrs, + }, + Some(index_box), + *value_reprs.get(index).ok_or_else(|| { + VmError::JitNative("SSA array-set index representation missing".to_string()) + })?, + values[index], + )?; + let value_box = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::MutationArgBox(output.id, 1), + )?; + let value_addr = ssa_ensure_boxed_value_addr( + b, + SsaBoxCtx { + exit_block, + pointer_type, + value_layout: layout.value, + helper_refs, + helper_addrs, + }, + Some(value_box), + *value_reprs.get(value).ok_or_else(|| { + VmError::JitNative("SSA array-set value representation missing".to_string()) + })?, + values[value], + )?; + let helper_ptr = iconst_ptr_from_addr(b, pointer_type, helper_addrs.collection_set)?; + let call = b.ins().call_indirect( + helper_refs.collection_set_ref, + helper_ptr, + &[out, values[array], index_addr, value_addr], + ); + let status = b.inst_results(call)[0]; + let fail = b.create_block(); + let cont = b.create_block(); + let is_error = b + .ins() + .icmp_imm(IntCC::Equal, status, i64::from(STATUS_ERROR)); + b.ins().brif(is_error, fail, &[], cont, &[]); + + b.switch_to_block(fail); + jump_with_status(b, exit_block, status); + + b.switch_to_block(cont); + out + } + SsaInstKind::ArrayPush { array, value } => { + let out = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::Output(output.id), + )?; + let value_box = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::MutationArgBox(output.id, 0), + )?; + let value_addr = ssa_ensure_boxed_value_addr( + b, + SsaBoxCtx { + exit_block, + pointer_type, + value_layout: layout.value, + helper_refs, + helper_addrs, + }, + Some(value_box), + *value_reprs.get(value).ok_or_else(|| { + VmError::JitNative("SSA array-push value representation missing".to_string()) + })?, + values[value], + )?; + let helper_ptr = iconst_ptr_from_addr(b, pointer_type, helper_addrs.array_push)?; + let call = b.ins().call_indirect( + helper_refs.array_push_ref, + helper_ptr, + &[out, values[array], value_addr], + ); + let status = b.inst_results(call)[0]; + let fail = b.create_block(); + let cont = b.create_block(); + let is_error = b + .ins() + .icmp_imm(IntCC::Equal, status, i64::from(STATUS_ERROR)); + b.ins().brif(is_error, fail, &[], cont, &[]); + + b.switch_to_block(fail); + jump_with_status(b, exit_block, status); + + b.switch_to_block(cont); + out + } SsaInstKind::MapLen { map } => { let map_ptr = ssa_load_heap_data_ptr(b, layout.value, values[map]); b.ins().load( @@ -2022,6 +2173,75 @@ fn lower_ssa_inst( b.switch_to_block(cont); b.ins().icmp_imm(IntCC::NotEqual, status, 0) } + SsaInstKind::MapSet { map, key, value } => { + let out = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::Output(output.id), + )?; + let key_box = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::MutationArgBox(output.id, 0), + )?; + let key_addr = ssa_ensure_boxed_value_addr( + b, + SsaBoxCtx { + exit_block, + pointer_type, + value_layout: layout.value, + helper_refs, + helper_addrs, + }, + Some(key_box), + *value_reprs.get(key).ok_or_else(|| { + VmError::JitNative("SSA map-set key representation missing".to_string()) + })?, + values[key], + )?; + let value_box = owned_value_temp_slot_addr( + b, + pointer_type, + owned_value_temps, + SsaTempValueSlotKey::MutationArgBox(output.id, 1), + )?; + let value_addr = ssa_ensure_boxed_value_addr( + b, + SsaBoxCtx { + exit_block, + pointer_type, + value_layout: layout.value, + helper_refs, + helper_addrs, + }, + Some(value_box), + *value_reprs.get(value).ok_or_else(|| { + VmError::JitNative("SSA map-set value representation missing".to_string()) + })?, + values[value], + )?; + let helper_ptr = iconst_ptr_from_addr(b, pointer_type, helper_addrs.collection_set)?; + let call = b.ins().call_indirect( + helper_refs.collection_set_ref, + helper_ptr, + &[out, values[map], key_addr, value_addr], + ); + let status = b.inst_results(call)[0]; + let fail = b.create_block(); + let cont = b.create_block(); + let is_error = b + .ins() + .icmp_imm(IntCC::Equal, status, i64::from(STATUS_ERROR)); + b.ins().brif(is_error, fail, &[], cont, &[]); + + b.switch_to_block(fail); + jump_with_status(b, exit_block, status); + + b.switch_to_block(cont); + out + } SsaInstKind::IntNeg { input } => b.ins().ineg(values[input]), SsaInstKind::IntAdd { lhs, rhs } => b.ins().iadd(values[lhs], values[rhs]), SsaInstKind::IntAddImm { lhs, imm } => b.ins().iadd_imm(values[lhs], *imm), diff --git a/src/vm/jit/recorder.rs b/src/vm/jit/recorder.rs index 89c2b6d0..84e70ea8 100644 --- a/src/vm/jit/recorder.rs +++ b/src/vm/jit/recorder.rs @@ -87,6 +87,7 @@ struct ValueInfo { const_float: Option, const_bool: Option, known_type: Option, + source_local: Option, } impl ValueInfo { @@ -97,6 +98,7 @@ impl ValueInfo { const_float: None, const_bool: None, known_type: None, + source_local: None, } } @@ -107,6 +109,7 @@ impl ValueInfo { const_float: None, const_bool: None, known_type: Some(known_type), + source_local: None, } } @@ -117,6 +120,7 @@ impl ValueInfo { const_float: None, const_bool: None, known_type: Some(ValueType::Int), + source_local: None, } } @@ -127,6 +131,7 @@ impl ValueInfo { const_float: value, const_bool: None, known_type: Some(ValueType::Float), + source_local: None, } } @@ -137,6 +142,7 @@ impl ValueInfo { const_float: None, const_bool: value, known_type: Some(ValueType::Bool), + source_local: None, } } @@ -147,6 +153,7 @@ impl ValueInfo { const_float: None, const_bool: None, known_type: Some(tag), + source_local: None, } } @@ -162,6 +169,11 @@ impl ValueInfo { Value::Map(_) => Self::tagged_typed(ValueType::Map), } } + + fn sourced_from(mut self, local: u8) -> Self { + self.source_local = Some(local); + self + } } #[derive(Clone, Debug, PartialEq)] @@ -462,9 +474,12 @@ enum SpecializedBuiltinKind { ArrayLen, ArrayGet, ArrayHas, + ArraySet, + ArrayPush, MapLen, MapGet, MapHas, + MapSet, } struct TraceCursor<'a> { @@ -702,6 +717,7 @@ pub(crate) fn record_trace( const_float: None, const_bool: None, known_type: loop_plan.stack_known_types[index], + source_local: None, }, }); entry_args.push(entry_arg); @@ -758,6 +774,7 @@ pub(crate) fn record_trace( const_float: None, const_bool: None, known_type: loop_plan.local_known_types[local], + source_local: None, }, }); entry_args.push(entry_arg); @@ -816,7 +833,9 @@ pub(crate) fn record_trace( } DecodedOp::Ldloc { index, .. } => { op_names.push("ldloc".to_string()); - frame.push(frame.local(index)?); + let mut value = frame.local(index)?; + value.info = value.info.sourced_from(index); + frame.push(value); } DecodedOp::Stloc { index, .. } => { op_names.push("stloc".to_string()); @@ -1082,9 +1101,22 @@ pub(crate) fn record_trace( && argc > 0 { let args = call_arg_slice(&frame.stack, usize::from(argc))?; - if let Some(kind) = - select_specialized_builtin_kind(program, ip, builtin, args[0].info) - { + let container_source = args[0].info.source_local; + let container_was_moved = container_source.is_some_and(|local| { + frame.dirty_locals[usize::from(local)] + && frame.locals[usize::from(local)].info.known_type + == Some(ValueType::Null) + && args[1..] + .iter() + .all(|arg| arg.info.source_local != Some(local)) + }); + if let Some(kind) = select_specialized_builtin_kind( + program, + ip, + builtin, + args[0].info, + container_was_moved, + ) { let (name, out) = emit_specialized_builtin_call( &mut builder, current_block, @@ -1169,7 +1201,7 @@ fn infer_loop_header_plan( { *state = EntryUseState::ReadBeforeWrite; } - frame.push(frame.local(index)?) + frame.push(frame.local(index)?.sourced_from(index)) } DecodedOp::Stloc { index, .. } => { if let Some(written) = local_written.get_mut(index as usize) { @@ -1264,8 +1296,18 @@ fn infer_loop_header_plan( return Ok(None); } let args = call_arg_slice(&frame.stack, usize::from(argc))?; - let Some(kind) = select_specialized_builtin_kind(program, ip, builtin, args[0]) - else { + let container_source = args[0].source_local; + let container_was_moved = container_source.is_some_and(|local| { + frame.locals[usize::from(local)].known_type == Some(ValueType::Null) + && args[1..].iter().all(|arg| arg.source_local != Some(local)) + }); + let Some(kind) = select_specialized_builtin_kind( + program, + ip, + builtin, + args[0], + container_was_moved, + ) else { return Ok(None); }; let _ = analyze_specialized_builtin_call(&mut frame, kind)?; @@ -2416,6 +2458,7 @@ fn select_specialized_builtin_kind( ip: usize, builtin: BuiltinFunction, container: ValueInfo, + container_was_moved: bool, ) -> Option { let observed_kind = observed_heap_container_kind(container); let container_kind = if matches!( @@ -2425,6 +2468,8 @@ fn select_specialized_builtin_kind( | BuiltinFunction::Get | BuiltinFunction::Has | BuiltinFunction::Concat + | BuiltinFunction::Set + | BuiltinFunction::ArrayPush ) { match operand_types(program, ip).0 { ValueType::String => Some(HeapContainerKind::String), @@ -2480,9 +2525,18 @@ fn select_specialized_builtin_kind( (BuiltinFunction::Len, HeapContainerKind::Array) => Some(SpecializedBuiltinKind::ArrayLen), (BuiltinFunction::Get, HeapContainerKind::Array) => Some(SpecializedBuiltinKind::ArrayGet), (BuiltinFunction::Has, HeapContainerKind::Array) => Some(SpecializedBuiltinKind::ArrayHas), + (BuiltinFunction::Set, HeapContainerKind::Array) if container_was_moved => { + Some(SpecializedBuiltinKind::ArraySet) + } + (BuiltinFunction::ArrayPush, HeapContainerKind::Array) if container_was_moved => { + Some(SpecializedBuiltinKind::ArrayPush) + } (BuiltinFunction::Len, HeapContainerKind::Map) => Some(SpecializedBuiltinKind::MapLen), (BuiltinFunction::Get, HeapContainerKind::Map) => Some(SpecializedBuiltinKind::MapGet), (BuiltinFunction::Has, HeapContainerKind::Map) => Some(SpecializedBuiltinKind::MapHas), + (BuiltinFunction::Set, HeapContainerKind::Map) if container_was_moved => { + Some(SpecializedBuiltinKind::MapSet) + } _ => None, } } @@ -2597,6 +2651,19 @@ fn analyze_specialized_builtin_call( frame.push(ValueInfo::bool(None)); Ok("array_has") } + SpecializedBuiltinKind::ArraySet => { + let _ = frame.pop()?; + let _ = frame.pop()?; + let _ = frame.pop()?; + frame.push(ValueInfo::tagged_typed(ValueType::Array)); + Ok("array_set") + } + SpecializedBuiltinKind::ArrayPush => { + let _ = frame.pop()?; + let _ = frame.pop()?; + frame.push(ValueInfo::tagged_typed(ValueType::Array)); + Ok("array_push") + } SpecializedBuiltinKind::MapLen => { let _ = frame.pop()?; frame.push(ValueInfo::int(None)); @@ -2614,6 +2681,13 @@ fn analyze_specialized_builtin_call( frame.push(ValueInfo::bool(None)); Ok("map_has") } + SpecializedBuiltinKind::MapSet => { + let _ = frame.pop()?; + let _ = frame.pop()?; + let _ = frame.pop()?; + frame.push(ValueInfo::tagged_typed(ValueType::Map)); + Ok("map_set") + } } } @@ -3076,6 +3150,60 @@ fn emit_specialized_builtin_call( .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?; Ok(("array_has", out)) } + SpecializedBuiltinKind::ArraySet => { + let value = frame.pop()?; + let index = ensure_int(builder, block, ip, frame.pop()?)?; + let array = frame.pop()?; + if array.info.repr != SsaValueRepr::Tagged { + return Err(TraceRecordError::TypeMismatch { + expected: "owned tagged array", + actual: array.info.repr, + }); + } + let out = builder + .append_value_inst( + block, + ip, + SsaValueRepr::Tagged, + SsaInstKind::ArraySet { + array: array.value.id, + index: index.value.id, + value: value.value.id, + }, + ) + .map(|value| SymbolicValue { + value, + info: ValueInfo::tagged_typed(ValueType::Array), + }) + .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?; + Ok(("array_set", out)) + } + SpecializedBuiltinKind::ArrayPush => { + let value = frame.pop()?; + let array = frame.pop()?; + if array.info.repr != SsaValueRepr::Tagged { + return Err(TraceRecordError::TypeMismatch { + expected: "owned tagged array", + actual: array.info.repr, + }); + } + let out = builder + .append_value_inst( + block, + ip, + SsaValueRepr::Tagged, + SsaInstKind::ArrayPush { + array: array.value.id, + value: value.value.id, + }, + ) + .map(|value| SymbolicValue { + value, + info: ValueInfo::tagged_typed(ValueType::Array), + }) + .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?; + Ok(("array_push", out)) + } SpecializedBuiltinKind::MapLen => { let map = frame.pop()?; let map = @@ -3138,6 +3266,34 @@ fn emit_specialized_builtin_call( .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?; Ok(("map_has", out)) } + SpecializedBuiltinKind::MapSet => { + let value = frame.pop()?; + let key = frame.pop()?; + let map = frame.pop()?; + if map.info.repr != SsaValueRepr::Tagged { + return Err(TraceRecordError::TypeMismatch { + expected: "owned tagged map", + actual: map.info.repr, + }); + } + let out = builder + .append_value_inst( + block, + ip, + SsaValueRepr::Tagged, + SsaInstKind::MapSet { + map: map.value.id, + key: key.value.id, + value: value.value.id, + }, + ) + .map(|value| SymbolicValue { + value, + info: ValueInfo::tagged_typed(ValueType::Map), + }) + .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?; + Ok(("map_set", out)) + } } } @@ -3666,6 +3822,7 @@ fn continue_with_frame( const_float: None, const_bool: None, known_type: value.info.known_type, + source_local: None, }, }); } @@ -3687,6 +3844,7 @@ fn continue_with_frame( const_float: None, const_bool: None, known_type: value.info.known_type, + source_local: None, }, }); } diff --git a/src/vm/native/bridge.rs b/src/vm/native/bridge.rs index e51ee7e7..fa7d83a7 100644 --- a/src/vm/native/bridge.rs +++ b/src/vm/native/bridge.rs @@ -257,6 +257,14 @@ pub(crate) fn map_get_entry_address() -> usize { pd_vm_native_map_get as *const () as usize } +pub(crate) fn collection_set_entry_address() -> usize { + pd_vm_native_collection_set as *const () as usize +} + +pub(crate) fn array_push_entry_address() -> usize { + pd_vm_native_array_push as *const () as usize +} + pub(crate) fn helper_entry_offset() -> i32 { i32::try_from(std::mem::offset_of!(Vm, native_helper_fn)) .expect("Vm::native_helper_fn offset must fit i32") @@ -669,6 +677,62 @@ pub(crate) extern "C" fn pd_vm_native_map_get( 1 } +pub(crate) extern "C" fn pd_vm_native_collection_set( + dst: *mut Value, + container: *mut Value, + key: *const Value, + value: *const Value, +) -> i32 { + if dst.is_null() || container.is_null() || key.is_null() || value.is_null() { + store_bridge_error(VmError::JitNative( + "native collection-set helper received null pointer".to_string(), + )); + return STATUS_ERROR; + } + + let container = unsafe { std::ptr::replace(container, Value::Null) }; + let key = unsafe { (&*key).clone() }; + let value = unsafe { (&*value).clone() }; + match crate::builtins::runtime::core::builtin_set_owned(container, key, value) { + Ok(result) => { + let previous = unsafe { std::ptr::replace(dst, result) }; + drop(previous); + STATUS_CONTINUE + } + Err(err) => { + store_bridge_error(err); + STATUS_ERROR + } + } +} + +pub(crate) extern "C" fn pd_vm_native_array_push( + dst: *mut Value, + array: *mut Value, + value: *const Value, +) -> i32 { + if dst.is_null() || array.is_null() || value.is_null() { + store_bridge_error(VmError::JitNative( + "native array-push helper received null pointer".to_string(), + )); + return STATUS_ERROR; + } + + let array = unsafe { std::ptr::replace(array, Value::Null) }; + let value = unsafe { (&*value).clone() }; + match crate::builtins::runtime::core::builtin_array_push_owned(array, value) { + Ok(result) => { + let previous = unsafe { std::ptr::replace(dst, result) }; + drop(previous); + STATUS_CONTINUE + } + Err(err) => { + store_bridge_error(err); + STATUS_ERROR + } + } +} + pub(crate) extern "C" fn pd_vm_native_step(vm: *mut Vm, op: i64, a: i64, b: i64, c: i64) -> i32 { run_step(vm, "step", |vm| { if op == OP_BUILTIN_CALL { diff --git a/src/vm/native/codegen.rs b/src/vm/native/codegen.rs index acaa2cf0..d8a3dd8b 100644 --- a/src/vm/native/codegen.rs +++ b/src/vm/native/codegen.rs @@ -159,6 +159,20 @@ pub(crate) fn collection_get_signature( sig } +#[cfg(feature = "cranelift-jit")] +pub(crate) fn collection_mutation_signature( + pointer_type: cranelift_codegen::ir::Type, + call_conv: cranelift_codegen::isa::CallConv, +) -> Signature { + let mut sig = Signature::new(call_conv); + sig.params.push(AbiParam::new(pointer_type)); + sig.params.push(AbiParam::new(pointer_type)); + sig.params.push(AbiParam::new(pointer_type)); + sig.params.push(AbiParam::new(pointer_type)); + sig.returns.push(AbiParam::new(types::I32)); + sig +} + #[cfg(feature = "cranelift-jit")] pub(crate) fn value_slot_signature( pointer_type: cranelift_codegen::ir::Type, diff --git a/src/vm/native/mod.rs b/src/vm/native/mod.rs index 78d9a648..0e800f8d 100644 --- a/src/vm/native/mod.rs +++ b/src/vm/native/mod.rs @@ -10,23 +10,23 @@ pub(crate) use bridge::{ STATUS_ERROR, STATUS_HALTED, STATUS_LINKED_CONTINUE, STATUS_OUT_OF_FUEL, STATUS_TRACE_EXIT, STATUS_WAITING, STATUS_YIELDED, alloc_byte_buffer_entry_address, alloc_value_buffer_entry_address, aot_call_boundary_interrupt_entry_address, - clear_bridge_error, clear_value_slot_entry_address, clone_value_to_slot_entry_address, - copy_bytes_entry_address, helper_entry_address, helper_entry_offset, - init_null_value_slot_entry_address, interrupt_helper_entry_address, - interrupt_helper_entry_offset, map_get_entry_address, map_has_entry_address, - restore_exit_state_entry_address, restore_sparse_exit_state_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, value_eq_entry_address, - write_heap_value_to_slot_entry_address, zero_bytes_entry_address, + array_push_entry_address, clear_bridge_error, clear_value_slot_entry_address, + clone_value_to_slot_entry_address, collection_set_entry_address, copy_bytes_entry_address, + helper_entry_address, helper_entry_offset, init_null_value_slot_entry_address, + interrupt_helper_entry_address, interrupt_helper_entry_offset, map_get_entry_address, + map_has_entry_address, restore_exit_state_entry_address, + restore_sparse_exit_state_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, + value_eq_entry_address, write_heap_value_to_slot_entry_address, zero_bytes_entry_address, }; #[cfg(feature = "cranelift-jit")] pub(crate) use codegen::{ alloc_buffer_signature, box_heap_value_signature, clone_value_signature, - collection_get_signature, collection_predicate_signature, copy_bytes_signature, - entry_signature, free_buffer_signature, helper_signature, jump_with_status, - pack_shared_signature, restore_exit_signature, sparse_restore_exit_signature, + collection_get_signature, collection_mutation_signature, collection_predicate_signature, + copy_bytes_signature, entry_signature, free_buffer_signature, helper_signature, + jump_with_status, pack_shared_signature, restore_exit_signature, sparse_restore_exit_signature, string_binary_transform_signature, string_contains_signature, string_replace_signature, string_unary_transform_signature, value_eq_signature, value_slot_signature, }; diff --git a/tests/jit/jit_tests.rs b/tests/jit/jit_tests.rs index a30ad75c..92cc9ad1 100644 --- a/tests/jit/jit_tests.rs +++ b/tests/jit/jit_tests.rs @@ -2209,6 +2209,380 @@ fn trace_jit_supports_array_len_get_has_in_ssa() { ); } +#[test] +fn trace_jit_specializes_same_local_array_set_through_loop_back() { + if !native_jit_supported() { + return; + } + + let source = r#" + let mut values = [1, 2, 3]; + let mut i = 0; + while i < 64 { + values[1] = i; + i = i + 1; + } + values[1]; + "#; + + let compiled = compile_source(source).expect("array-set trace compile should succeed"); + let program = force_local_types( + compiled.program, + &[(0, ValueType::Array), (1, ValueType::Int)], + ); + let mut vm = Vm::new(program); + vm.set_jit_config(JitConfig { + enabled: true, + hot_loop_threshold: 1, + max_trace_len: 512, + }); + + let status = vm.run().expect("array-set trace vm should run"); + assert_eq!(status, VmStatus::Halted); + assert_eq!(vm.stack(), &[Value::Int(63)]); + + let dump = vm.dump_jit_info(); + let snapshot = vm.jit_snapshot(); + assert!( + any_trace_op(&snapshot, "array_set"), + "expected array set to remain in SSA, dump:\n{dump}" + ); + let mutation_trace = snapshot + .traces + .iter() + .find(|trace| trace.op_names.iter().any(|op| op == "array_set")) + .expect("array-set trace should be present"); + assert!( + !mutation_trace.has_call && mutation_trace.terminal == JitTraceTerminal::LoopBack, + "array-set trace should remain native through its backedge, dump:\n{dump}" + ); + assert!( + snapshot.metrics.native_loop_back_count > 0, + "array-set loop should stay native across backedges, dump:\n{dump}" + ); + assert_eq!( + snapshot.metrics.helper_fallback_count, 0, + "array-set loop should not need interpreter fallback, dump:\n{dump}" + ); +} + +#[test] +fn trace_jit_array_set_preserves_cow_alias() { + if !native_jit_supported() { + return; + } + + let mut bc = BytecodeBuilder::new(); + let set_call = builtin_call_index("set").expect("set builtin should exist"); + let get_call = builtin_call_index("get").expect("get builtin should exist"); + bc.ldc(0); + bc.stloc(0); + bc.ldloc(0); + bc.stloc(1); + bc.ldc(1); + bc.stloc(2); + + let loop_ip = bc.position(); + bc.ldloc(2); + bc.ldc(2); + bc.clt(); + let exit_branch_ip = bc.position(); + bc.brfalse(0); + + bc.ldloc(0); + bc.ldc(3); + bc.ldloc(2); + bc.ldc(4); + bc.stloc(0); + bc.call(set_call, 3); + bc.stloc(0); + bc.ldloc(2); + bc.ldc(5); + bc.add(); + bc.stloc(2); + bc.br(loop_ip); + + let exit_ip = bc.position(); + bc.ldloc(1); + bc.ldc(3); + bc.call(get_call, 2); + bc.ldloc(0); + bc.ldc(3); + bc.call(get_call, 2); + bc.add(); + bc.ret(); + + let mut code = bc.finish(); + patch_branch_target(&mut code, exit_branch_ip, exit_ip); + let program = Program::new( + vec![ + Value::array(vec![Value::Int(1), Value::Int(2), Value::Int(3)]), + Value::Int(0), + Value::Int(64), + Value::Int(1), + Value::Null, + Value::Int(1), + ], + code, + ) + .with_local_count(3); + let program = force_local_types( + program, + &[ + (0, ValueType::Array), + (1, ValueType::Array), + (2, ValueType::Int), + ], + ); + 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("COW array-set trace should run"), + VmStatus::Halted + ); + assert_eq!(vm.stack(), &[Value::Int(65)]); + let snapshot = vm.jit_snapshot(); + let mutation_trace = snapshot + .traces + .iter() + .find(|trace| trace.op_names.iter().any(|op| op == "array_set")) + .expect("array-set trace should be present"); + assert_eq!(mutation_trace.terminal, JitTraceTerminal::LoopBack); + assert!(!mutation_trace.has_call); +} + +#[test] +fn trace_jit_does_not_consume_non_moved_array_set_container() { + if !native_jit_supported() { + return; + } + + let mut bc = BytecodeBuilder::new(); + let set_call = builtin_call_index("set").expect("set builtin should exist"); + let get_call = builtin_call_index("get").expect("get builtin should exist"); + bc.ldc(0); + bc.stloc(0); + bc.ldc(1); + bc.stloc(1); + + let loop_ip = bc.position(); + bc.ldloc(1); + bc.ldc(2); + bc.clt(); + let exit_branch_ip = bc.position(); + bc.brfalse(0); + + bc.ldloc(0); + bc.ldc(3); + bc.ldloc(1); + bc.call(set_call, 3); + bc.pop(); + bc.ldloc(1); + bc.ldc(4); + bc.add(); + bc.stloc(1); + bc.br(loop_ip); + + let exit_ip = bc.position(); + bc.ldloc(0); + bc.ldc(3); + bc.call(get_call, 2); + bc.ret(); + + let mut code = bc.finish(); + patch_branch_target(&mut code, exit_branch_ip, exit_ip); + let program = Program::new( + vec![ + Value::array(vec![Value::Int(1), Value::Int(2), Value::Int(3)]), + Value::Int(0), + Value::Int(64), + Value::Int(1), + Value::Int(1), + ], + code, + ) + .with_local_count(2); + let program = force_local_types(program, &[(0, ValueType::Array), (1, ValueType::Int)]); + 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("non-moved array-set trace should run"), + VmStatus::Halted + ); + assert_eq!(vm.stack(), &[Value::Int(2)]); + let snapshot = vm.jit_snapshot(); + assert!( + snapshot + .traces + .iter() + .all(|trace| !trace.op_names.iter().any(|op| op == "array_set")), + "non-moved Set must not consume the local container: {}", + vm.dump_jit_info() + ); + assert!(snapshot.traces.iter().any(|trace| trace.has_call)); +} + +#[test] +fn trace_jit_specializes_same_local_array_push_through_loop_back() { + if !native_jit_supported() { + return; + } + + let mut bc = BytecodeBuilder::new(); + let push_call = vm::BuiltinFunction::ArrayPush.call_index(); + let get_call = builtin_call_index("get").expect("get builtin should exist"); + bc.ldc(0); + bc.stloc(0); + bc.ldc(1); + bc.stloc(1); + + let loop_ip = bc.position(); + bc.ldloc(1); + bc.ldc(2); + bc.clt(); + let exit_branch_ip = bc.position(); + bc.brfalse(0); + + bc.ldloc(0); + bc.ldloc(1); + bc.ldc(3); + bc.stloc(0); + bc.call(push_call, 2); + bc.stloc(0); + bc.ldloc(1); + bc.ldc(4); + bc.add(); + bc.stloc(1); + bc.br(loop_ip); + + let exit_ip = bc.position(); + bc.ldloc(0); + bc.ldc(5); + bc.call(get_call, 2); + bc.ret(); + + let mut code = bc.finish(); + patch_branch_target(&mut code, exit_branch_ip, exit_ip); + let program = Program::new( + vec![ + Value::array(Vec::new()), + Value::Int(0), + Value::Int(64), + Value::Null, + Value::Int(1), + Value::Int(63), + ], + code, + ) + .with_local_count(2); + let program = force_local_types(program, &[(0, ValueType::Array), (1, ValueType::Int)]); + let mut vm = Vm::new(program); + vm.set_jit_config(JitConfig { + enabled: true, + hot_loop_threshold: 1, + max_trace_len: 512, + }); + + let status = vm.run().expect("array-push trace vm should run"); + assert_eq!(status, VmStatus::Halted); + assert_eq!(vm.stack(), &[Value::Int(63)]); + + let dump = vm.dump_jit_info(); + let snapshot = vm.jit_snapshot(); + let mutation_trace = snapshot + .traces + .iter() + .find(|trace| trace.op_names.iter().any(|op| op == "array_push")) + .expect("array-push trace should be present"); + assert!( + !mutation_trace.has_call && mutation_trace.terminal == JitTraceTerminal::LoopBack, + "array-push trace should remain native through its backedge, dump:\n{dump}" + ); + assert!( + snapshot.metrics.native_loop_back_count > 0, + "array-push loop should stay native across backedges, dump:\n{dump}" + ); + assert_eq!( + snapshot.metrics.helper_fallback_count, 0, + "array-push loop should not need interpreter fallback, dump:\n{dump}" + ); +} + +#[test] +fn trace_jit_specializes_same_local_map_set_through_loop_back() { + if !native_jit_supported() { + return; + } + + let source = r#" + let mut values = {"key": 0, "drop": 1}; + let mut i = 0; + while i < 64 { + values["key"] = i; + values["drop"] = null; + i = i + 1; + } + values["key"]; + "#; + + let compiled = compile_source(source).expect("map-set trace compile should succeed"); + let program = force_local_types( + compiled.program, + &[(0, ValueType::Map), (1, ValueType::Int)], + ); + let mut vm = Vm::new(program); + vm.set_jit_config(JitConfig { + enabled: true, + hot_loop_threshold: 1, + max_trace_len: 512, + }); + + let status = vm + .run() + .unwrap_or_else(|err| panic!("map-set trace vm failed: {err:?}\n{}", vm.dump_jit_info())); + assert_eq!(status, VmStatus::Halted); + assert_eq!(vm.stack(), &[Value::Int(63)]); + + let dump = vm.dump_jit_info(); + let snapshot = vm.jit_snapshot(); + let mutation_trace = snapshot + .traces + .iter() + .find(|trace| trace.op_names.iter().any(|op| op == "map_set")) + .expect("map-set trace should be present"); + assert_eq!( + mutation_trace + .op_names + .iter() + .filter(|op| op.as_str() == "map_set") + .count(), + 2, + "map-set trace should include overwrite and null-delete mutations, dump:\n{dump}" + ); + assert!( + !mutation_trace.has_call && mutation_trace.terminal == JitTraceTerminal::LoopBack, + "map-set trace should remain native through its backedge, dump:\n{dump}" + ); + assert!( + snapshot.metrics.native_loop_back_count > 0, + "map-set loop should stay native across backedges, dump:\n{dump}" + ); + assert_eq!( + snapshot.metrics.helper_fallback_count, 0, + "map-set loop should not need interpreter fallback, dump:\n{dump}" + ); +} + #[test] fn trace_jit_supports_map_len_get_has_in_ssa() { if !native_jit_supported() {