Get load_store optimization working

jacob-shops · jacob-shops · commit 875d6000a3cc · 2026-02-24T19:43:38.000-05:00
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -1172,7 +1172,7 @@ impl Insn {
             Insn::LoadSelf { .. } => Effect::read_write(abstract_heaps::Frame, abstract_heaps::Empty),
             Insn::LoadField { .. } => Effect::read_write(abstract_heaps::Memory, abstract_heaps::Empty),
             Insn::StoreField { .. } => effects::Any,
-            Insn::WriteBarrier { .. } => effects::Any,
+            Insn::WriteBarrier { .. } => effects::Allocator,
             Insn::GetLocal   { .. } => Effect::read_write(abstract_heaps::Locals, abstract_heaps::Empty),
             Insn::SetLocal { .. } => effects::Any,
             Insn::GetSpecialSymbol { .. } => effects::Any,
@@ -1262,6 +1262,9 @@ impl Insn {
     /// Note: These are restrictions on the `write` `EffectSet` only. Even instructions with
     /// `read: effects::Any` could potentially be omitted.
     fn is_elidable(&self) -> bool {
+        if let Insn::WriteBarrier { .. } = self {
+            return false
+        }
         abstract_heaps::Allocator.includes(self.effects_of().write_bits())
     }
 }
@@ -4486,41 +4489,35 @@ impl Function {
         // The key for the hashmap should be type and offset, with a value of value
         // This lets us to index in with both load and store fields since insn_ids are probably always going to be different and we can't easily match on that
         // So... how do we match against and store the enum label without all the data?? not sure yet :/
-        eprintln!("{}", FunctionPrinter::with_snapshot(self));
         let mut compile_time_heap: HashMap<(InsnId, i32), InsnId>  = HashMap::new();
         for block in self.rpo() {
             let old_insns = std::mem::take(&mut self.blocks[block.0].insns);
             let mut new_insns = vec![];
             for insn_id in old_insns {
                 let replacement_insn: InsnId = match self.find(insn_id) {
                     Insn::StoreField { recv, offset, val, .. } => {
-                        let key = (recv, offset);
-                        eprintln!("hi jane");
+                        let key = (self.chase_insn(recv), offset);
                         let heap_entry = compile_time_heap.get(&key).copied();
-                        eprintln!("{heap_entry:?}");
                         // TODO(Jacob): Switch from actual to partial equality
                         if Some(val) == heap_entry {
-                            eprintln!("Matched value {val}, {heap_entry:?}");
                             // TODO(Jacob): Add TBAA to avoid removing so many entries
-                            eprintln!("Erasing aliasing offsets {offset}");
                             compile_time_heap.retain(|(_, off), _| *off != offset);
                             // If the value is already stored, short circuit and don't add an instruction to the block
                             continue
                         }
                         // TODO(Jacob): Add TBAA to avoid removing so many entries
-                        eprintln!("Erasing aliasing offsets {offset}");
                         compile_time_heap.retain(|(_, off), _| *off != offset);
                         compile_time_heap.insert(key, val);
-                        eprintln!("Inserted into heap {key:?}, {val}");
                         insn_id
                     },
                     Insn::LoadField { recv, offset, .. } => {
-                        let key = (recv, offset);
+                        let key = (self.chase_insn(recv), offset);
                         match compile_time_heap.entry(key) {
                             std::collections::hash_map::Entry::Occupied(entry) => {
                                 // If the value is already saved, we can't short circuit like we can with a store.
                                 // However, we can avoid the load with a reference to the representative to the union from SSA.
                                 self.make_equal_to(insn_id, *entry.get());
+                                continue
                             }
                             std::collections::hash_map::Entry::Vacant(_) => {
                                 // TODO(Jacob): Make sure this is correct, could be wrong?
@@ -4532,7 +4529,6 @@ impl Function {
                     insn => {
                         // If an instruction affects memory and we haven't modeled it, the compile_time_heap is invalidated
                         if insn.effects_of().includes(Effect::write(abstract_heaps::Memory)) {
-                            eprintln!("Clearing... {insn:?}");
                             compile_time_heap.clear();
                         }
                         insn_id
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -5007,8 +5007,6 @@ mod hir_opt_tests {
           v14:HeapBasicObject = RefineType v6, HeapBasicObject
           v17:Fixnum[2] = Const Value(2)
           PatchPoint SingleRactorMode
-          v36:CShape = LoadField v14, :_shape_id@0x1000
-          v37:CShape[0x1003] = GuardBitEquals v36, CShape(0x1003)
           StoreField v14, :@bar@0x1004, v17
           WriteBarrier v14, v17
           v40:CShape[0x1005] = Const CShape(0x1005)
@@ -8390,8 +8388,7 @@ mod hir_opt_tests {
           v32:ArrayExact = GuardType v9, ArrayExact
           v33:CUInt64 = LoadField v32, :_rbasic_flags@0x1038
           v34:CUInt64 = GuardNoBitsSet v33, RUBY_FL_FREEZE=CUInt64(2048)
-          v35:CUInt64 = LoadField v32, :_rbasic_flags@0x1038
-          v36:CUInt64 = GuardNoBitsSet v35, RUBY_ELTS_SHARED=CUInt64(4096)
+          v36:CUInt64 = GuardNoBitsSet v33, RUBY_ELTS_SHARED=CUInt64(4096)
           v37:CInt64[1] = UnboxFixnum v16
           v38:CInt64 = ArrayLength v32
           v39:CInt64[1] = GuardLess v37, v38
@@ -8436,8 +8433,7 @@ mod hir_opt_tests {
           v37:Fixnum = GuardType v14, Fixnum
           v38:CUInt64 = LoadField v36, :_rbasic_flags@0x1038
           v39:CUInt64 = GuardNoBitsSet v38, RUBY_FL_FREEZE=CUInt64(2048)
-          v40:CUInt64 = LoadField v36, :_rbasic_flags@0x1038
-          v41:CUInt64 = GuardNoBitsSet v40, RUBY_ELTS_SHARED=CUInt64(4096)
+          v41:CUInt64 = GuardNoBitsSet v38, RUBY_ELTS_SHARED=CUInt64(4096)
           v42:CInt64 = UnboxFixnum v37
           v43:CInt64 = ArrayLength v36
           v44:CInt64 = GuardLess v42, v43
@@ -13032,17 +13028,13 @@ mod hir_opt_tests {
           v14:HeapBasicObject = RefineType v6, HeapBasicObject
           v17:Fixnum[2] = Const Value(2)
           PatchPoint SingleRactorMode
-          v50:CShape = LoadField v14, :_shape_id@0x1000
-          v51:CShape[0x1003] = GuardBitEquals v50, CShape(0x1003)
           StoreField v14, :@b@0x1004, v17
           WriteBarrier v14, v17
           v54:CShape[0x1005] = Const CShape(0x1005)
           StoreField v14, :_shape_id@0x1000, v54
           v21:HeapBasicObject = RefineType v14, HeapBasicObject
           v24:Fixnum[3] = Const Value(3)
           PatchPoint SingleRactorMode
-          v57:CShape = LoadField v21, :_shape_id@0x1000
-          v58:CShape[0x1005] = GuardBitEquals v57, CShape(0x1005)
           StoreField v21, :@c@0x1006, v24
           WriteBarrier v21, v24
           v61:CShape[0x1007] = Const CShape(0x1007)
@@ -13247,9 +13239,6 @@ mod hir_opt_tests {
           v20:HeapBasicObject = RefineType v8, HeapBasicObject
           PatchPoint NoEPEscape(initialize)
           PatchPoint SingleRactorMode
-          v43:CShape = LoadField v20, :_shape_id@0x1000
-          v44:CShape[0x1003] = GuardBitEquals v43, CShape(0x1003)
-          StoreField v20, :@a@0x1002, v13
           WriteBarrier v20, v13
           CheckInterrupts
           Return v13
