From d4bbee7475269972cf56ebd3b5e4dc9905789663 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Tue, 12 Dec 2023 17:10:14 -0500
Subject: [PATCH] YJIT: Fix off-by-one in Kernel#send type handling (#9212)

Previously, if the method ID argument happens to be on one below the top
of the stack, we didn't overwrite the type of the stack slot, which
leaves an incorrect type for the stack slot. The included script tripped
asserts both with and without --yjit-verify-ctx.
---
 bootstraptest/test_yjit.rb | 13 +++++++++++++
 yjit/src/core.rs           | 29 +++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 43ed9f2803..b9c332b37a 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,3 +1,16 @@
+# regression test for send stack shifting
+assert_normal_exit %q{
+  def foo(a, b)
+    a.singleton_methods(b)
+  end
+
+  def call_foo
+    [1, 1, 1, 1, 1, 1, send(:foo, 1, 1)]
+  end
+
+  call_foo
+}
+
 # regression test for arity check with splat
 assert_equal '[:ae, :ae]', %q{
   def req_one(a_, b_ = 1) = raise
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index f84286bd2f..064a7b5e8f 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2180,8 +2180,13 @@ impl Assembler {
         let method_name_index = (self.ctx.stack_size as usize) - argc - 1;
 
         for i in method_name_index..(self.ctx.stack_size - 1) as usize {
-            if i + 1 < MAX_TEMP_TYPES {
-                self.ctx.set_temp_mapping(i, self.ctx.get_temp_mapping(i + 1));
+            if i < MAX_TEMP_TYPES {
+                let next_arg_mapping = if i + 1 < MAX_TEMP_TYPES {
+                    self.ctx.get_temp_mapping(i + 1)
+                } else {
+                    TempMapping::map_to_stack(Type::Unknown)
+                };
+                self.ctx.set_temp_mapping(i, next_arg_mapping);
             }
         }
         self.stack_pop(1);
@@ -3500,6 +3505,26 @@ mod tests {
         // TODO: write more tests for Context type diff
     }
 
+    #[test]
+    fn shift_stack_for_send() {
+        let mut asm = Assembler::new();
+
+        // Push values to simulate send(:name, arg) with 6 items already on-stack
+        for _ in 0..6 {
+            asm.stack_push(Type::Fixnum);
+        }
+        asm.stack_push(Type::Unknown);
+        asm.stack_push(Type::ImmSymbol);
+        asm.stack_push(Type::Unknown);
+
+        // This method takes argc of the sendee, not argc of send
+        asm.shift_stack(1);
+
+        // The symbol should be gone
+        assert_eq!(Type::Unknown, asm.ctx.get_opnd_type(StackOpnd(0)));
+        assert_eq!(Type::Unknown, asm.ctx.get_opnd_type(StackOpnd(1)));
+    }
+
     #[test]
     fn test_miri_ref_unchecked() {
         let blockid = BlockId {