From 1936278461cf9aec1495596bf25a2963721f21ee Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Thu, 8 Feb 2024 09:03:47 -0800
Subject: [PATCH] YJIT: Maintain MapToLocal that is just upgraded (#9876)

---
 yjit/src/core.rs | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index e46de01f2f..b34a455b2a 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -1947,6 +1947,9 @@ impl Context {
                         let mut new_type = self.get_local_type(idx);
                         new_type.upgrade(opnd_type);
                         self.set_local_type(idx, new_type);
+                        // Re-attach MapToLocal for this StackOpnd(idx). set_local_type() detaches
+                        // all MapToLocal mappings, including the one we're upgrading here.
+                        self.set_opnd_mapping(opnd, mapping);
                     }
                 }
             }
@@ -3581,6 +3584,14 @@ mod tests {
         // TODO: write more tests for Context type diff
     }
 
+    #[test]
+    fn context_upgrade_local() {
+        let mut asm = Assembler::new();
+        asm.stack_push_local(0);
+        asm.ctx.upgrade_opnd_type(StackOpnd(0), Type::Nil);
+        assert_eq!(Type::Nil, asm.ctx.get_opnd_type(StackOpnd(0)));
+    }
+
     #[test]
     fn context_chain_depth() {
         let mut ctx = Context::default();