ZJIT: Fix backtraces on opt_new (#14461)

13 files changed, 131 insertions, 43 deletions

diff --git a/.github/auto_request_review.yml b/.github/auto_request_review.yml
index 264e6ef177..c4c94681f0 100644
--- a/.github/auto_request_review.yml
+++ b/.github/auto_request_review.yml
@@ -10,7 +10,6 @@ files:
   'zjit/src/cruby_bindings.inc.rs': []
   'doc/zjit*': [team:jit]
   'test/ruby/test_zjit*': [team:jit]
-  'test/.excludes-zjit/*': [team:jit]
   'defs/jit.mk': [team:jit]
 options:
   ignore_draft: true
diff --git a/.github/workflows/zjit-macos.yml b/.github/workflows/zjit-macos.yml
index bb244dd1b2..892b605ce8 100644
--- a/.github/workflows/zjit-macos.yml
+++ b/.github/workflows/zjit-macos.yml
@@ -111,7 +111,6 @@ jobs:
           RUN_OPTS="$RUN_OPTS"
           SPECOPTS="$SPECOPTS"
           TESTOPTS="$TESTOPTS"
-          TEST_EXCLUDES="$TEST_EXCLUDES"
         timeout-minutes: 60
         env:
           RUBY_TESTOPTS: '-q --tty=no'
@@ -119,7 +118,6 @@ jobs:
           SYNTAX_SUGGEST_TIMEOUT: '5'
           PRECHECK_BUNDLED_GEMS: 'no'
           TESTS: ${{ matrix.tests }}
-          TEST_EXCLUDES: '--excludes-dir=../src/test/.excludes-zjit --name=!/memory_leak/'
         continue-on-error: ${{ matrix.continue-on-test_task || false }}
 
   result:
diff --git a/.github/workflows/zjit-ubuntu.yml b/.github/workflows/zjit-ubuntu.yml
index 3861e2481f..b9efcaaf5f 100644
--- a/.github/workflows/zjit-ubuntu.yml
+++ b/.github/workflows/zjit-ubuntu.yml
@@ -149,13 +149,12 @@ jobs:
         run: >-
           make -s ${{ matrix.test_task }} ${TESTS:+TESTS="$TESTS"}
           RUN_OPTS="$RUN_OPTS" MSPECOPT=--debug SPECOPTS="$SPECOPTS"
-          TESTOPTS="$TESTOPTS" TEST_EXCLUDES="$TEST_EXCLUDES"
+          TESTOPTS="$TESTOPTS"
           ZJIT_BINDGEN_DIFF_OPTS="$ZJIT_BINDGEN_DIFF_OPTS"
         timeout-minutes: 90
         env:
           RUBY_TESTOPTS: '-q --tty=no'
           TEST_BUNDLED_GEMS_ALLOW_FAILURES: ''
-          TEST_EXCLUDES: '--excludes-dir=../src/test/.excludes-zjit --name=!/memory_leak/'
           PRECHECK_BUNDLED_GEMS: 'no'
           SYNTAX_SUGGEST_TIMEOUT: '5'
           ZJIT_BINDGEN_DIFF_OPTS: '--exit-code'
diff --git a/doc/zjit.md b/doc/zjit.md
index 04ca0a8bb4..b5b605d5cb 100644
--- a/doc/zjit.md
+++ b/doc/zjit.md
@@ -96,16 +96,6 @@ use `make`.
 
 </details>
 
-### make zjit-test-all
-
-```
-make zjit-test-all
-```
-
-This command runs all Ruby tests under `/test/ruby/` with ZJIT enabled.
-
-Certain tests are excluded under `/test/.excludes-zjit`.
-
 ### test/ruby/test\_zjit.rb
 
 This command runs Ruby execution tests.
diff --git a/test/.excludes-zjit/TestERBCore.rb b/test/.excludes-zjit/TestERBCore.rb
deleted file mode 100644
index 9ab398de6f..0000000000
--- a/test/.excludes-zjit/TestERBCore.rb
+++ /dev/null
@@ -1 +0,0 @@
-exclude(:test_invalid_trim_mode, 'Test fails with ZJIT')
diff --git a/test/.excludes-zjit/TestERBCoreWOStrScan.rb b/test/.excludes-zjit/TestERBCoreWOStrScan.rb
deleted file mode 100644
index 9ab398de6f..0000000000
--- a/test/.excludes-zjit/TestERBCoreWOStrScan.rb
+++ /dev/null
@@ -1 +0,0 @@
-exclude(:test_invalid_trim_mode, 'Test fails with ZJIT')
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 9bfe2c0c00..280a7503d4 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -766,6 +766,37 @@ class TestZJIT < Test::Unit::TestCase
     }, insns: [:opt_ge], call_threshold: 2
   end
 
