RJIT: Avoid retaining unrelated local variables in memory

1 files changed, 87 insertions, 63 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index e03db73069..feca1b1cb4 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -1924,26 +1924,30 @@ module RubyVM::RJIT
         end
 
         # Jump to target0 on jnz
-        branch_stub.compile = proc do |branch_asm|
-          branch_asm.comment("branchif #{branch_stub.shape}")
-          branch_asm.stub(branch_stub) do
-            case branch_stub.shape
-            in Default
-              branch_asm.jnz(branch_stub.target0.address)
-              branch_asm.jmp(branch_stub.target1.address)
-            in Next0
-              branch_asm.jz(branch_stub.target1.address)
-            in Next1
-              branch_asm.jnz(branch_stub.target0.address)
-            end
-          end
-        end
+        branch_stub.compile = compile_branchif(branch_stub)
         branch_stub.compile.call(asm)
       end
 
       EndBlock
     end
 
+    def compile_branchif(branch_stub) # Proc escapes arguments in memory
+      proc do |branch_asm|
+        branch_asm.comment("branchif #{branch_stub.shape}")
+        branch_asm.stub(branch_stub) do
+          case branch_stub.shape
+          in Default
+            branch_asm.jnz(branch_stub.target0.address)
+            branch_asm.jmp(branch_stub.target1.address)
+          in Next0
+            branch_asm.jz(branch_stub.target1.address)
+          in Next1
+            branch_asm.jnz(branch_stub.target0.address)
+          end
+        end
+      end
+    end
+
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
@@ -1985,26 +1989,30 @@ module RubyVM::RJIT
         end
 
         # Jump to target0 on jz
-        branch_stub.compile = proc do |branch_asm|
-          branch_asm.comment("branchunless #{branch_stub.shape}")
-          branch_asm.stub(branch_stub) do
-            case branch_stub.shape
-            in Default
-              branch_asm.jz(branch_stub.target0.address)
-              branch_asm.jmp(branch_stub.target1.address)
-            in Next0
-              branch_asm.jnz(branch_stub.target1.address)
-            in Next1
-              branch_asm.jz(branch_stub.target0.address)
-            end
-          end
-        end
+        branch_stub.compile = compile_branchunless(branch_stub)
         branch_stub.compile.call(asm)
       end
 
       EndBlock
     end
 
+    def compile_branchunless(branch_stub) # Proc escapes arguments in memory
+      proc do |branch_asm|
+        branch_asm.comment("branchunless #{branch_stub.shape}")
+        branch_asm.stub(branch_stub) do
+          case branch_stub.shape
+          in Default
+            branch_asm.jz(branch_stub.target0.address)
+            branch_asm.jmp(branch_stub.target1.address)
+          in Next0
+            branch_asm.jnz(branch_stub.target1.address)
+          in Next1
+            branch_asm.jz(branch_stub.target0.address)
+          end
+        end
+      end
+    end
+
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
@@ -2045,26 +2053,30 @@ module RubyVM::RJIT
         end
 
         # Jump to target0 on je
-        branch_stub.compile = proc do |branch_asm|
-          branch_asm.comment("branchnil #{branch_stub.shape}")
-          branch_asm.stub(branch_stub) do
-            case branch_stub.shape
-            in Default
-              branch_asm.je(branch_stub.target0.address)
-              branch_asm.jmp(branch_stub.target1.address)
-            in Next0
-              branch_asm.jne(branch_stub.target1.address)
-            in Next1
-              branch_asm.je(branch_stub.target0.address)
-            end
-          end
-        end
+        branch_stub.compile = compile_branchnil(branch_stub)
         branch_stub.compile.call(asm)
       end
 
       EndBlock
     end
 
