ZJIT: Inline Hash#key? and its aliases

Author: st0012

test/ruby/test_zjit.rb

@@ -1397,6 +1397,46 @@ def test(arr) = arr.pop(2)
     }, call_threshold: 2
   end
 
+  def test_opt_hash_key_p
+    assert_compiles 'true', %q{
+      def test(h, k) = h.key?(k)
+      test({a: 1, b: 2}, :a)
+      test({a: 1, b: 2}, :a)
+    }, call_threshold: 2
+  end
+
+  def test_opt_hash_key_p_false
+    assert_compiles 'false', %q{
+      def test(h, k) = h.key?(k)
+      test({a: 1, b: 2}, :c)
+      test({a: 1, b: 2}, :c)
+    }, call_threshold: 2
+  end
+
+  def test_opt_hash_has_key
+    assert_compiles 'true', %q{
+      def test(h, k) = h.has_key?(k)
+      test({a: 1}, :a)
+      test({a: 1}, :a)
+    }, call_threshold: 2
+  end
+
+  def test_opt_hash_include
+    assert_compiles 'true', %q{
+      def test(h, k) = h.include?(k)
+      test({a: 1}, :a)
+      test({a: 1}, :a)
+    }, call_threshold: 2
+  end
+
+  def test_opt_hash_member
+    assert_compiles 'true', %q{
+      def test(h, k) = h.member?(k)
+      test({a: 1}, :a)
+      test({a: 1}, :a)
+    }, call_threshold: 2
+  end
+
   def test_new_range_inclusive
     assert_compiles '1..5', %q{
       def test(a, b) = a..b

zjit/bindgen/src/main.rs

@@ -118,6 +118,7 @@ fn main() {
         .allowlist_type("st_retval")
         .allowlist_function("rb_hash_aset")
         .allowlist_function("rb_hash_aref")
+        .allowlist_function("rb_hash_has_key")
         .allowlist_function("rb_hash_bulk_insert")
         .allowlist_function("rb_hash_stlike_lookup")
         .allowlist_function("rb_ary_new_capa")

zjit/src/codegen.rs

@@ -466,6 +466,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::CheckInterrupts { state } => no_output!(gen_check_interrupts(jit, asm, &function.frame_state(state))),
         &Insn::HashDup { val, state } => { gen_hash_dup(asm, opnd!(val), &function.frame_state(state)) },
         &Insn::HashAref { hash, key, state } => { gen_hash_aref(jit, asm, opnd!(hash), opnd!(key), &function.frame_state(state)) },
+        &Insn::HashKeyP { hash, key, state } => { gen_hash_key_p(jit, asm, opnd!(hash), opnd!(key), &function.frame_state(state)) },
         &Insn::ArrayPush { array, val, state } => { no_output!(gen_array_push(asm, opnd!(array), opnd!(val), &function.frame_state(state))) },
         &Insn::ToNewArray { val, state } => { gen_to_new_array(jit, asm, opnd!(val), &function.frame_state(state)) },
         &Insn::ToArray { val, state } => { gen_to_array(jit, asm, opnd!(val), &function.frame_state(state)) },
@@ -1048,6 +1049,11 @@ fn gen_hash_aref(jit: &mut JITState, asm: &mut Assembler, hash: Opnd, key: Opnd,
     asm_ccall!(asm, rb_hash_aref, hash, key)
 }
 
+fn gen_hash_key_p(jit: &mut JITState, asm: &mut Assembler, hash: Opnd, key: Opnd, state: &FrameState) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+    asm_ccall!(asm, rb_hash_has_key, hash, key)
+}
+
 fn gen_array_push(asm: &mut Assembler, array: Opnd, val: Opnd, state: &FrameState) {
     gen_prepare_leaf_call_with_gc(asm, state);
     asm_ccall!(asm, rb_ary_push, array, val);

zjit/src/cruby_bindings.inc.rs

@@ -232,10 +232,10 @@ pub fn init() -> Annotations {
     annotate!(rb_cHash, "size", types::Fixnum, no_gc, leaf, elidable);
     annotate!(rb_cHash, "empty?", types::BoolExact, no_gc, leaf, elidable);
     // These 4 are all aliases for rb_hash_has_key, which can call #hash and #eql? on the key
-    annotate!(rb_cHash, "key?", types::BoolExact);
-    annotate!(rb_cHash, "has_key?", types::BoolExact);
-    annotate!(rb_cHash, "include?", types::BoolExact);
-    annotate!(rb_cHash, "member?", types::BoolExact);
+    annotate!(rb_cHash, "key?", inline_hash_key_p, types::BoolExact);
+    annotate!(rb_cHash, "has_key?", inline_hash_key_p, types::BoolExact);
+    annotate!(rb_cHash, "include?", inline_hash_key_p, types::BoolExact);
+    annotate!(rb_cHash, "member?", inline_hash_key_p, types::BoolExact);
     annotate!(rb_cNilClass, "nil?", inline_nilclass_nil_p);
     annotate!(rb_mKernel, "nil?", inline_kernel_nil_p);
     annotate!(rb_mKernel, "respond_to?", inline_kernel_respond_to_p);
@@ -361,6 +361,13 @@ fn inline_hash_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::Ins
     None
 }
 
+fn inline_hash_key_p(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    if let &[key] = args {
+        let result = fun.push_insn(block, hir::Insn::HashKeyP { hash: recv, key, state });
+        return Some(result);
+    }
+    None
+}
 
 fn inline_string_bytesize(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
     if args.is_empty() && fun.likely_a(recv, types::String, state) {

zjit/src/cruby_methods.rs

@@ -685,6 +685,7 @@ pub enum Insn {
     ArrayLength { array: InsnId },
 
     HashAref { hash: InsnId, key: InsnId, state: InsnId },
+    HashKeyP { hash: InsnId, key: InsnId, state: InsnId },
     HashDup { val: InsnId, state: InsnId },
 
     /// Allocate an instance of the `val` object without calling `#initialize` on it.
@@ -1127,6 +1128,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
             Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
             Insn::HashAref { hash, key, .. } => { write!(f, "HashAref {hash}, {key}")}
+            Insn::HashKeyP { hash, key, .. } => { write!(f, "HashKeyP {hash}, {key}")}
             Insn::ObjectAlloc { val, .. } => { write!(f, "ObjectAlloc {val}") }
             &Insn::ObjectAllocClass { class, .. } => {
                 let class_name = get_class_name(class);
@@ -1960,6 +1962,7 @@ impl Function {
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
             &HashAref { hash, key, state } => HashAref { hash: find!(hash), key: find!(key), state },
+            &HashKeyP { hash, key, state } => HashKeyP { hash: find!(hash), key: find!(key), state },
             &ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
             &ObjectAllocClass { class, state } => ObjectAllocClass { class, state: find!(state) },
             &CCall { cfunc, recv, ref args, name, return_type, elidable } => CCall { cfunc, recv: find!(recv), args: find_vec!(args), name, return_type, elidable },
@@ -2096,6 +2099,7 @@ impl Function {
             Insn::ArrayPop { .. } => types::BasicObject,
             Insn::ArrayLength { .. } => types::CInt64,
             Insn::HashAref { .. } => types::BasicObject,
+            Insn::HashKeyP { .. } => types::BoolExact,
             Insn::NewHash { .. } => types::HashExact,
             Insn::HashDup { .. } => types::HashExact,
             Insn::NewRange { .. } => types::RangeExact,
@@ -3778,7 +3782,7 @@ impl Function {
             &Insn::ArrayLength { array } => {
                 worklist.push_back(array);
             }
-            &Insn::HashAref { hash, key, state } => {
+            &Insn::HashAref { hash, key, state } | &Insn::HashKeyP { hash, key, state } => {
                 worklist.push_back(hash);
                 worklist.push_back(key);
                 worklist.push_back(state);
@@ -4471,6 +4475,7 @@ impl Function {
             }
             // Instructions with Hash operands
             Insn::HashAref { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),
+            Insn::HashKeyP { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),
             Insn::HashDup { val, .. } => self.assert_subtype(insn_id, val, types::HashExact),
             // Other
             Insn::ObjectAllocClass { class, .. } => {

zjit/src/hir.rs

@@ -5139,7 +5139,8 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Hash@0x1000, key?@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Hash@0x1000)
           v26:HashExact = GuardType v11, HashExact
-          v27:BoolExact = CCallWithFrame v26, :Hash#key?@0x1038, v12
+          v27:BoolExact = HashKeyP v26, v12
+          IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
           Return v27
         ");