+  def test_opt_new_does_not_push_frame
+    assert_compiles 'nil', %q{
+      class Foo
+        attr_reader :backtrace
+
+        def initialize
+          @backtrace = caller
+        end
+      end
+      def test = Foo.new
+
+      foo = test
+      foo.backtrace.find do |frame|
+        frame.include?('Class#new')
+      end
+    }, insns: [:opt_new]
+  end
+
+  def test_opt_new_with_redefinition
+    assert_compiles '"foo"', %q{
+      class Foo
+        def self.new = "foo"
+
+        def initialize = raise("unreachable")
+      end
+      def test = Foo.new
+
+      test
+    }, insns: [:opt_new]
+  end
+
   def test_new_hash_empty
     assert_compiles '{}', %q{
       def test = {}
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index eee20cd9d4..cc22e4a2e0 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -2417,6 +2417,12 @@ vm_method_cfunc_is(const rb_iseq_t *iseq, CALL_DATA cd, VALUE recv, cfunc_type f
     return check_cfunc(cme, func);
 }
 
+int
+rb_vm_method_cfunc_is(const rb_iseq_t *iseq, CALL_DATA cd, VALUE recv, cfunc_type func)
+{
+    return vm_method_cfunc_is(iseq, cd, recv, func);
+}
+
 #define check_cfunc(me, func) check_cfunc(me, make_cfunc_type(func))
 #define vm_method_cfunc_is(iseq, cd, recv, func) vm_method_cfunc_is(iseq, cd, recv, make_cfunc_type(func))
 
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index e3609e237e..56400e20cd 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -417,6 +417,8 @@ fn main() {
         // From internal/object.h
         .allowlist_function("rb_class_allocate_instance")
         .allowlist_function("rb_obj_equal")
+        .allowlist_function("rb_class_new_instance_pass_kw")
+        .allowlist_function("rb_obj_alloc")
 
         // From gc.h and internal/gc.h
         .allowlist_function("rb_obj_info")
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 5c13d278fb..1db1ddc510 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -349,6 +349,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::NewRange { low, high, flag, state } => gen_new_range(jit, asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
+        Insn::ObjectAlloc { val, state } => gen_object_alloc(asm, opnd!(val), &function.frame_state(*state)),
         Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)),
         // concatstrings shouldn't have 0 strings
         // If it happens we abort the compilation for now
@@ -385,6 +386,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::FixnumAnd { left, right } => gen_fixnum_and(asm, opnd!(left), opnd!(right)),
         Insn::FixnumOr { left, right } => gen_fixnum_or(asm, opnd!(left), opnd!(right)),
         Insn::IsNil { val } => gen_isnil(asm, opnd!(val)),
+        &Insn::IsMethodCfunc { val, cd, cfunc } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc),
         Insn::Test { val } => gen_test(asm, opnd!(val)),
         Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state)),
         Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(jit, asm, opnd!(val), *expected, &function.frame_state(*state)),
@@ -1190,6 +1192,11 @@ fn gen_new_range_fixnum(
     asm_ccall!(asm, rb_range_new, low, high, (flag as i64).into())
 }
 
