ZJIT: Only optimize [] and []= for exact Hash, not Hash subclasses

st0012 · st0012 · commit 4f900e3ce9ce · 2025-12-12T18:26:31.000Z
diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs
@@ -350,18 +350,30 @@ fn inline_array_pop(fun: &mut hir::Function, block: hir::BlockId, recv: hir::Ins
 }
 
 fn inline_hash_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
-    if let &[key] = args  {
+    let &[key] = args else { return None; };
+
+    // Only optimize exact Hash, not subclasses
+    if fun.likely_a(recv, types::HashExact, state) {
+        let recv = fun.coerce_to(block, recv, types::HashExact, state);
         let result = fun.push_insn(block, hir::Insn::HashAref { hash: recv, key, state });
-        return Some(result);
+        Some(result)
+    } else {
+        None
     }
-    None
 }
 
 fn inline_hash_aset(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
     let &[key, val] = args else { return None; };
-    let _ = fun.push_insn(block, hir::Insn::HashAset { hash: recv, key, val, state });
-    // Hash#[]= returns the value, not the hash
-    Some(val)
+
+    // Only optimize exact Hash, not subclasses
+    if fun.likely_a(recv, types::HashExact, state) {
+        let recv = fun.coerce_to(block, recv, types::HashExact, state);
+        let _ = fun.push_insn(block, hir::Insn::HashAset { hash: recv, key, val, state });
+        // Hash#[]= returns the value, not the hash
+        Some(val)
+    } else {
+        None
+    }
 }
 
 fn inline_string_bytesize(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -4710,7 +4710,7 @@ impl Function {
             }
             // Instructions with Hash operands
             Insn::HashAref { hash, .. }
-            | Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),
+            | Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::HashExact),
             Insn::HashDup { val, .. } => self.assert_subtype(insn_id, val, types::HashExact),
             // Other
             Insn::ObjectAllocClass { class, .. } => {
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -6678,7 +6678,7 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn test_hash_aref_subclass() {
+    fn test_no_optimize_hash_aref_subclass() {
         eval("
             class C < Hash; end
             def test(hash, key)
@@ -6701,8 +6701,7 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(C@0x1000, []@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(C@0x1000)
           v27:HashSubclass[class_exact:C] = GuardType v11, HashSubclass[class_exact:C]
-          v28:BasicObject = HashAref v27, v12
-          IncrCounter inline_cfunc_optimized_send_count
+          v28:BasicObject = CCallWithFrame v27, :Hash#[]@0x1038, v12
           CheckInterrupts
           Return v28
         ");
@@ -6803,7 +6802,7 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn test_hash_aset_subclass() {
+    fn test_no_optimize_hash_aset_subclass() {
         eval("
             class C < Hash; end
             def test(hash, key, val)
@@ -6827,8 +6826,7 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(C@0x1000, []=@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(C@0x1000)
           v35:HashSubclass[class_exact:C] = GuardType v13, HashSubclass[class_exact:C]
-          HashAset v35, v14, v15
-          IncrCounter inline_cfunc_optimized_send_count
+          v36:BasicObject = CCallWithFrame v35, :Hash#[]=@0x1038, v14, v15
           CheckInterrupts
           Return v15
         ");

Original file line number	Diff line number	Diff line change
`@@ -4710,7 +4710,7 @@ impl Function {`
`4710`	`4710`	`}`
`4711`	`4711`	`// Instructions with Hash operands`
`4712`	`4712`	`Insn::HashAref { hash, .. }`
`4713`		`- \| Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),`
	`4713`	`+ \| Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::HashExact),`
`4714`	`4714`	`Insn::HashDup { val, .. } => self.assert_subtype(insn_id, val, types::HashExact),`
`4715`	`4715`	`// Other`
`4716`	`4716`	`Insn::ObjectAllocClass { class, .. } => {`