3 files changed, 117 insertions, 9 deletions
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index f274a64ca6..286f3f39b4 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -7,7 +7,7 @@ use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX, SpecialObjectType};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
 use crate::options::get_option;
@@ -252,6 +252,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::NewArray { elements, state } => gen_new_array(jit, asm, elements, &function.frame_state(*state)),
         Insn::NewRange { low, high, flag, state } => gen_new_range(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
+        Insn::StringCopy { val, chilled } => gen_string_copy(asm, opnd!(val), *chilled),
         Insn::Param { idx } => unreachable!("block.insns should not have Insn::Param({idx})"),
         Insn::Snapshot { .. } => return Some(()), // we don't need to do anything for this instruction at the moment
         Insn::Jump(branch) => return gen_jump(jit, asm, branch),
@@ -279,6 +280,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GetGlobal { id, state: _ } => gen_getglobal(asm, *id),
         Insn::SetIvar { self_val, id, val, state: _ } => return gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
         Insn::SideExit { state } => return gen_side_exit(jit, asm, &function.frame_state(*state)),
+        Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
             return None;
@@ -347,6 +349,20 @@ fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, state: &FrameState) ->
     Some(())
 }
 
+/// Emit a special object lookup
+fn gen_putspecialobject(asm: &mut Assembler, value_type: SpecialObjectType) -> Opnd {
+    asm_comment!(asm, "call rb_vm_get_special_object");
+
+    // Get the EP of the current CFP and load it into a register
+    let ep_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP);
+    let ep_reg = asm.load(ep_opnd);
+
+    asm.ccall(
+        rb_vm_get_special_object as *const u8,
+        vec![ep_reg, Opnd::UImm(u64::from(value_type))],
+    )
+}
+
 /// Compile an interpreter entry block to be inserted into an ISEQ
 fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
     asm_comment!(asm, "ZJIT entry point: {}", iseq_get_location(iseq, 0));
@@ -596,6 +612,17 @@ fn gen_send_without_block_direct(
     Some(ret)
 }
 