+fn gen_object_alloc(asm: &mut Assembler, val: lir::Opnd, state: &FrameState) -> lir::Opnd {
+    gen_prepare_call_with_gc(asm, state);
+    asm_ccall!(asm, rb_obj_alloc, val)
+}
+
 /// Compile code that exits from JIT code with a return value
 fn gen_return(asm: &mut Assembler, val: lir::Opnd) {
     // Pop the current frame (ec->cfp++)
@@ -1292,6 +1299,13 @@ fn gen_isnil(asm: &mut Assembler, val: lir::Opnd) -> lir::Opnd {
     asm.csel_e(Opnd::Imm(1), Opnd::Imm(0))
 }
 
+fn gen_is_method_cfunc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, cd: *const rb_call_data, cfunc: *const u8) -> lir::Opnd {
+    unsafe extern "C" {
+        fn rb_vm_method_cfunc_is(iseq: IseqPtr, cd: *const rb_call_data, recv: VALUE, cfunc: *const u8) -> VALUE;
+    }
+    asm_ccall!(asm, rb_vm_method_cfunc_is, VALUE(jit.iseq as usize).into(), (cd as usize).into(), val, (cfunc as usize).into())
+}
+
 fn gen_anytostring(asm: &mut Assembler, val: lir::Opnd, str: lir::Opnd, state: &FrameState) -> lir::Opnd {
     gen_prepare_call_with_gc(asm, state);
 
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index 0523c00ede..ae394f9c60 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -806,7 +806,13 @@ unsafe extern "C" {
     pub fn rb_id2str(id: ID) -> VALUE;
     pub fn rb_sym2str(symbol: VALUE) -> VALUE;
     pub fn rb_class2name(klass: VALUE) -> *const ::std::os::raw::c_char;
+    pub fn rb_class_new_instance_pass_kw(
+        argc: ::std::os::raw::c_int,
+        argv: *const VALUE,
+        klass: VALUE,
+    ) -> VALUE;
     pub fn rb_obj_is_kind_of(obj: VALUE, klass: VALUE) -> VALUE;
+    pub fn rb_obj_alloc(klass: VALUE) -> VALUE;
     pub fn rb_obj_frozen_p(obj: VALUE) -> VALUE;
     pub fn rb_class_inherited_p(scion: VALUE, ascendant: VALUE) -> VALUE;
     pub fn rb_backref_get() -> VALUE;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index af44c86cbe..6bec2ab552 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -519,11 +519,16 @@ pub enum Insn {
 
     HashDup { val: InsnId, state: InsnId },
 
+    /// Allocate an instance of the `val` class without calling `#initialize` on it.
+    ObjectAlloc { val: InsnId, state: InsnId },
+
     /// Check if the value is truthy and "return" a C boolean. In reality, we will likely fuse this
     /// with IfTrue/IfFalse in the backend to generate jcc.
     Test { val: InsnId },
     /// Return C `true` if `val` is `Qnil`, else `false`.
     IsNil { val: InsnId },
+    /// Return C `true` if `val`'s method on cd resolves to the cfunc.
+    IsMethodCfunc { val: InsnId, cd: *const rb_call_data, cfunc: *const u8 },
     Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, state: InsnId },
     GetConstantPath { ic: *const iseq_inline_constant_cache, state: InsnId },
 
@@ -761,6 +766,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             }
             Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
             Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
+            Insn::ObjectAlloc { val, .. } => { write!(f, "ObjectAlloc {val}") }
             Insn::StringCopy { val, .. } => { write!(f, "StringCopy {val}") }
             Insn::StringConcat { strings, .. } => {
                 write!(f, "StringConcat")?;
@@ -796,6 +802,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             }
             Insn::Test { val } => { write!(f, "Test {val}") }
             Insn::IsNil { val } => { write!(f, "IsNil {val}") }
+            Insn::IsMethodCfunc { val, cd, .. } => { write!(f, "IsMethodCFunc {val}, :{}", ruby_call_method_name(*cd)) }
             Insn::Jump(target) => { write!(f, "Jump {target}") }
             Insn::IfTrue { val, target } => { write!(f, "IfTrue {val}, {target}") }
             Insn::IfFalse { val, target } => { write!(f, "IfFalse {val}, {target}") }
@@ -1273,6 +1280,7 @@ impl Function {
             &ToRegexp { opt, ref values, state } => ToRegexp { opt, values: find_vec!(values), state },
             &Test { val } => Test { val: find!(val) },
             &IsNil { val } => IsNil { val: find!(val) },
+            &IsMethodCfunc { val, cd, cfunc } => IsMethodCfunc { val: find!(val), cd, cfunc },
             Jump(target) => Jump(find_branch_edge!(target)),
             &IfTrue { val, ref target } => IfTrue { val: find!(val), target: find_branch_edge!(target) },
             &IfFalse { val, ref target } => IfFalse { val: find!(val), target: find_branch_edge!(target) },
@@ -1333,6 +1341,7 @@ impl Function {
             &InvokeBuiltin { bf, ref args, state, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, return_type },
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
+            &ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
             &CCall { cfun, ref args, name, return_type, elidable } => CCall { cfun, args: find_vec!(args), name, return_type, elidable },
             &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) },
             &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state },
