feat(jit): build native trace graph regions

src/lib.rs

@@ -9,8 +9,8 @@ pub mod debugger;
 #[cfg(feature = "runtime")]
 pub mod jit {
     pub use crate::vm::jit::{
-        JitAttempt, JitConfig, JitMetrics, JitNyiDoc, JitNyiReason, JitSnapshot, JitTrace,
-        JitTraceTerminal, TraceJitEngine,
+        JitAttempt, 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, JitMetrics, JitNyiDoc, JitNyiReason, JitSnapshot, JitTrace,
-    JitTraceTerminal, TraceJitEngine,
+    JitAttempt, JitConfig, JitExitProfile, JitMetrics, JitNyiDoc, JitNyiReason, JitSnapshot,
+    JitTrace, JitTraceTerminal, TraceJitEngine,
 };
 #[cfg(feature = "runtime")]
 pub use vm::diagnostics::render_vm_error;

src/vm/host.rs

@@ -313,9 +313,10 @@ impl HostFunctionRegistry {
 
     /// Registers a static args-only host function that always returns one value synchronously.
     ///
-    /// Returning no value, `Halt`, `Yield`, or `Pending` violates the contract and is reported as
-    /// a host error. When appropriate, the native JIT may keep traces active across the call
-    /// boundary.
+    /// The returned [`Value`] must match the return type declared by the corresponding host
+    /// import. Returning a different type is reported as [`VmError::TypeMismatch`]. Returning no
+    /// value, `Halt`, `Yield`, or `Pending` violates the contract and is reported as a host error.
+    /// When appropriate, the native JIT may keep traces active across the call boundary.
     pub fn register_static_non_yielding_args(
         &mut self,
         name: impl Into<String>,
@@ -494,6 +495,41 @@ pub(crate) fn require_non_yielding_host_value(outcome: CallOutcome) -> VmResult<
     }
 }
 
+pub(crate) fn validate_non_yielding_host_value(
+    value: Value,
+    expected: Option<ValueType>,
+) -> VmResult<Value> {
+    let valid = matches!(
+        (expected, &value),
+        (None | Some(ValueType::Unknown), _)
+            | (Some(ValueType::Null), Value::Null)
+            | (Some(ValueType::Int), Value::Int(_))
+            | (Some(ValueType::Float), Value::Float(_))
+            | (Some(ValueType::Bool), Value::Bool(_))
+            | (Some(ValueType::String), Value::String(_))
+            | (Some(ValueType::Bytes), Value::Bytes(_))
+            | (Some(ValueType::Array), Value::Array(_))
+            | (Some(ValueType::Map), Value::Map(_))
+            | (Some(ValueType::Callable), Value::Callable(_))
+    );
+    if valid {
+        return Ok(value);
+    }
+    let expected = match expected.expect("known expected host return type") {
+        ValueType::Unknown => unreachable!(),
+        ValueType::Null => "null",
+        ValueType::Int => "int",
+        ValueType::Float => "float",
+        ValueType::Bool => "bool",
+        ValueType::String => "string",
+        ValueType::Bytes => "bytes",
+        ValueType::Array => "array",
+        ValueType::Map => "map",
+        ValueType::Callable => "callable",
+    };
+    Err(VmError::TypeMismatch(expected))
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(super) struct WaitingHostOp {
     pub(super) op_id: HostOpId,
@@ -573,7 +609,10 @@ impl Vm {
 
     /// Registers a static args-only host function that always returns one value synchronously.
     ///
-    /// Returning no value, `Halt`, `Yield`, or `Pending` violates the contract and is a host error.
+    /// When used to resolve a declared host import, the returned [`Value`] must match that
+    /// import's return type. Returning a different type is reported as
+    /// [`VmError::TypeMismatch`]. Returning no value, `Halt`, `Yield`, or `Pending` violates the
+    /// contract and is a host error.
     pub fn register_static_non_yielding_args_function(
         &mut self,
         function: StaticHostArgsFunction,
@@ -723,8 +762,9 @@ impl Vm {
     /// Binds a static args-only host function that always returns one value synchronously.
     ///
     /// This is equivalent to [`Vm::bind_static_args_function`] except that the VM may keep
-    /// native JIT traces active across the call boundary. Returning no value, `Halt`, `Yield`,
-    /// or `Pending` violates the contract and is reported as a host error.
+    /// native JIT traces active across the call boundary. The returned [`Value`] must match the
+    /// return type declared by the corresponding host import. Returning a different type, no
+    /// value, `Halt`, `Yield`, or `Pending` violates the contract and is reported as a host error.
     pub fn bind_static_non_yielding_args_function(
         &mut self,
         name: impl Into<String>,
@@ -948,9 +988,19 @@ impl Vm {
             return self.execute_builtin_call_from_stack(builtin, argc, call_ip);
         }
 
+        let expected_return_type = self
+            .program
+            .imports
+            .get(usize::from(index))
+            .map(|import| import.return_type);
         let resolved_index = self.resolve_call_target(index, argc_u8)?;
         if self.bound_host_function_uses_args_slice(resolved_index)? {
-            self.execute_bound_args_host_function(resolved_index, argc, call_ip)
+            self.execute_bound_args_host_function(
+                resolved_index,
+                argc,
+                call_ip,
+                expected_return_type,
+            )
         } else if self.bound_host_function_uses_stack_borrow(resolved_index)? {
             self.execute_bound_stack_host_function(resolved_index, argc, call_ip)
         } else {
@@ -975,7 +1025,7 @@ impl Vm {
             })?;
         let argc = argc_u8 as usize;
         if self.bound_host_function_uses_args_slice(resolved_index)? {
-            self.execute_bound_args_host_function(resolved_index, argc, call_ip)
+            self.execute_bound_args_host_function(resolved_index, argc, call_ip, None)
         } else if self.bound_host_function_uses_stack_borrow(resolved_index)? {
             self.execute_bound_stack_host_function(resolved_index, argc, call_ip)
         } else {
@@ -1492,6 +1542,7 @@ impl Vm {
         resolved_index: u16,
         argc: usize,
         call_ip: usize,
+        expected_return_type: Option<ValueType>,
     ) -> VmResult<HostCallExecOutcome> {
         let arg_start = self
             .stack
@@ -1520,6 +1571,7 @@ impl Vm {
         let outcome = outcome?;
         if non_yielding {
             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);
             return Ok(HostCallExecOutcome::Returned);

src/vm/jit/deopt.rs

@@ -1,5 +1,7 @@
 #![allow(dead_code)]
-use super::ir::{SsaExit, SsaMaterialization, SsaValue, SsaValueId, SsaValueRepr};
+use super::ir::{
+    SsaExit, SsaExitId, SsaMaterialization, SsaTrace, SsaValue, SsaValueId, SsaValueRepr,
+};
 
 pub(crate) fn materialize_ssa_values(
     values: impl IntoIterator<Item = SsaValue>,
@@ -41,3 +43,211 @@ pub(crate) fn exit_inputs(exit: &SsaExit) -> Vec<SsaValueId> {
     }
     out
 }
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) struct SideTraceImport {
+    pub(crate) parent_exit: SsaExitId,
+    pub(crate) stack_depth: usize,
+    pub(crate) local_count: usize,
+    pub(crate) dirty_locals: Vec<bool>,
+    pub(crate) args: Vec<SsaMaterialization>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum SideTraceImportError {
+    UnknownParentExit(SsaExitId),
+    ExitIpMismatch { parent: usize, child: usize },
+    StackDepthMismatch { parent: usize, child: usize },
+    LocalCountMismatch { parent: usize, child: usize },
+    InvalidChildEntry,
+}
+
+pub(crate) fn side_trace_import(
+    parent: &SsaTrace,
+    parent_exit: SsaExitId,
+    child: &SsaTrace,
+) -> Result<SideTraceImport, SideTraceImportError> {
+    let exit = parent
+        .exits
+        .iter()
+        .find(|exit| exit.id == parent_exit)
+        .ok_or(SideTraceImportError::UnknownParentExit(parent_exit))?;
+    if exit.exit_ip != child.root_ip {
+        return Err(SideTraceImportError::ExitIpMismatch {
+            parent: exit.exit_ip,
+            child: child.root_ip,
+        });
+    }
+    if exit.stack.len() != child.entry_stack_depth {
+        return Err(SideTraceImportError::StackDepthMismatch {
+            parent: exit.stack.len(),
+            child: child.entry_stack_depth,
+        });
+    }
+    let child_entry = child
+        .blocks
+        .get(child.entry.index())
+        .ok_or(SideTraceImportError::InvalidChildEntry)?;
+    let child_local_count = child_entry
+        .params
+        .len()
+        .checked_sub(child.entry_stack_depth)
+        .ok_or(SideTraceImportError::InvalidChildEntry)?;
+    if exit.locals.len() != child_local_count {
+        return Err(SideTraceImportError::LocalCountMismatch {
+            parent: exit.locals.len(),
+            child: child_local_count,
+        });
+    }
+
+    Ok(SideTraceImport {
+        parent_exit,
+        stack_depth: exit.stack.len(),
+        local_count: exit.locals.len(),
+        dirty_locals: exit.dirty_locals.clone(),
+        args: exit.stack.iter().chain(&exit.locals).cloned().collect(),
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::vm::jit::ir::{SsaTerminator, SsaTraceBuilder, SsaValueRepr};
+
+    #[test]
+    fn side_trace_import_maps_parent_stack_then_locals_to_child_entry() {
+        let mut parent = SsaTraceBuilder::new(0, 1);
+        let parent_entry = parent.entry();
+        let stack = parent
+            .append_param(parent_entry, SsaValueRepr::Tagged, "stack0".to_string())
+            .unwrap();
+        let local = parent
+            .append_param(parent_entry, SsaValueRepr::I64, "local0".to_string())
+            .unwrap();
+        let exit_id = parent.add_exit(
+            12,
+            vec![SsaMaterialization::Value(stack.id)],
+            vec![SsaMaterialization::BoxInt(local.id)],
+            vec![true],
+        );
+        parent
+            .set_terminator(parent_entry, SsaTerminator::Exit { exit: exit_id })
+            .unwrap();
+        let parent = parent.finish();
+
+        let mut child = SsaTraceBuilder::new(12, 1);
+        let child_entry = child.entry();
+        child
+            .append_param(child_entry, SsaValueRepr::Tagged, "stack0".to_string())
+            .unwrap();
+        child
+            .append_param(child_entry, SsaValueRepr::Tagged, "local0".to_string())
+            .unwrap();
+        let child_exit = child.add_exit(13, Vec::new(), Vec::new(), Vec::new());
+        child
+            .set_terminator(child_entry, SsaTerminator::Exit { exit: child_exit })
+            .unwrap();
+        let child = child.finish();
+
+        let import = side_trace_import(&parent, exit_id, &child).unwrap();
+
+        assert_eq!(import.parent_exit, exit_id);
+        assert_eq!(import.stack_depth, 1);
+        assert_eq!(import.local_count, 1);
+        assert_eq!(import.dirty_locals, vec![true]);
+        assert_eq!(
+            import.args,
+            vec![
+                SsaMaterialization::Value(stack.id),
+                SsaMaterialization::BoxInt(local.id),
+            ]
+        );
+    }
+
+    #[test]
+    fn side_trace_import_rejects_mismatched_exit_ip() {
+        let mut parent = SsaTraceBuilder::new(0, 0);
+        let parent_entry = parent.entry();
+        let exit_id = parent.add_exit(12, Vec::new(), Vec::new(), Vec::new());
+        parent
+            .set_terminator(parent_entry, SsaTerminator::Exit { exit: exit_id })
+            .unwrap();
+        let parent = parent.finish();
+
+        let mut child = SsaTraceBuilder::new(13, 0);
+        let child_entry = child.entry();
+        let child_exit = child.add_exit(14, Vec::new(), Vec::new(), Vec::new());
+        child
+            .set_terminator(child_entry, SsaTerminator::Exit { exit: child_exit })
+            .unwrap();
+        let child = child.finish();
+
+        assert_eq!(
+            side_trace_import(&parent, exit_id, &child),
+            Err(SideTraceImportError::ExitIpMismatch {
+                parent: 12,
+                child: 13,
+            })
+        );
+    }
+
+    #[test]
+    fn side_trace_import_rejects_mismatched_stack_depth() {
+        let mut parent = SsaTraceBuilder::new(0, 0);
+        let parent_entry = parent.entry();
+        let exit_id = parent.add_exit(12, Vec::new(), Vec::new(), Vec::new());
+        parent
+            .set_terminator(parent_entry, SsaTerminator::Exit { exit: exit_id })
+            .unwrap();
+        let parent = parent.finish();
+
+        let mut child = SsaTraceBuilder::new(12, 1);
+        let child_entry = child.entry();
+        child
+            .append_param(child_entry, SsaValueRepr::Tagged, "stack0".to_string())
+            .unwrap();
+        let child_exit = child.add_exit(13, Vec::new(), Vec::new(), Vec::new());
+        child
+            .set_terminator(child_entry, SsaTerminator::Exit { exit: child_exit })
+            .unwrap();
+        let child = child.finish();
+
+        assert_eq!(
+            side_trace_import(&parent, exit_id, &child),
+            Err(SideTraceImportError::StackDepthMismatch {
+                parent: 0,
+                child: 1,
+            })
+        );
+    }
+
+    #[test]
+    fn side_trace_import_rejects_mismatched_local_count() {
+        let mut parent = SsaTraceBuilder::new(0, 0);
+        let parent_entry = parent.entry();
+        let exit_id = parent.add_exit(12, Vec::new(), Vec::new(), Vec::new());
+        parent
+            .set_terminator(parent_entry, SsaTerminator::Exit { exit: exit_id })
+            .unwrap();
+        let parent = parent.finish();
+
+        let mut child = SsaTraceBuilder::new(12, 0);
+        let child_entry = child.entry();
+        child
+            .append_param(child_entry, SsaValueRepr::Tagged, "local0".to_string())
+            .unwrap();
+        let child_exit = child.add_exit(13, Vec::new(), Vec::new(), Vec::new());
+        child
+            .set_terminator(child_entry, SsaTerminator::Exit { exit: child_exit })
+            .unwrap();
+        let child = child.finish();
+
+        assert_eq!(
+            side_trace_import(&parent, exit_id, &child),
+            Err(SideTraceImportError::LocalCountMismatch {
+                parent: 0,
+                child: 1,
+            })
+        );
+    }
+}