From 1bb7638e7a864db9b635c5d41f18653e8451dbc7 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Mon, 22 Apr 2024 11:16:46 -0400
Subject: [PATCH] YJIT: Fix shrinking block with assumption too much (#10585)

* YJIT: Fix shrinking block with assumption too much

Under the very specific circumstances, discovered by a test case in
`ruby/spec`, an `expandarray` block can contain just a branch and carry
a method lookup assumption. Previously, when we regenerated the branch,
we allowed it to shrink to empty, since we put the code at the jump
target immediately after it. That was incorrect and caused a crash while
the block is invalidated, since that left no room to patch in an exit.

When regenerating a branch that makes up a block entirely, and the block
could be invalidated, we need to ensure there is room for invalidation.
When there is code before the branch, they should act as padding, so we
don't need to worry about those cases.

* skip on RJIT
---
 bootstraptest/test_yjit.rb | 13 +++++++++++++
 yjit/src/core.rs           |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 6b7d483926..c7a5883986 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -4770,6 +4770,19 @@ assert_equal '[:ok, :ok, :ok]', %q{
   tests
 }
 
+# regression test for invalidating an empty block
+assert_equal '0', %q{
+  def foo = (* = 1).pred
+
+  foo # compile it
+
+  class Integer
+    def to_ary = [] # invalidate
+  end
+
+  foo # try again
+} unless rjit_enabled? # doesn't work on RJIT
+
 # test integer left shift with constant rhs
 assert_equal [0x80000000000, 'a+', :ok].inspect, %q{
   def shift(val) = val << 43
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index fb7d52cc5d..4152eab02c 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2643,6 +2643,12 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &Branch) {
         branch.get_target_address(1).map(|addr| Target::CodePtr(addr)),
     );
 
+    // If the entire block is the branch and the block could be invalidated,
+    // we need to pad to ensure there is room for invalidation patching.
+    if branch.start_addr == block.start_addr && branch_terminates_block && block.entry_exit.is_some() {
+        asm.pad_inval_patch();
+    }
+
     // Rewrite the branch
     let old_write_pos = cb.get_write_pos();
     let old_dropped_bytes = cb.has_dropped_bytes();