@@ -1407,6 +1416,7 @@ impl Function {
             Insn::IsNil { val } if self.is_a(*val, types::NilClass) => Type::from_cbool(true),
             Insn::IsNil { val } if !self.type_of(*val).could_be(types::NilClass) => Type::from_cbool(false),
             Insn::IsNil { .. } => types::CBool,
+            Insn::IsMethodCfunc { .. } => types::CBool,
             Insn::StringCopy { .. } => types::StringExact,
             Insn::StringIntern { .. } => types::Symbol,
             Insn::StringConcat { .. } => types::StringExact,
@@ -1417,6 +1427,7 @@ impl Function {
             Insn::HashDup { .. } => types::HashExact,
             Insn::NewRange { .. } => types::RangeExact,
             Insn::NewRangeFixnum { .. } => types::RangeExact,
+            Insn::ObjectAlloc { .. } => types::HeapObject,
             Insn::CCall { return_type, .. } => *return_type,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
@@ -2173,12 +2184,14 @@ impl Function {
             | &Insn::Return { val }
             | &Insn::Test { val }
             | &Insn::SetLocal { val, .. }
-            | &Insn::IsNil { val } =>
+            | &Insn::IsNil { val }
+            | &Insn::IsMethodCfunc { val, .. } =>
                 worklist.push_back(val),
             &Insn::SetGlobal { val, state, .. }
             | &Insn::Defined { v: val, state, .. }
             | &Insn::StringIntern { val, state }
             | &Insn::StringCopy { val, state, .. }
+            | &Insn::ObjectAlloc { val, state }
             | &Insn::GuardType { val, state, .. }
             | &Insn::GuardBitEquals { val, state, .. }
             | &Insn::GuardShape { val, state, .. }
@@ -3340,16 +3353,30 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     state.stack_pop()?;
                 }
                 YARVINSN_opt_new => {
-                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
-                    fun.push_insn(block, Insn::CheckInterrupts { state: exit_id });
-                    let offset = get_arg(pc, 1).as_i64();
-                    let target_idx = insn_idx_at_offset(insn_idx, offset);
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let dst = get_arg(pc, 1).as_i64();
+
+                    // Check if #new resolves to rb_class_new_instance_pass_kw.
+                    // TODO: Guard on a profiled class and add a patch point for #new redefinition
+                    let argc = unsafe { vm_ci_argc((*cd).ci) } as usize;
+                    let val = state.stack_topn(argc)?;
+                    let test_id = fun.push_insn(block, Insn::IsMethodCfunc { val, cd, cfunc: rb_class_new_instance_pass_kw as *const u8 });
+
+                    // Jump to the fallback block if it's not the expected function.
+                    // Skip CheckInterrupts since the #new call will do it very soon anyway.
+                    let target_idx = insn_idx_at_offset(insn_idx, dst);
                     let target = insn_idx_to_block[&target_idx];
-                    // Skip the fast-path and go straight to the fallback code. We will let the
-                    // optimizer take care of the converting Class#new->alloc+initialize instead.
-                    fun.push_insn(block, Insn::Jump(BranchEdge { target, args: state.as_args(self_param) }));
+                    let _branch_id = fun.push_insn(block, Insn::IfFalse {
+                        val: test_id,
+                        target: BranchEdge { target, args: state.as_args(self_param) }
+                    });
                     queue.push_back((state.clone(), target, target_idx, local_inval));
-                    break;  // Don't enqueue the next block as a successor
+
+                    // Move on to the fast path
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let insn_id = fun.push_insn(block, Insn::ObjectAlloc { val, state: exit_id });
+                    state.stack_setn(argc, insn_id);
+                    state.stack_setn(argc + 1, insn_id);
                 }
                 YARVINSN_jump => {
                     let offset = get_arg(pc, 0).as_i64();
@@ -5201,14 +5228,18 @@ mod tests {
         bb0(v0:BasicObject):
           v5:BasicObject = GetConstantPath 0x1000
           v6:NilClass = Const Value(nil)
+          v7:CBool = IsMethodCFunc v5, :new
+          IfFalse v7, bb1(v0, v6, v5)
+          v10:HeapObject = ObjectAlloc v5
+          v12:BasicObject = SendWithoutBlock v10, :initialize
           CheckInterrupts
-          Jump bb1(v0, v6, v5)
-        bb1(v10:BasicObject, v11:NilClass, v12:BasicObject):
-          v15:BasicObject = SendWithoutBlock v12, :new
-          Jump bb2(v10, v15, v11)
-        bb2(v17:BasicObject, v18:BasicObject, v19:NilClass):
+          Jump bb2(v0, v10, v12)
+        bb1(v16:BasicObject, v17:NilClass, v18:BasicObject):
+          v21:BasicObject = SendWithoutBlock v18, :new
+          Jump bb2(v16, v21, v17)
+        bb2(v23:BasicObject, v24:BasicObject, v25:BasicObject):
           CheckInterrupts
-          Return v18
+          Return v24
         ");
     }
 
@@ -7842,12 +7873,20 @@ mod opt_tests {
         bb0(v0:BasicObject):
           PatchPoint SingleRactorMode
           PatchPoint StableConstantNames(0x1000, C)
-          v28:Class[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+          v34:Class[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           v6:NilClass = Const Value(nil)
+          v7:CBool = IsMethodCFunc v34, :new
+          IfFalse v7, bb1(v0, v6, v34)
+          v10:HeapObject = ObjectAlloc v34
+          v12:BasicObject = SendWithoutBlock v10, :initialize
           CheckInterrupts
-          v15:BasicObject = SendWithoutBlock v28, :new
+          Jump bb2(v0, v10, v12)
+        bb1(v16:BasicObject, v17:NilClass, v18:Class[VALUE(0x1008)]):
+          v21:BasicObject = SendWithoutBlock v18, :new
+          Jump bb2(v16, v21, v17)
+        bb2(v23:BasicObject, v24:BasicObject, v25:BasicObject):
           CheckInterrupts
-          Return v15
+          Return v24
         ");
     }
 
@@ -7867,13 +7906,23 @@ mod opt_tests {
         bb0(v0:BasicObject):
           PatchPoint SingleRactorMode
           PatchPoint StableConstantNames(0x1000, C)
-          v30:Class[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+          v36:Class[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           v6:NilClass = Const Value(nil)
           v7:Fixnum[1] = Const Value(1)
+          v8:CBool = IsMethodCFunc v36, :new
+          IfFalse v8, bb1(v0, v6, v36, v7)
+          v11:HeapObject = ObjectAlloc v36
+          PatchPoint MethodRedefined(C@0x1008, initialize@0x1010, cme:0x1018)
+          v38:HeapObject[class_exact:C] = GuardType v11, HeapObject[class_exact:C]
+          v39:BasicObject = SendWithoutBlockDirect v38, :initialize (0x1040), v7
+          CheckInterrupts
+          Jump bb2(v0, v11, v39)
+        bb1(v17:BasicObject, v18:NilClass, v19:Class[VALUE(0x1008)], v20:Fixnum[1]):
+          v23:BasicObject = SendWithoutBlock v19, :new, v20
+          Jump bb2(v17, v23, v18)
+        bb2(v25:BasicObject, v26:BasicObject, v27:BasicObject):
           CheckInterrupts
-          v17:BasicObject = SendWithoutBlock v30, :new, v7
-          CheckInterrupts
-          Return v17
+          Return v26
         ");
     }
 
diff --git a/zjit/zjit.mk b/zjit/zjit.mk
index c39abf2826..be989bdecd 100644
--- a/zjit/zjit.mk
+++ b/zjit/zjit.mk
@@ -50,10 +50,6 @@ zjit-check:
 	$(MAKE) zjit-test
 	$(MAKE) test-all TESTS='$(top_srcdir)/test/ruby/test_zjit.rb'
 
-.PHONY: zjit-test-all
-zjit-test-all:
-	$(MAKE) test-all RUST_BACKTRACE=1 TEST_EXCLUDES='--excludes-dir=$(top_srcdir)/test/.excludes-zjit --name=!/memory_leak/' RUN_OPTS='--zjit-call-threshold=1'
-
 ZJIT_BINDGEN_DIFF_OPTS =
 
 # Generate Rust bindings. See source for details.