ZJIT: Constant-fold LoadField/GuardBitEquals from known-frozen object

Resolves https://github.com/Shopify/ruby/issues/915

When we have `LoadField` with a `Shape` return type, we can fold it similar to the object case.
`GuardBitEquals` can be removed when the argument is `Const` and the values are equal.

The behaviors for loading instances variables from frozen/dynamic objects are already covered in existing tests so no new tests were added.

2 files changed, 45 insertions, 39 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index a959ab07f4..8c7c838900 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -3798,6 +3798,30 @@ impl Function {
                             _ => insn_id,
                         }
                     }
+                    Insn::LoadField { recv, offset, return_type, .. } if return_type.is_subtype(types::CShape) &&
+                            u32::try_from(offset).is_ok() => {
+                        let offset = (offset as u32).to_usize();
+                        let recv_type = self.type_of(recv);
+                        match recv_type.ruby_object() {
+                            Some(recv_obj) if recv_obj.is_frozen() => {
+                                let recv_ptr = recv_obj.as_ptr() as *const u32;
+                                let val = unsafe { recv_ptr.byte_add(offset).read() };
+                                self.new_insn(Insn::Const { val: Const::CShape(ShapeId(val)) })
+                            }
+                            _ => insn_id,
+                        }
+                    }
+                    Insn::GuardBitEquals { val, expected, .. } => {
+                        match self.find(val) {
+                            // TODO: Refactor this into a more general method like
+                            // has_value(Const) that can check on the value specialization
+                            // of the Type instead
+                            Insn::Const { val: const_val } if const_val == expected => {
+                                continue;
+                            }
+                            _ => insn_id
+                        }
+                    }
                     Insn::AnyToString { str, .. } if self.is_a(str, types::String) => {
                         self.make_equal_to(insn_id, str);
                         // Don't bother re-inferring the type of str; we already know it.
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 1d360bed71..8d49353b30 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -9927,11 +9927,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestFrozen@0x1010, a@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestFrozen@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:Fixnum[1] = Const Value(1)
+          v29:Fixnum[1] = Const Value(1)
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -9969,11 +9967,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestMultiIvars@0x1010, b@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestMultiIvars@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:Fixnum[20] = Const Value(20)
+          v29:Fixnum[20] = Const Value(20)
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10009,11 +10005,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestFrozenStr@0x1010, name@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestFrozenStr@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:StringExact[VALUE(0x1050)] = Const Value(VALUE(0x1050))
+          v29:StringExact[VALUE(0x1048)] = Const Value(VALUE(0x1048))
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10049,11 +10043,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestFrozenNil@0x1010, value@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestFrozenNil@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:NilClass = Const Value(nil)
+          v29:NilClass = Const Value(nil)
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10129,11 +10121,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestAttrReader@0x1010, value@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestAttrReader@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:Fixnum[42] = Const Value(42)
+          v29:Fixnum[42] = Const Value(42)
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10169,11 +10159,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestFrozenSym@0x1010, sym@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestFrozenSym@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:StaticSymbol[:hello] = Const Value(VALUE(0x1050))
+          v29:StaticSymbol[:hello] = Const Value(VALUE(0x1048))
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10209,11 +10197,9 @@ mod hir_opt_tests {
           v20:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestFrozenBool@0x1010, flag@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestFrozenBool@0x1010)
-          v25:CShape = LoadField v20, :_shape_id@0x1048
-          v26:CShape[0x1049] = GuardBitEquals v25, CShape(0x1049)
-          v28:TrueClass = Const Value(true)
+          v29:TrueClass = Const Value(true)
           CheckInterrupts
-          Return v28
+          Return v29
         ");
     }
 
@@ -10288,22 +10274,18 @@ mod hir_opt_tests {
           v28:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(TestNestedAccess@0x1010, x@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(TestNestedAccess@0x1010)
-          v39:CShape = LoadField v28, :_shape_id@0x1048
-          v40:CShape[0x1049] = GuardBitEquals v39, CShape(0x1049)
-          v52:Fixnum[100] = Const Value(100)
+          v53:Fixnum[100] = Const Value(100)
           PatchPoint SingleRactorMode
-          PatchPoint StableConstantNames(0x1050, NESTED_FROZEN)
+          PatchPoint StableConstantNames(0x1048, NESTED_FROZEN)
           v34:HeapObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
-          PatchPoint MethodRedefined(TestNestedAccess@0x1010, y@0x1058, cme:0x1060)
+          PatchPoint MethodRedefined(TestNestedAccess@0x1010, y@0x1050, cme:0x1058)
           PatchPoint NoSingletonClass(TestNestedAccess@0x1010)
-          v43:CShape = LoadField v34, :_shape_id@0x1048
-          v44:CShape[0x1049] = GuardBitEquals v43, CShape(0x1049)
-          v53:Fixnum[200] = Const Value(200)
-          PatchPoint MethodRedefined(Integer@0x1088, +@0x1090, cme:0x1098)
-          v54:Fixnum[300] = Const Value(300)
+          v55:Fixnum[200] = Const Value(200)
+          PatchPoint MethodRedefined(Integer@0x1080, +@0x1088, cme:0x1090)
+          v56:Fixnum[300] = Const Value(300)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v54
+          Return v56
         ");
     }