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

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 43ed9f2803036b..b9c332b37aed1a 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 f84286bd2fb681..064a7b5e8f772c 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 {