perf(jit): short function hot loop inline

src/builtins/runtime/core.rs

@@ -713,7 +713,7 @@ pub(crate) fn builtin_set_map_shared_impl(
     entries
 }
 
-fn ensure_supported_map_key(key: &Value) -> VmResult<()> {
+pub(crate) fn ensure_supported_map_key(key: &Value) -> VmResult<()> {
     if matches!(key, Value::Callable(_)) {
         return Err(VmError::HostError(
             "callable values are not supported as map keys".to_string(),

src/bytecode.rs

@@ -539,13 +539,15 @@ pub struct HostImport {
 pub(crate) struct DecodedInstructionData {
     pub(crate) ldc_values: Box<[Option<Value>]>,
     pub(crate) jump_targets: Box<[Option<usize>]>,
+    pub(crate) valid_jump_targets: Box<[bool]>,
     pub(crate) local_indices: Box<[Option<u8>]>,
 }
 
 impl DecodedInstructionData {
     fn build(program: &Program) -> Self {
         let mut ldc_values = vec![None; program.code.len()];
         let mut jump_targets = vec![None; program.code.len()];
+        let mut valid_jump_targets = vec![false; program.code.len()];
         let mut local_indices = vec![None; program.code.len()];
         let mut ip = 0usize;
         while ip < program.code.len() {
@@ -562,8 +564,32 @@ impl DecodedInstructionData {
                     }
                 }
                 OpCode::Br | OpCode::Brfalse => {
-                    if let Some(target) = read_u32_at(&program.code, ip + 1) {
-                        jump_targets[ip] = Some(target as usize);
+                    if let Some(target) =
+                        read_u32_at(&program.code, ip + 1).map(|target| target as usize)
+                    {
+                        jump_targets[ip] = Some(target);
+                        if target >= program.code.len() {
+                            ip = ip.saturating_add(1 + opcode.operand_len());
+                            continue;
+                        }
+                        let source_owner = program
+                            .function_regions
+                            .iter()
+                            .find(|region| {
+                                region.start_ip as usize <= ip && ip < region.end_ip as usize
+                            })
+                            .and_then(|region| region.prototype_id);
+                        let target_owner = program
+                            .function_regions
+                            .iter()
+                            .find(|region| {
+                                region.start_ip as usize <= target
+                                    && target < region.end_ip as usize
+                            })
+                            .and_then(|region| region.prototype_id);
+                        if source_owner == target_owner {
+                            valid_jump_targets[ip] = true;
+                        }
                     }
                 }
                 OpCode::Ldloc | OpCode::Stloc => {
@@ -578,6 +604,7 @@ impl DecodedInstructionData {
         Self {
             ldc_values: ldc_values.into_boxed_slice(),
             jump_targets: jump_targets.into_boxed_slice(),
+            valid_jump_targets: valid_jump_targets.into_boxed_slice(),
             local_indices: local_indices.into_boxed_slice(),
         }
     }

src/lib.rs

@@ -9,8 +9,8 @@ pub mod debugger;
 #[cfg(feature = "runtime")]
 pub mod jit {
     pub use crate::vm::jit::{
-        JitAttempt, JitConfig, JitExitProfile, JitMetrics, JitNyiDoc, JitNyiReason, JitSnapshot,
-        JitTrace, JitTraceTerminal, TraceJitEngine,
+        JitAttempt, JitCallSiteProfile, JitConfig, JitExitProfile, JitMetrics, JitNyiDoc,
+        JitNyiReason, JitSnapshot, JitTrace, JitTraceTerminal, TraceJitEngine,
     };
 }
 #[cfg(feature = "cli")]
@@ -81,8 +81,8 @@ pub use debugger::{
 };
 #[cfg(feature = "runtime")]
 pub use jit::{
-    JitAttempt, JitConfig, JitExitProfile, JitMetrics, JitNyiDoc, JitNyiReason, JitSnapshot,
-    JitTrace, JitTraceTerminal, TraceJitEngine,
+    JitAttempt, JitCallSiteProfile, JitConfig, JitExitProfile, JitMetrics, JitNyiDoc, JitNyiReason,
+    JitSnapshot, JitTrace, JitTraceTerminal, TraceJitEngine,
 };
 #[cfg(feature = "runtime")]
 pub use vm::diagnostics::render_vm_error;

src/vm/host.rs

@@ -994,6 +994,20 @@ impl Vm {
             .get(usize::from(index))
             .map(|import| import.return_type);
         let resolved_index = self.resolve_call_target(index, argc_u8)?;
+        if let Some(function) =
+            self.host_functions
+                .get(resolved_index as usize)
+                .and_then(|function| match function {
+                    VmHostFunction::ArgsStaticNonYielding(function) => Some(*function),
+                    _ => None,
+                })
+        {
+            return self.execute_static_non_yielding_args_host_function(
+                function,
+                argc,
+                expected_return_type,
+            );
+        }
         if self.bound_host_function_uses_args_slice(resolved_index)? {
             self.execute_bound_args_host_function(
                 resolved_index,
@@ -1537,6 +1551,28 @@ impl Vm {
         ))
     }
 
+    #[inline(always)]
+    fn execute_static_non_yielding_args_host_function(
+        &mut self,
+        function: StaticHostArgsFunction,
+        argc: usize,
+        expected_return_type: Option<ValueType>,
+    ) -> VmResult<HostCallExecOutcome> {
+        let arg_start = self
+            .stack
+            .len()
+            .checked_sub(argc)
+            .ok_or(VmError::StackUnderflow)?;
+        self.call_depth += 1;
+        let outcome = function(&self.stack[arg_start..]);
+        self.call_depth = self.call_depth.saturating_sub(1);
+        let value = require_non_yielding_host_value(outcome?)?;
+        let value = validate_non_yielding_host_value(value, expected_return_type)?;
+        self.stack.truncate(arg_start);
+        self.stack.push(value);
+        Ok(HostCallExecOutcome::Returned)
+    }
+
     pub(super) fn execute_bound_args_host_function(
         &mut self,
         resolved_index: u16,

src/vm/jit/deopt.rs

@@ -41,6 +41,20 @@ pub(crate) fn exit_inputs(exit: &SsaExit) -> Vec<SsaValueId> {
             out.push(value);
         }
     }
+    for frame in &exit.virtual_frames {
+        for materialization in frame.operand_stack.iter().chain(&frame.locals) {
+            let value = match materialization {
+                SsaMaterialization::Value(value)
+                | SsaMaterialization::BoxInt(value)
+                | SsaMaterialization::BoxFloat(value)
+                | SsaMaterialization::BoxBool(value) => *value,
+                SsaMaterialization::BoxHeapPtr { value, .. } => *value,
+            };
+            if !out.contains(&value) {
+                out.push(value);
+            }
+        }
+    }
     out
 }
 
@@ -59,6 +73,7 @@ pub(crate) enum SideTraceImportError {
     ExitIpMismatch { parent: usize, child: usize },
     StackDepthMismatch { parent: usize, child: usize },
     LocalCountMismatch { parent: usize, child: usize },
+    VirtualFramesUnsupported { count: usize },
     InvalidChildEntry,
 }
 
@@ -72,6 +87,11 @@ pub(crate) fn side_trace_import(
         .iter()
         .find(|exit| exit.id == parent_exit)
         .ok_or(SideTraceImportError::UnknownParentExit(parent_exit))?;
+    if !exit.virtual_frames.is_empty() {
+        return Err(SideTraceImportError::VirtualFramesUnsupported {
+            count: exit.virtual_frames.len(),
+        });
+    }
     if exit.exit_ip != child.root_ip {
         return Err(SideTraceImportError::ExitIpMismatch {
             parent: exit.exit_ip,
@@ -250,4 +270,47 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn exit_inputs_include_virtual_frame_values_once_in_frame_order() {
+        use crate::vm::jit::ir::VirtualFrameSnapshot;
+
+        let mut builder = SsaTraceBuilder::new(0, 0);
+        let entry = builder.entry();
+        let caller = builder
+            .append_param(entry, SsaValueRepr::Tagged, "caller")
+            .unwrap();
+        let callee_stack = builder
+            .append_param(entry, SsaValueRepr::I64, "callee_stack")
+            .unwrap();
+        let callee_local = builder
+            .append_param(entry, SsaValueRepr::Bool, "callee_local")
+            .unwrap();
+        let exit_id = builder.add_exit_with_virtual_frames(
+            20,
+            vec![SsaMaterialization::Value(caller.id)],
+            Vec::new(),
+            Vec::new(),
+            vec![VirtualFrameSnapshot {
+                prototype_id: 1,
+                call_ip: 10,
+                return_ip: 12,
+                resume_ip: 20,
+                operand_stack: vec![SsaMaterialization::BoxInt(callee_stack.id)],
+                locals: vec![
+                    SsaMaterialization::Value(caller.id),
+                    SsaMaterialization::BoxBool(callee_local.id),
+                ],
+                dirty_locals: vec![true, true],
+            }],
+        );
+        builder
+            .set_terminator(entry, SsaTerminator::Exit { exit: exit_id })
+            .unwrap();
+        let trace = builder.finish();
+        assert_eq!(
+            exit_inputs(&trace.exits[0]),
+            vec![caller.id, callee_stack.id, callee_local.id]
+        );
+    }
 }