ZJIT: Add support for putspecialobject

Author: st0012

test/ruby/test_zjit.rb

@@ -568,6 +568,17 @@ def test() = @foo = 1
     }
   end
 
+  def test_putspecialobject_method
+    assert_compiles 'nil', %q{
+      def test
+        alias bar __callee__
+      end
+
+      test
+      bar
+    }, insns: [:putspecialobject]
+  end
+
   # tool/ruby_vm/views/*.erb relies on the zjit instructions a) being contiguous and
   # b) being reliably ordered after all the other instructions.
   def test_instruction_order

vm_insnhelper.c

@@ -5550,6 +5550,13 @@ vm_get_special_object(const VALUE *const reg_ep,
     }
 }
 
+// ZJIT implementation is using the C function
+// and needs to call a non-static function
+VALUE rb_vm_get_special_object(const VALUE *reg_ep, enum vm_special_object_type type)
+{
+    return vm_get_special_object(reg_ep, type);
+}
+
 static VALUE
 vm_concat_array(VALUE ary1, VALUE ary2st)
 {

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;
@@ -276,6 +276,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(jit, asm, *cfun, args)?,
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
         Insn::SetIvar { self_val, id, val, state: _ } => gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
+        Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
             return None;
@@ -317,6 +318,20 @@ fn gen_setivar(asm: &mut Assembler, recv: Opnd, id: ID, val: Opnd) -> Opnd {
     )
 }
 
+/// 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(value_type as u32 as u64)],
+    )
+}
+
 /// 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));

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 {

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 {
+            1 => SpecialObjectType::VMCore,
+            2 => SpecialObjectType::CBase,
+            3 => SpecialObjectType::ConstBase,
+            _ => panic!("Invalid special object type: {}", value),
+        }
+    }
+}
+
+impl From<SpecialObjectType> for u32 {
+    fn from(special_type: SpecialObjectType) -> Self {
+        special_type as u32
+    }
+}
+
+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,
@@ -365,6 +399,9 @@ pub enum Insn {
     StringCopy { val: InsnId },
     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
@@ -482,6 +519,7 @@ impl Insn {
         match self {
             Insn::Const { .. } => false,
             Insn::Param { .. } => false,
+            Insn::PutSpecialObject { .. } => false,
             Insn::StringCopy { .. } => false,
             Insn::NewArray { .. } => false,
             Insn::NewHash { .. } => false,
@@ -630,6 +668,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:?}") }
         }
     }
@@ -943,6 +984,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(),
@@ -1056,6 +1098,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,
@@ -1542,7 +1585,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 } => {
@@ -2074,6 +2118,10 @@ 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_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 | YARVINSN_putchilledstring => {
                     // TODO(max): Do something different for chilled string
                     let val = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
@@ -3806,6 +3854,25 @@ mod tests {
         "#]]);
     }
 
+    #[test]
+    fn test_putspecialobject_callee() {
+        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("