fix(jit): guard inlined callable identity

Author: fffonion

src/vm/jit/inline.rs

@@ -48,11 +48,14 @@ pub(crate) fn classify_static_inline_candidate(
         return Err(InlineRejectReason::NonRootCaller);
     }
     let source_local = source_local.ok_or(InlineRejectReason::UnknownTarget)?;
-    let binding = program
+    let mut bindings = program
         .root_callable_bindings
         .iter()
-        .find(|binding| binding.local_slot == u16::from(source_local))
-        .ok_or(InlineRejectReason::UnknownTarget)?;
+        .filter(|binding| binding.local_slot == u16::from(source_local));
+    let binding = bindings.next().ok_or(InlineRejectReason::UnknownTarget)?;
+    if bindings.next().is_some() {
+        return Err(InlineRejectReason::PolymorphicTarget);
+    }
     if caller_prototype_id == Some(binding.prototype_id) {
         return Err(InlineRejectReason::Recursive);
     }

src/vm/jit/recorder.rs

@@ -1,5 +1,6 @@
-use std::fmt;
+use std::{fmt, sync::Arc};
 
+use crate::CallableValue;
 use crate::builtins::BuiltinFunction;
 use crate::vm::{OpCode, Program, Value, ValueType, checked_int_div};
 
@@ -1245,6 +1246,55 @@ pub(crate) fn record_trace_with_local_count(
                 if inline_frame.is_none()
                     && let Ok(candidate) = candidate
                 {
+                    let prototype = &program.callable_prototypes[candidate.prototype_id as usize];
+                    let expected_callable = builder
+                        .append_value_inst(
+                            current_block,
+                            ip,
+                            SsaValueRepr::Tagged,
+                            SsaInstKind::Constant(Value::Callable(Arc::new(CallableValue {
+                                prototype_id: candidate.prototype_id,
+                                kind: prototype.kind,
+                                env: None,
+                            }))),
+                        )
+                        .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?;
+                    let callable_matches = builder
+                        .append_value_inst(
+                            current_block,
+                            ip,
+                            SsaValueRepr::Bool,
+                            SsaInstKind::ValueCmpEq {
+                                lhs: callable.value.id,
+                                rhs: expected_callable.id,
+                            },
+                        )
+                        .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?;
+                    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(
+                        &mut builder,
+                        &frame,
+                        &mut inline_frame,
+                        "inline_callable_identity",
+                    )?;
+                    builder
+                        .set_terminator(
+                            current_block,
+                            SsaTerminator::BranchBool {
+                                condition: callable_matches.id,
+                                if_true: SsaBranchTarget::Block {
+                                    target: guarded_block,
+                                    args: guard_args,
+                                },
+                                if_false: SsaBranchTarget::Exit(identity_exit),
+                            },
+                        )
+                        .map_err(|err| TraceRecordError::InvalidIr(err.to_string()))?;
+                    current_block = guarded_block;
+                    frame = guarded_frame;
+                    op_names.push("inline_callable_identity_guard".to_string());
+
                     let operand_base = frame.stack.len() - usize::from(argc) - 1;
                     let mut operands = frame.stack.split_off(operand_base);
                     let _callable = operands.remove(0);

tests/jit/jit_tests.rs

@@ -6453,6 +6453,51 @@ fn trace_jit_inlines_static_leaf_in_root_loop() {
     assert!(vm.dump_jit_info().contains("interpreter fallbacks: 0"));
 }
 
+#[test]
+fn trace_jit_guards_static_inline_callable_identity() {
+    if !native_jit_supported() {
+        return;
+    }
+    let source = r#"
+        fn add_one(value: int) -> int { value + 1 }
+        fn add_ten(value: int) -> int { value + 10 }
+        let mut i = 0;
+        let mut total = 0;
+        while i < 100 {
+            total = add_one(total);
+            i = i + 1;
+        }
+        total;
+    "#;
+    let compiled = compile_source(source).expect("callable identity source should compile");
+    let bindings = compiled.program.root_callable_bindings.clone();
+    let replaced_slot = bindings.first().expect("add_one binding").local_slot;
+    let replacement_id = bindings.get(1).expect("add_ten binding").prototype_id;
+    let replacement_kind = compiled.program.callable_prototypes[replacement_id as usize].kind;
+    let mut vm = Vm::new(compiled.program.with_local_count(compiled.locals));
+    vm.set_local(
+        u8::try_from(replaced_slot).expect("root callable slot should fit u8"),
+        Value::Callable(Arc::new(vm::CallableValue {
+            prototype_id: replacement_id,
+            kind: replacement_kind,
+            env: None,
+        })),
+    )
+    .expect("callable replacement should succeed");
+    vm.set_jit_config(JitConfig {
+        enabled: true,
+        hot_loop_threshold: 1,
+        max_trace_len: 512,
+    });
+
+    assert_eq!(
+        vm.run().expect("identity guard source should run"),
+        VmStatus::Halted
+    );
+    assert_eq!(vm.stack(), &[Value::Int(1_000)]);
+    assert!(vm.jit_native_exec_count() > 0, "{}", vm.dump_jit_info());
+}
+
 #[test]
 fn trace_jit_inlines_array_swap_leaf() {
     if !native_jit_supported() {