From bd0d0d3b387b4ed87dde8a13a62c83eb9aafdb3a Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Sat, 4 Mar 2023 22:13:20 -0800
Subject: [PATCH] Reset chain_depth in more places

---
 lib/ruby_vm/mjit/compiler.rb      | 4 ++++
 lib/ruby_vm/mjit/insn_compiler.rb | 1 +
 2 files changed, 5 insertions(+)

diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 2f67136102..85e6b5fa06 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -239,6 +239,10 @@ module RubyVM::MJIT
 
         case status = @insn_compiler.compile(jit, ctx, asm, insn)
         when KeepCompiling
+          # For now, reset the chain depth after each instruction as only the
+          # first instruction in the block can concern itself with the depth.
+          ctx.chain_depth = 0
+
           index += insn.len
         when EndBlock
           # TODO: pad nops if entry exit exists (not needed for x86_64?)
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index ecdd00a617..a37c1c1b1c 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -3623,6 +3623,7 @@ module RubyVM::MJIT
         return_ctx = ctx.dup
         return_ctx.stack_size -= argc # Pop args. blockarg has been popped
         return_ctx.sp_offset = 1 # SP is in the position after popping a receiver and arguments
+        return_ctx.chain_depth = 0
         branch_stub = BranchStub.new(
           iseq: jit.iseq,
           shape: Default,