ZJIT: Specialize monomorphic DefinedIvar (#15281)

This lets us constant-fold common monomorphic cases.

5 files changed, 172 insertions, 28 deletions

diff --git a/insns.def b/insns.def
index f282b5a8e7..3d13f4cb64 100644
--- a/insns.def
+++ b/insns.def
@@ -745,6 +745,7 @@ definedivar
 ()
 (VALUE val)
 // attr bool leaf  = false;
+// attr bool zjit_profile = true;
 {
     val = Qnil;
     if (!UNDEF_P(vm_getivar(GET_SELF(), id, GET_ISEQ(), ic, NULL, FALSE, Qundef))) {
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index 66126b627c..fe2055d4cc 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -1744,34 +1744,35 @@ pub const YARVINSN_trace_setlocal_WC_1: ruby_vminsn_type = 215;
 pub const YARVINSN_trace_putobject_INT2FIX_0_: ruby_vminsn_type = 216;
 pub const YARVINSN_trace_putobject_INT2FIX_1_: ruby_vminsn_type = 217;
 pub const YARVINSN_zjit_getinstancevariable: ruby_vminsn_type = 218;
-pub const YARVINSN_zjit_send: ruby_vminsn_type = 219;
-pub const YARVINSN_zjit_opt_send_without_block: ruby_vminsn_type = 220;
-pub const YARVINSN_zjit_objtostring: ruby_vminsn_type = 221;
-pub const YARVINSN_zjit_opt_nil_p: ruby_vminsn_type = 222;
-pub const YARVINSN_zjit_invokeblock: ruby_vminsn_type = 223;
-pub const YARVINSN_zjit_opt_plus: ruby_vminsn_type = 224;
-pub const YARVINSN_zjit_opt_minus: ruby_vminsn_type = 225;
-pub const YARVINSN_zjit_opt_mult: ruby_vminsn_type = 226;
-pub const YARVINSN_zjit_opt_div: ruby_vminsn_type = 227;
-pub const YARVINSN_zjit_opt_mod: ruby_vminsn_type = 228;
-pub const YARVINSN_zjit_opt_eq: ruby_vminsn_type = 229;
-pub const YARVINSN_zjit_opt_neq: ruby_vminsn_type = 230;
-pub const YARVINSN_zjit_opt_lt: ruby_vminsn_type = 231;
-pub const YARVINSN_zjit_opt_le: ruby_vminsn_type = 232;
-pub const YARVINSN_zjit_opt_gt: ruby_vminsn_type = 233;
-pub const YARVINSN_zjit_opt_ge: ruby_vminsn_type = 234;
-pub const YARVINSN_zjit_opt_ltlt: ruby_vminsn_type = 235;
-pub const YARVINSN_zjit_opt_and: ruby_vminsn_type = 236;
-pub const YARVINSN_zjit_opt_or: ruby_vminsn_type = 237;
-pub const YARVINSN_zjit_opt_aref: ruby_vminsn_type = 238;
-pub const YARVINSN_zjit_opt_aset: ruby_vminsn_type = 239;
-pub const YARVINSN_zjit_opt_length: ruby_vminsn_type = 240;
-pub const YARVINSN_zjit_opt_size: ruby_vminsn_type = 241;
-pub const YARVINSN_zjit_opt_empty_p: ruby_vminsn_type = 242;
-pub const YARVINSN_zjit_opt_succ: ruby_vminsn_type = 243;
-pub const YARVINSN_zjit_opt_not: ruby_vminsn_type = 244;
-pub const YARVINSN_zjit_opt_regexpmatch2: ruby_vminsn_type = 245;
-pub const VM_INSTRUCTION_SIZE: ruby_vminsn_type = 246;
+pub const YARVINSN_zjit_definedivar: ruby_vminsn_type = 219;
+pub const YARVINSN_zjit_send: ruby_vminsn_type = 220;
+pub const YARVINSN_zjit_opt_send_without_block: ruby_vminsn_type = 221;
+pub const YARVINSN_zjit_objtostring: ruby_vminsn_type = 222;
+pub const YARVINSN_zjit_opt_nil_p: ruby_vminsn_type = 223;
+pub const YARVINSN_zjit_invokeblock: ruby_vminsn_type = 224;
+pub const YARVINSN_zjit_opt_plus: ruby_vminsn_type = 225;
+pub const YARVINSN_zjit_opt_minus: ruby_vminsn_type = 226;
+pub const YARVINSN_zjit_opt_mult: ruby_vminsn_type = 227;
+pub const YARVINSN_zjit_opt_div: ruby_vminsn_type = 228;
+pub const YARVINSN_zjit_opt_mod: ruby_vminsn_type = 229;
+pub const YARVINSN_zjit_opt_eq: ruby_vminsn_type = 230;
+pub const YARVINSN_zjit_opt_neq: ruby_vminsn_type = 231;
+pub const YARVINSN_zjit_opt_lt: ruby_vminsn_type = 232;
+pub const YARVINSN_zjit_opt_le: ruby_vminsn_type = 233;
+pub const YARVINSN_zjit_opt_gt: ruby_vminsn_type = 234;
+pub const YARVINSN_zjit_opt_ge: ruby_vminsn_type = 235;
+pub const YARVINSN_zjit_opt_ltlt: ruby_vminsn_type = 236;
+pub const YARVINSN_zjit_opt_and: ruby_vminsn_type = 237;
+pub const YARVINSN_zjit_opt_or: ruby_vminsn_type = 238;
+pub const YARVINSN_zjit_opt_aref: ruby_vminsn_type = 239;
+pub const YARVINSN_zjit_opt_aset: ruby_vminsn_type = 240;
+pub const YARVINSN_zjit_opt_length: ruby_vminsn_type = 241;
+pub const YARVINSN_zjit_opt_size: ruby_vminsn_type = 242;
+pub const YARVINSN_zjit_opt_empty_p: ruby_vminsn_type = 243;
+pub const YARVINSN_zjit_opt_succ: ruby_vminsn_type = 244;
+pub const YARVINSN_zjit_opt_not: ruby_vminsn_type = 245;
+pub const YARVINSN_zjit_opt_regexpmatch2: ruby_vminsn_type = 246;
+pub const VM_INSTRUCTION_SIZE: ruby_vminsn_type = 247;
 pub type ruby_vminsn_type = u32;
 pub type rb_iseq_callback = ::std::option::Option<
     unsafe extern "C" fn(arg1: *const rb_iseq_t, arg2: *mut ::std::os::raw::c_void),
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 3866da52b9..40c6092e56 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -2820,6 +2820,38 @@ impl Function {
                         };
                         self.make_equal_to(insn_id, replacement);
                     }
+                    Insn::DefinedIvar { self_val, id, pushval, state } => {
+                        let frame_state = self.frame_state(state);
+                        let Some(recv_type) = self.profiled_type_of_at(self_val, frame_state.insn_idx) else {
+                            // No (monomorphic/skewed polymorphic) profile info
+                            self.push_insn_id(block, insn_id); continue;
+                        };
+                        if recv_type.flags().is_immediate() {
+                            // Instance variable lookups on immediate values are always nil
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        assert!(recv_type.shape().is_valid());
+                        if !recv_type.flags().is_t_object() {
+                            // Check if the receiver is a T_OBJECT
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        if recv_type.shape().is_too_complex() {
+                            // too-complex shapes can't use index access
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        let self_val = self.push_insn(block, Insn::GuardType { val: self_val, guard_type: types::HeapBasicObject, state });
+                        let _ = self.push_insn(block, Insn::GuardShape { val: self_val, shape: recv_type.shape(), state });
+                        let mut ivar_index: u16 = 0;
+                        let replacement = if unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
+                            self.push_insn(block, Insn::Const { val: Const::Value(pushval) })
+                        } else {
+                            // If there is no IVAR index, then the ivar was undefined when we
+                            // entered the compiler.  That means we can just return nil for this
+                            // shape + iv name
+                            self.push_insn(block, Insn::Const { val: Const::Value(Qnil) })
+                        };
+                        self.make_equal_to(insn_id, replacement);
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -4839,6 +4871,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
             // profiled cfp->self.
             if opcode == YARVINSN_getinstancevariable || opcode == YARVINSN_trace_getinstancevariable {
                 profiles.profile_self(&exit_state, self_param);
+            } else if opcode == YARVINSN_definedivar || opcode == YARVINSN_trace_definedivar {
+                profiles.profile_self(&exit_state, self_param);
             } else if opcode == YARVINSN_invokeblock || opcode == YARVINSN_trace_invokeblock {
                 if get_option!(stats) {
                     let iseq_insn_idx = exit_state.insn_idx;
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 083820c0da..b1ab8a0605 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -3374,6 +3374,113 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn test_specialize_monomorphic_definedivar_true() {
+        eval("
+            @foo = 4
+            def test = defined?(@foo)
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v15:HeapBasicObject = GuardType v6, HeapBasicObject
+          v16:HeapBasicObject = GuardShape v15, 0x1000
+          v17:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+          CheckInterrupts
+          Return v17
+        ");
+    }
+
+    #[test]
+    fn test_specialize_monomorphic_definedivar_false() {
+        eval("
+            def test = defined?(@foo)
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v15:HeapBasicObject = GuardType v6, HeapBasicObject
+          v16:HeapBasicObject = GuardShape v15, 0x1000
+          v17:NilClass = Const Value(nil)
+          CheckInterrupts
+          Return v17
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_definedivar_with_t_data() {
+        eval("
+            class C < Range
+              def test = defined?(@a)
+            end
+            obj = C.new 0, 1
+            obj.instance_variable_set(:@a, 1)
+            obj.test
+            TEST = C.instance_method(:test)
+        ");
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v10:StringExact|NilClass = DefinedIvar v6, :@a
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_dont_specialize_polymorphic_definedivar() {
+        set_call_threshold(3);
+        eval("
+            class C
+              def test = defined?(@a)
+            end
+            obj = C.new
+            obj.instance_variable_set(:@a, 1)
+            obj.test
+            obj = C.new
+            obj.instance_variable_set(:@b, 1)
+            obj.test
+            TEST = C.instance_method(:test)
+        ");
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v10:StringExact|NilClass = DefinedIvar v6, :@a
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
     fn test_elide_freeze_with_frozen_hash() {
         eval("
             def test = {}.freeze
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
index 8c8190609d..10afdf2cc6 100644
--- a/zjit/src/profile.rs
+++ b/zjit/src/profile.rs
@@ -82,6 +82,7 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
         YARVINSN_opt_aset  => profile_operands(profiler, profile, 3),
         YARVINSN_opt_not   => profile_operands(profiler, profile, 1),
         YARVINSN_getinstancevariable => profile_self(profiler, profile),
+        YARVINSN_definedivar   => profile_self(profiler, profile),
         YARVINSN_opt_regexpmatch2    => profile_operands(profiler, profile, 2),
         YARVINSN_objtostring   => profile_operands(profiler, profile, 1),
         YARVINSN_opt_length    => profile_operands(profiler, profile, 1),