+/// Compile a string resurrection
+fn gen_string_copy(asm: &mut Assembler, recv: Opnd, chilled: bool) -> Opnd {
+    asm_comment!(asm, "call rb_ec_str_resurrect");
+    // TODO: split rb_ec_str_resurrect into separate functions
+    let chilled = if chilled { Opnd::Imm(1) } else { Opnd::Imm(0) };
+    asm.ccall(
+        rb_ec_str_resurrect as *const u8,
+        vec![EC, recv, chilled],
+    )
+}
+
 /// Compile an array duplication instruction
 fn gen_array_dup(
     asm: &mut Assembler,
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index de1c86e8d6..e0334ed44d 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -133,6 +133,7 @@ unsafe extern "C" {
     pub fn rb_str_setbyte(str: VALUE, index: VALUE, value: VALUE) -> VALUE;
     pub fn rb_vm_splat_array(flag: VALUE, ary: VALUE) -> VALUE;
     pub fn rb_vm_concat_array(ary1: VALUE, ary2st: VALUE) -> VALUE;
+    pub fn rb_vm_get_special_object(reg_ep: *const VALUE, value_type: vm_special_object_type) -> VALUE;
     pub fn rb_vm_concat_to_array(ary1: VALUE, ary2st: VALUE) -> VALUE;
     pub fn rb_vm_defined(
         ec: EcPtr,
@@ -213,6 +214,7 @@ pub use rb_vm_ci_flag as vm_ci_flag;
 pub use rb_vm_ci_kwarg as vm_ci_kwarg;
 pub use rb_METHOD_ENTRY_VISI as METHOD_ENTRY_VISI;
 pub use rb_RCLASS_ORIGIN as RCLASS_ORIGIN;
+pub use rb_vm_get_special_object as vm_get_special_object;
 
 /// Helper so we can get a Rust string for insn_name()
 pub fn insn_name(opcode: usize) -> String {
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 45a9024ca9..c67f25451a 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -138,6 +138,40 @@ impl Invariant {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum SpecialObjectType {
+    VMCore = 1,
+    CBase = 2,
+    ConstBase = 3,
+}
+
+impl From<u32> for SpecialObjectType {
+    fn from(value: u32) -> Self {
+        match value {
+            VM_SPECIAL_OBJECT_VMCORE => SpecialObjectType::VMCore,
+            VM_SPECIAL_OBJECT_CBASE => SpecialObjectType::CBase,
+            VM_SPECIAL_OBJECT_CONST_BASE => SpecialObjectType::ConstBase,
+            _ => panic!("Invalid special object type: {}", value),
+        }
+    }
+}
+
+impl From<SpecialObjectType> for u64 {
+    fn from(special_type: SpecialObjectType) -> Self {
+        special_type as u64
+    }
+}
+
+impl std::fmt::Display for SpecialObjectType {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            SpecialObjectType::VMCore => write!(f, "VMCore"),
+            SpecialObjectType::CBase => write!(f, "CBase"),
+            SpecialObjectType::ConstBase => write!(f, "ConstBase"),
+        }
+    }
+}
+
 /// Print adaptor for [`Invariant`]. See [`PtrPrintMap`].
 pub struct InvariantPrinter<'a> {
     inner: Invariant,
@@ -362,9 +396,12 @@ pub enum Insn {
     /// SSA block parameter. Also used for function parameters in the function's entry block.
     Param { idx: usize },
 
-    StringCopy { val: InsnId },
+    StringCopy { val: InsnId, chilled: bool },
     StringIntern { val: InsnId },
 
+    /// Put special object (VMCORE, CBASE, etc.) based on value_type
+    PutSpecialObject { value_type: SpecialObjectType },
+
     /// Call `to_a` on `val` if the method is defined, or make a new array `[val]` otherwise.
     ToArray { val: InsnId, state: InsnId },
     /// Call `to_a` on `val` if the method is defined, or make a new array `[val]` otherwise. If we
@@ -565,7 +602,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ArraySet { array, idx, val } => { write!(f, "ArraySet {array}, {idx}, {val}") }
             Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
             Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
-            Insn::StringCopy { val } => { write!(f, "StringCopy {val}") }
+            Insn::StringCopy { val, .. } => { write!(f, "StringCopy {val}") }
             Insn::Test { val } => { write!(f, "Test {val}") }
             Insn::IsNil { val } => { write!(f, "IsNil {val}") }
             Insn::Jump(target) => { write!(f, "Jump {target}") }
@@ -645,6 +682,9 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ArrayExtend { left, right, .. } => write!(f, "ArrayExtend {left}, {right}"),
             Insn::ArrayPush { array, val, .. } => write!(f, "ArrayPush {array}, {val}"),
             Insn::SideExit { .. } => write!(f, "SideExit"),
+            Insn::PutSpecialObject { value_type } => {
+                write!(f, "PutSpecialObject {}", value_type)
+            }
             insn => { write!(f, "{insn:?}") }
         }
     }
@@ -938,7 +978,7 @@ impl Function {
                     }
                 },
             Return { val } => Return { val: find!(*val) },
-            StringCopy { val } => StringCopy { val: find!(*val) },
+            StringCopy { val, chilled } => StringCopy { val: find!(*val), chilled: *chilled },
             StringIntern { val } => StringIntern { val: find!(*val) },
             Test { val } => Test { val: find!(*val) },
             &IsNil { val } => IsNil { val: find!(val) },
@@ -958,6 +998,7 @@ impl Function {
             FixnumGe { left, right } => FixnumGe { left: find!(*left), right: find!(*right) },
             FixnumLt { left, right } => FixnumLt { left: find!(*left), right: find!(*right) },
             FixnumLe { left, right } => FixnumLe { left: find!(*left), right: find!(*right) },
+            PutSpecialObject { value_type } => PutSpecialObject { value_type: *value_type },
             SendWithoutBlock { self_val, call_info, cd, args, state } => SendWithoutBlock {
                 self_val: find!(*self_val),
                 call_info: call_info.clone(),
@@ -1074,6 +1115,7 @@ impl Function {
             Insn::FixnumLe   { .. } => types::BoolExact,
             Insn::FixnumGt   { .. } => types::BoolExact,
             Insn::FixnumGe   { .. } => types::BoolExact,
+            Insn::PutSpecialObject { .. } => types::BasicObject,
             Insn::SendWithoutBlock { .. } => types::BasicObject,
             Insn::SendWithoutBlockDirect { .. } => types::BasicObject,
             Insn::Send { .. } => types::BasicObject,
@@ -1561,7 +1603,8 @@ impl Function {
             necessary[insn_id.0] = true;
             match self.find(insn_id) {
                 Insn::Const { .. } | Insn::Param { .. }
-                | Insn::PatchPoint(..) | Insn::GetConstantPath { .. } =>
+                | Insn::PatchPoint(..) | Insn::GetConstantPath { .. }
+                | Insn::PutSpecialObject { .. } =>
                     {}
                 Insn::ArrayMax { elements, state }
                 | Insn::NewArray { elements, state } => {
@@ -1580,7 +1623,7 @@ impl Function {
                     worklist.push_back(high);
                     worklist.push_back(state);
                 }
-                Insn::StringCopy { val }
+                Insn::StringCopy { val, .. }
                 | Insn::StringIntern { val }
                 | Insn::Return { val }
                 | Insn::Defined { v: val, .. }
@@ -2095,10 +2138,18 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 YARVINSN_nop => {},
                 YARVINSN_putnil => { state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(Qnil) })); },
                 YARVINSN_putobject => { state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) })); },