+    def compile_branchnil(branch_stub) # Proc escapes arguments in memory
+      proc do |branch_asm|
+        branch_asm.comment("branchnil #{branch_stub.shape}")
+        branch_asm.stub(branch_stub) do
+          case branch_stub.shape
+          in Default
+            branch_asm.je(branch_stub.target0.address)
+            branch_asm.jmp(branch_stub.target1.address)
+          in Next0
+            branch_asm.jne(branch_stub.target1.address)
+          in Next1
+            branch_asm.je(branch_stub.target0.address)
+          end
+        end
+      end
+    end
+
     # once
 
     # @param jit [RubyVM::RJIT::JITState]
@@ -3625,21 +3637,25 @@ module RubyVM::RJIT
           @exit_compiler.compile_branch_stub(deeper, ocb_asm, branch_stub, true)
           @ocb.write(ocb_asm)
         end
-        branch_stub.compile = proc do |branch_asm|
-          # Not using `asm.comment` here since it's usually put before cmp/test before this.
-          branch_asm.stub(branch_stub) do
-            case branch_stub.shape
-            in Default
-              branch_asm.public_send(opcode, branch_stub.target0.address)
-            end
-          end
-        end
+        branch_stub.compile = compile_jit_chain_guard(branch_stub, opcode:)
         branch_stub.compile.call(asm)
       else
         asm.public_send(opcode, side_exit)
       end
     end
 
+    def compile_jit_chain_guard(branch_stub, opcode:) # Proc escapes arguments in memory
+      proc do |branch_asm|
+        # Not using `asm.comment` here since it's usually put before cmp/test before this.
+        branch_asm.stub(branch_stub) do
+          case branch_stub.shape
+          in Default
+            branch_asm.public_send(opcode, branch_stub.target0.address)
+          end
+        end
+      end
+    end
+
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
@@ -5628,16 +5644,7 @@ module RubyVM::RJIT
           @exit_compiler.compile_branch_stub(return_ctx, ocb_asm, branch_stub, true)
           @ocb.write(ocb_asm)
         end
-        branch_stub.compile = proc do |branch_asm|
-          branch_asm.comment('set jit_return to callee CFP')
-          branch_asm.stub(branch_stub) do
-            case branch_stub.shape
-            in Default
-              branch_asm.mov(:rax, branch_stub.target0.address)
-              branch_asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:jit_return)], :rax)
-            end
-          end
-        end
+        branch_stub.compile = compile_jit_return(branch_stub, cfp_offset:)
         branch_stub.compile.call(asm)
       end
 
@@ -5648,6 +5655,19 @@ module RubyVM::RJIT
       asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], cfp_reg)
     end
 
+    def compile_jit_return(branch_stub, cfp_offset:) # Proc escapes arguments in memory
+      proc do |branch_asm|
+        branch_asm.comment('set jit_return to callee CFP')
+        branch_asm.stub(branch_stub) do
+          case branch_stub.shape
+          in Default
+            branch_asm.mov(:rax, branch_stub.target0.address)
+            branch_asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:jit_return)], :rax)
+          end
+        end
+      end
+    end
+
     # CALLER_SETUP_ARG: Return CantCompile if not supported
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
@@ -5868,7 +5888,12 @@ module RubyVM::RJIT
         @exit_compiler.compile_branch_stub(ctx, ocb_asm, branch_stub, true)
         @ocb.write(ocb_asm)
       end
-      branch_stub.compile = proc do |branch_asm|
+      branch_stub.compile = compile_jit_direct_jump(branch_stub, comment:)
+      branch_stub.compile.call(asm)
+    end
+
+    def compile_jit_direct_jump(branch_stub, comment:) # Proc escapes arguments in memory
+      proc do |branch_asm|
         branch_asm.comment(comment)
         branch_asm.stub(branch_stub) do
           case branch_stub.shape
@@ -5879,7 +5904,6 @@ module RubyVM::RJIT
           end
         end
       end
-      branch_stub.compile.call(asm)
     end
 
     # @param jit [RubyVM::RJIT::JITState]