fix(jit): guard inlined callable schemas

Author: fffonion

src/vm/jit/inline.rs

@@ -16,6 +16,7 @@ pub(crate) enum InlineRejectReason {
     HostTarget,
     CapturedCallable,
     ArityMismatch,
+    SchemaUnproven,
     Recursive,
     NestedScriptCall,
     YieldingCall,

src/vm/jit/ir.rs

@@ -105,6 +105,10 @@ pub(crate) enum SsaInstKind {
     CloneTagged {
         input: SsaValueId,
     },
+    ValueIsType {
+        input: SsaValueId,
+        tag: ValueType,
+    },
     UnboxInt {
         input: SsaValueId,
     },
@@ -394,6 +398,7 @@ impl SsaInstKind {
             Self::HostCall { args, .. } => args.clone(),
 
             Self::CloneTagged { input }
+            | Self::ValueIsType { input, .. }
             | Self::UnboxInt { input }
             | Self::UnboxFloat { input }
             | Self::UnboxBool { input }
@@ -1121,6 +1126,9 @@ fn render_inst_kind(kind: &SsaInstKind) -> String {
     match kind {
         SsaInstKind::Constant(value) => format!("const {value:?}"),
         SsaInstKind::CloneTagged { input } => format!("clone_tagged {input}"),
+        SsaInstKind::ValueIsType { input, tag } => {
+            format!("value_is_type {input}, {tag:?}")
+        }
         SsaInstKind::UnboxInt { input } => format!("unbox_int {input}"),
         SsaInstKind::UnboxFloat { input } => format!("unbox_float {input}"),
         SsaInstKind::UnboxBool { input } => format!("unbox_bool {input}"),

src/vm/jit/native/lower.rs

@@ -1274,6 +1274,7 @@ fn ssa_trace_supported(ssa: &SsaTrace) -> bool {
                 inst.kind,
                 SsaInstKind::Constant(_)
                     | SsaInstKind::CloneTagged { .. }
+                    | SsaInstKind::ValueIsType { .. }
                     | SsaInstKind::UnboxHeapPtr { .. }
                     | SsaInstKind::UnboxInt { .. }
                     | SsaInstKind::UnboxFloat { .. }
@@ -2090,6 +2091,29 @@ fn lower_ssa_inst(
             )?;
             out
         }
+        SsaInstKind::ValueIsType { input, tag } => {
+            let input = *values.get(input).ok_or_else(|| {
+                VmError::JitNative("SSA type predicate input missing".to_string())
+            })?;
+            let expected_tag = match tag {
+                ValueType::Null => layout.value.null_tag,
+                ValueType::Int => layout.value.int_tag,
+                ValueType::Float => layout.value.float_tag,
+                ValueType::Bool => layout.value.bool_tag,
+                ValueType::String => layout.value.string_tag,
+                ValueType::Bytes => layout.value.bytes_tag,
+                ValueType::Array => layout.value.array_tag,
+                ValueType::Map => layout.value.map_tag,
+                ValueType::Callable | ValueType::Unknown => {
+                    return Err(VmError::JitNative(format!(
+                        "unsupported SSA type predicate tag {tag:?}"
+                    )));
+                }
+            };
+            let actual_tag = ssa_load_tag_i32(b, layout.value, input);
+            b.ins()
+                .icmp_imm(IntCC::Equal, actual_tag, i64::from(expected_tag))
+        }
         SsaInstKind::UnboxInt { input } => {
             let input = *values
                 .get(input)

src/vm/jit/recorder.rs

@@ -2,14 +2,15 @@ use std::{fmt, sync::Arc};
 
 use crate::CallableValue;
 use crate::builtins::BuiltinFunction;
+use crate::compiler::TypeSchema;
 use crate::vm::{OpCode, Program, Value, ValueType, checked_int_div};
 
 use super::JitTraceTerminal;
 use super::deopt::materialize_ssa_values;
-use super::inline::{InlineCandidate, classify_static_inline_candidate};
+use super::inline::{InlineCandidate, InlineRejectReason, classify_static_inline_candidate};
 use super::ir::{
     SsaBranchTarget, SsaInstKind, SsaMaterialization, SsaTerminator, SsaTrace, SsaTraceBuilder,
-    SsaValue, SsaValueRepr, VirtualFrameSnapshot,
+    SsaValue, SsaValueId, SsaValueRepr, VirtualFrameSnapshot,
 };
 
 #[derive(Clone, Debug, PartialEq)]
@@ -263,6 +264,99 @@ struct SymbolicValue {
     info: ValueInfo,
 }
 
+fn inline_schema_guard_type(schema: &TypeSchema) -> Option<Option<ValueType>> {
+    match schema {
+        TypeSchema::Unknown | TypeSchema::GenericParam(_) => Some(None),
+        TypeSchema::Int => Some(Some(ValueType::Int)),
+        TypeSchema::Float => Some(Some(ValueType::Float)),
+        TypeSchema::Bool => Some(Some(ValueType::Bool)),
+        TypeSchema::String => Some(Some(ValueType::String)),
+        TypeSchema::Bytes => Some(Some(ValueType::Bytes)),
+        TypeSchema::Named(_, _) | TypeSchema::Map(_) | TypeSchema::Object(_) => {
+            Some(Some(ValueType::Map))
+        }
+        TypeSchema::Array(_) | TypeSchema::ArrayTuple(_) | TypeSchema::ArrayTupleRest { .. } => {
+            Some(Some(ValueType::Array))
+        }
+        TypeSchema::Null
+        | TypeSchema::Number
+        | TypeSchema::Optional(_)
+        | TypeSchema::Callable { .. } => None,
+    }
+}
+
+fn inline_argument_schemas_supported(
+    arguments: &[SymbolicValue],
+    schema: Option<&TypeSchema>,
+) -> bool {
+    let Some(TypeSchema::Callable { params, .. }) = schema else {
+        return schema.is_none();
+    };
+    params.len() == arguments.len()
+        && params.iter().zip(arguments).all(|(schema, argument)| {
+            let Some(guard_type) = inline_schema_guard_type(schema) else {
+                return false;
+            };
+            match (guard_type, argument.info.repr) {
+                (None, _) | (Some(_), SsaValueRepr::Tagged) => true,
+                (Some(ValueType::Int), SsaValueRepr::I64)
+                | (Some(ValueType::Float), SsaValueRepr::F64)
+                | (Some(ValueType::Bool), SsaValueRepr::Bool) => true,
+                (Some(expected), SsaValueRepr::HeapPtr(actual)) => expected == actual,
+                _ => false,
+            }
+        })
+}
+
+fn append_inline_argument_schema_guards(
+    builder: &mut SsaTraceBuilder,
+    block: super::ir::SsaBlockId,
+    ip: usize,
+    arguments: &[SymbolicValue],
+    schema: Option<&TypeSchema>,
+) -> Result<Option<SsaValueId>, TraceRecordError> {
+    let Some(TypeSchema::Callable { params, .. }) = schema else {
+        return Ok(None);
+    };
+    let mut guard = None;
+    for (schema, argument) in params.iter().zip(arguments) {
+        let Some(Some(expected)) = inline_schema_guard_type(schema) else {
+            continue;
+        };
+        if argument.info.repr != SsaValueRepr::Tagged {
+            continue;
+        }
+        let predicate = builder
+            .append_value_inst(
+                block,
+                ip,
+                SsaValueRepr::Bool,
+                SsaInstKind::ValueIsType {
+                    input: argument.value.id,
+                    tag: expected,
+                },
+            )
+            .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?;
+        guard = Some(if let Some(previous) = guard {
+            builder
+                .append_value_inst(
+                    block,
+                    ip,
+                    SsaValueRepr::Bool,
+                    SsaInstKind::BoolAnd {
+                        lhs: previous,
+                        rhs: predicate.id,
+                    },
+                )
+                .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?
+                .id
+        } else {
+            predicate.id
+        });
+    }
+    Ok(guard)
+}
+
 #[derive(Clone, Debug, PartialEq)]
 struct SymbolicFrame {
     stack: Vec<SymbolicValue>,
@@ -1241,12 +1335,30 @@ pub(crate) fn record_trace_with_local_count(
                     callable.info.source_local,
                     argc,
                     max_trace_len.saturating_sub(cursor.recorded_ops),
-                );
+                )
+                .and_then(|candidate| {
+                    let prototype = &program.callable_prototypes[candidate.prototype_id as usize];
+                    let argument_start = frame.stack.len() - usize::from(argc);
+                    inline_argument_schemas_supported(
+                        &frame.stack[argument_start..],
+                        prototype.schema.as_ref(),
+                    )
+                    .then_some(candidate)
+                    .ok_or(InlineRejectReason::SchemaUnproven)
+                });
                 let inline_reject_reason = candidate.as_ref().err().copied();
                 if inline_frame.is_none()
                     && let Ok(candidate) = candidate
                 {
                     let prototype = &program.callable_prototypes[candidate.prototype_id as usize];
+                    let argument_start = frame.stack.len() - usize::from(argc);
+                    let schema_guard = append_inline_argument_schema_guards(
+                        &mut builder,
+                        current_block,
+                        ip,
+                        &frame.stack[argument_start..],
+                        prototype.schema.as_ref(),
+                    )?;
                     let expected_callable = builder
                         .append_value_inst(
                             current_block,
@@ -1270,6 +1382,22 @@ pub(crate) fn record_trace_with_local_count(
                             },
                         )
                         .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?;
+                    let inline_guard = if let Some(schema_guard) = schema_guard {
+                        builder
+                            .append_value_inst(
+                                current_block,
+                                ip,
+                                SsaValueRepr::Bool,
+                                SsaInstKind::BoolAnd {
+                                    lhs: schema_guard,
+                                    rhs: callable_matches.id,
+                                },
+                            )
+                            .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?
+                            .id
+                    } else {
+                        callable_matches.id
+                    };
                     let identity_exit =
                         add_symbolic_exit(&mut builder, ip, &frame, inline_frame.as_ref());
                     let (guarded_block, guarded_frame, guard_args) = continue_with_inline_frame(
@@ -1282,7 +1410,7 @@ pub(crate) fn record_trace_with_local_count(
                         .set_terminator(
                             current_block,
                             SsaTerminator::BranchBool {
-                                condition: callable_matches.id,
+                                condition: inline_guard,
                                 if_true: SsaBranchTarget::Block {
                                     target: guarded_block,
                                     args: guard_args,

src/vm/jit/region.rs

@@ -212,6 +212,7 @@ fn remap_inst_inputs(
             }
         }
         SsaInstKind::CloneTagged { input }
+        | SsaInstKind::ValueIsType { input, .. }
         | SsaInstKind::UnboxInt { input }
         | SsaInstKind::UnboxFloat { input }
         | SsaInstKind::UnboxBool { input }

tests/jit/jit_tests.rs

@@ -6498,6 +6498,60 @@ fn trace_jit_guards_static_inline_callable_identity() {
     assert!(vm.jit_native_exec_count() > 0, "{}", vm.dump_jit_info());
 }
 
+#[test]
+fn trace_jit_preserves_inline_callable_argument_schema_checks() {
+    if !native_jit_supported() {
+        return;
+    }
+    let source = r#"
+        fn ignore(value: int) -> int { 1 }
+        let mut i = 0;
+        let value: int = 7;
+        while i < 100 {
+            i = i + ignore(value);
+        }
+        i;
+    "#;
+    let compiled = compile_source(source).expect("callable schema source should compile");
+    let value_slot = compiled
+        .program
+        .debug
+        .as_ref()
+        .expect("debug metadata")
+        .locals
+        .iter()
+        .find(|local| local.name == "value")
+        .expect("value local")
+        .index;
+    let mut vm = Vm::new(compiled.program.with_local_count(compiled.locals));
+    vm.set_fuel_check_interval(1).expect("fuel interval");
+    vm.set_fuel(1);
+    loop {
+        assert_eq!(
+            vm.run().expect("initialization step should run"),
+            VmStatus::Yielded
+        );
+        if vm.locals().get(value_slot as usize) == Some(&Value::Int(7)) {
+            break;
+        }
+        vm.recharge_fuel(1).expect("fuel recharge");
+    }
+    vm.set_local(value_slot, Value::Bool(true))
+        .expect("dynamic value replacement should succeed");
+    vm.clear_fuel();
+    vm.set_jit_config(JitConfig {
+        enabled: true,
+        hot_loop_threshold: 1,
+        max_trace_len: 512,
+    });
+
+    let error = vm.run().expect_err("invalid callable argument must fail");
+    assert!(
+        matches!(error, vm::VmError::TypeMismatch("callable argument schema")),
+        "unexpected error: {error:?}"
+    );
+}
+
 #[test]
 fn trace_jit_inlines_array_swap_leaf() {
     if !native_jit_supported() {