-                YARVINSN_putstring | YARVINSN_putchilledstring => {
-                    // TODO(max): Do something different for chilled string
+                YARVINSN_putspecialobject => {
+                    let value_type = SpecialObjectType::from(get_arg(pc, 0).as_u32());
+                    state.stack_push(fun.push_insn(block, Insn::PutSpecialObject { value_type }));
+                }
+                YARVINSN_putstring => {
                     let val = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
-                    let insn_id = fun.push_insn(block, Insn::StringCopy { val });
+                    let insn_id = fun.push_insn(block, Insn::StringCopy { val, chilled: false });
+                    state.stack_push(insn_id);
+                }
+                YARVINSN_putchilledstring => {
+                    let val = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
+                    let insn_id = fun.push_insn(block, Insn::StringCopy { val, chilled: true });
                     state.stack_push(insn_id);
                 }
                 YARVINSN_putself => { state.stack_push(self_param); }
@@ -3523,6 +3574,13 @@ mod tests {
         assert_method_hir("test",  expect![[r#"
             fn test:
             bb0(v0:BasicObject, v1:BasicObject):
+              v3:BasicObject = PutSpecialObject VMCore
+              v5:HashExact = NewHash
+              v7:BasicObject = SendWithoutBlock v3, :core#hash_merge_kwd, v5, v1
+              v8:BasicObject = PutSpecialObject VMCore
+              v9:StaticSymbol[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v10:Fixnum[1] = Const Value(1)
+              v12:BasicObject = SendWithoutBlock v8, :core#hash_merge_ptr, v7, v9, v10
               SideExit
         "#]]);
     }
@@ -3957,6 +4015,27 @@ mod tests {
     }
 
     #[test]
+    // Tests for ConstBase requires either constant or class definition, both
+    // of which can't be performed inside a method.
+    fn test_putspecialobject_vm_core_and_cbase() {
+        eval("
+            def test
+              alias aliased __callee__
+            end
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_putspecialobject, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:BasicObject = PutSpecialObject VMCore
+              v3:BasicObject = PutSpecialObject CBase
+              v4:StaticSymbol[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v5:StaticSymbol[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              v7:BasicObject = SendWithoutBlock v2, :core#set_method_alias, v3, v4, v5
+              Return v7
+        "#]]);
+    }
+
+    #[test]
     fn test_branchnil() {
         eval("
         def test(x) = x&.itself