YJIT: Specialize splatkw on T_HASH (#9764)

* YJIT: Specialize splatkw on T_HASH

* Fix a typo

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

* Fix a few more comments

---------

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

bootstraptest/test_yjit.rb

@@ -2535,6 +2535,18 @@ assert_equal '[1, 2]', %q{
 
   entry { 2 }
 }
+assert_equal '[1, 2]', %q{
+  def foo(a:) = [a, yield]
+
+  def entry(obj, &block)
+    foo(**obj, &block)
+  end
+
+  entry({ a: 3 }) { 2 }
+  obj = Object.new
+  def obj.to_hash = { a: 1 }
+  entry(obj) { 2 }
+}
 
 assert_equal '[1, 1, 2, 1, 2, 3]', %q{
   def expandarray

yjit.rb

@@ -287,6 +287,7 @@ module RubyVM::YJIT
         opt_plus
         opt_succ
         setlocal
+        splatkw
       ].each do |insn|
         print_counters(stats, out: out, prefix: "#{insn}_", prompt: "#{insn} exit reasons:", optional: true)
       end

yjit/src/codegen.rs

@@ -1404,7 +1404,7 @@ fn gen_duphash(
     // call rb_hash_resurrect(VALUE hash);
     let hash = asm.ccall(rb_hash_resurrect as *const u8, vec![hash.into()]);
 
-    let stack_ret = asm.stack_push(Type::Hash);
+    let stack_ret = asm.stack_push(Type::THash);
     asm.mov(stack_ret, hash);
 
     Some(KeepCompiling)
@@ -1440,24 +1440,39 @@ fn gen_splatarray(
 fn gen_splatkw(
     jit: &mut JITState,
     asm: &mut Assembler,
-    _ocb: &mut OutlinedCb,
+    ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
-    // Save the PC and SP because the callee may allocate
+    // Defer compilation so we can specialize on a runtime hash operand
-    jit_prepare_routine_call(jit, asm);
+    if !jit.at_current_insn() {
+        defer_compilation(jit, asm, ocb);
+        return Some(EndBlock);
+    }
 
-    // Get the operands from the stack
+    let comptime_hash = jit.peek_at_stack(&asm.ctx, 1);
-    let block_opnd = asm.stack_opnd(0);
+    if comptime_hash.hash_p() {
-    let block_type = asm.ctx.get_opnd_type(block_opnd.into());
+        // If a compile-time hash operand is T_HASH, just guard that it's T_HASH.
-    let hash_opnd = asm.stack_opnd(1);
+        let hash_opnd = asm.stack_opnd(1);
+        guard_object_is_hash(asm, hash_opnd, hash_opnd.into(), Counter::splatkw_not_hash);
+    } else {
+        // Otherwise, call #to_hash operand to get T_HASH.
 
-    let hash = asm.ccall(rb_to_hash_type as *const u8, vec![hash_opnd]);
+        // Save the PC and SP because the callee may allocate
-    asm.stack_pop(2); // Keep it on stack during ccall for GC
+        jit_prepare_routine_call(jit, asm);
 
-    let stack_ret = asm.stack_push(Type::Hash);
+        // Get the operands from the stack
-    asm.mov(stack_ret, hash);
+        let block_opnd = asm.stack_opnd(0);
-    asm.stack_push(block_type);
+        let block_type = asm.ctx.get_opnd_type(block_opnd.into());
-    // Leave block_opnd spilled by ccall as is
+        let hash_opnd = asm.stack_opnd(1);
-    asm.ctx.dealloc_temp_reg(asm.ctx.get_stack_size() - 1);
+
+        let hash = asm.ccall(rb_to_hash_type as *const u8, vec![hash_opnd]);
+        asm.stack_pop(2); // Keep it on stack during ccall for GC
+
+        let stack_ret = asm.stack_push(Type::THash);
+        asm.mov(stack_ret, hash);
+        asm.stack_push(block_type);
+        // Leave block_opnd spilled by ccall as is
+        asm.ctx.dealloc_temp_reg(asm.ctx.get_stack_size() - 1);
+    }
 
     Some(KeepCompiling)
 }
@@ -1620,6 +1635,38 @@ fn guard_object_is_array(
     }
 }
 
+fn guard_object_is_hash(
+    asm: &mut Assembler,
+    object: Opnd,
+    object_opnd: YARVOpnd,
+    counter: Counter,
+) {
+    let object_type = asm.ctx.get_opnd_type(object_opnd);
+    if object_type.is_hash() {
+        return;
+    }
+
+    let object_reg = match object {
+        Opnd::InsnOut { .. } => object,
+        _ => asm.load(object),
+    };
+    guard_object_is_heap(asm, object_reg, object_opnd, counter);
+
+    asm_comment!(asm, "guard object is hash");
+
+    // Pull out the type mask
+    let flags_opnd = Opnd::mem(VALUE_BITS, object_reg, RUBY_OFFSET_RBASIC_FLAGS);
+    let flags_opnd = asm.and(flags_opnd, (RUBY_T_MASK as u64).into());
+
+    // Compare the result with T_HASH
+    asm.cmp(flags_opnd, (RUBY_T_HASH as u64).into());
+    asm.jne(Target::side_exit(counter));
+
+    if Type::UnknownHeap.diff(object_type) != TypeDiff::Incompatible {
+        asm.ctx.upgrade_opnd_type(object_opnd, Type::THash);
+    }
+}
+
 fn guard_object_is_string(
     asm: &mut Assembler,
     object: Opnd,
@@ -2096,12 +2143,12 @@ fn gen_newhash(
         asm.cpop_into(new_hash); // x86 alignment
 
         asm.stack_pop(num.try_into().unwrap());
-        let stack_ret = asm.stack_push(Type::Hash);
+        let stack_ret = asm.stack_push(Type::THash);
         asm.mov(stack_ret, new_hash);
     } else {
         // val = rb_hash_new();
         let new_hash = asm.ccall(rb_hash_new as *const u8, vec![]);
-        let stack_ret = asm.stack_push(Type::Hash);
+        let stack_ret = asm.stack_push(Type::THash);
         asm.mov(stack_ret, new_hash);
     }
 

yjit/src/core.rs

@@ -50,7 +50,6 @@ pub enum Type {
     False,
     Fixnum,
     Flonum,
-    Hash,
     ImmSymbol,
 
     #[allow(unused)]
@@ -59,6 +58,7 @@ pub enum Type {
     TString, // An object with the T_STRING flag set, possibly an rb_cString
     CString, // An un-subclassed string of type rb_cString (can have instance vars in some cases)
     TArray, // An object with the T_ARRAY flag set, possibly an rb_cArray
+    THash, // An object with the T_HASH flag set, possibly an rb_cHash
 
     TProc, // A proc object. Could be an instance of a subclass of ::rb_cProc
 
@@ -111,7 +111,7 @@ impl Type {
             }
             match val.builtin_type() {
                 RUBY_T_ARRAY => Type::TArray,
-                RUBY_T_HASH => Type::Hash,
+                RUBY_T_HASH => Type::THash,
                 RUBY_T_STRING => Type::TString,
                 #[cfg(not(test))]
                 RUBY_T_DATA if unsafe { rb_obj_is_proc(val).test() } => Type::TProc,
@@ -154,7 +154,7 @@ impl Type {
         match self {
             Type::UnknownHeap => true,
             Type::TArray => true,
-            Type::Hash => true,
+            Type::THash => true,
             Type::HeapSymbol => true,
             Type::TString => true,
             Type::CString => true,
@@ -169,6 +169,11 @@ impl Type {
         matches!(self, Type::TArray)
     }
 
+    /// Check if it's a T_HASH object
+    pub fn is_hash(&self) -> bool {
+        matches!(self, Type::THash)
+    }
+
     /// Check if it's a T_STRING object (both TString and CString are T_STRING)
     pub fn is_string(&self) -> bool {
         match self {
@@ -187,7 +192,7 @@ impl Type {
             Type::Fixnum => Some(RUBY_T_FIXNUM),
             Type::Flonum => Some(RUBY_T_FLOAT),
             Type::TArray => Some(RUBY_T_ARRAY),
-            Type::Hash => Some(RUBY_T_HASH),
+            Type::THash => Some(RUBY_T_HASH),
             Type::ImmSymbol | Type::HeapSymbol => Some(RUBY_T_SYMBOL),
             Type::TString | Type::CString => Some(RUBY_T_STRING),
             Type::TProc => Some(RUBY_T_DATA),

yjit/src/cruby.rs

@@ -379,6 +379,11 @@ impl VALUE {
         }
     }
 
+    /// Returns true if the value is T_HASH
+    pub fn hash_p(self) -> bool {
+        !self.special_const_p() && self.builtin_type() == RUBY_T_HASH
+    }
+
     /// Returns true or false depending on whether the value is nil
     pub fn nil_p(self) -> bool {
         self == Qnil

yjit/src/stats.rs

@@ -524,6 +524,8 @@ make_counters! {
 
     objtostring_not_string,
 
+    splatkw_not_hash,
+
     binding_allocations,
     binding_set,