Added debug statements in load_store and added a test

Author: jacob-shops

zjit/src/hir.rs

@@ -4604,12 +4604,19 @@ impl Function {
         self.infer_types();
     }
 
+    // struct Heap {
+    //     fn load(obj, offset)
+    //     fn kill(obj, offset)
+    //     fn store(obj, offset, val)
+    // }
+
     fn optimize_load_store(&mut self) {
         // TODO: Add specific tests for load_store
         // TODO: Add dead store elimination
         // 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);
@@ -4618,18 +4625,23 @@ impl Function {
                 let replacement_insn: InsnId = match self.find(insn_id) {
                     Insn::StoreField { recv, offset, val, .. } => {
                         let key = (recv, offset);
-
+                        eprintln!("hi jane");
                         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
-                            compile_time_heap.retain(|(_, off), _| *off == offset);
+                            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
-                        compile_time_heap.retain(|(_, off), _| *off == offset);
+                        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, .. } => {
@@ -4650,6 +4662,7 @@ 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
@@ -5477,6 +5490,8 @@ impl Function {
                               || ident_equal!($name, optimize_getivar)
                               || ident_equal!($name, optimize_c_calls) {
                     Counter::compile_hir_strength_reduce_time_ns
+                } else if ident_equal!($name, optimize_load_store) {
+                    Counter::compile_hir_optimize_load_store_time_ns
                 } else if ident_equal!($name, fold_constants) {
                     Counter::compile_hir_fold_constants_time_ns
                 } else if ident_equal!($name, clean_cfg) {

zjit/src/hir/opt_tests.rs

@@ -13696,4 +13696,51 @@ mod hir_opt_tests {
           Jump bb8(v67, v94)
         ");
     }
+
+    #[test]
+    fn test_double_store_removal() {
+        eval("
+            class C
+              def initialize
+                a = 1
+                @a = a
+                @a = a
+              end
+            end
+
+            C.new
+        ");
+        assert_snapshot!(hir_string_proc("C.instance_method(:initialize)"), @r"
+        fn initialize@<compiled>:4:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          Jump bb3(v1, v2)
+        bb2():
+          EntryPoint JIT(0)
+          v5:BasicObject = LoadArg :self@0
+          v6:NilClass = Const Value(nil)
+          Jump bb3(v5, v6)
+        bb3(v8:BasicObject, v9:NilClass):
+          v13:Fixnum[1] = Const Value(1)
+          PatchPoint SingleRactorMode
+          v35:HeapBasicObject = GuardType v8, HeapBasicObject
+          v36:CShape = LoadField v35, :_shape_id@0x1000
+          v37:CShape[0x1001] = GuardBitEquals v36, CShape(0x1001)
+          StoreField v35, :@a@0x1002, v13
+          WriteBarrier v35, v13
+          v40:CShape[0x1003] = Const CShape(0x1003)
+          StoreField v35, :_shape_id@0x1000, v40
+          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
+        ");
+    }
 }

zjit/src/stats.rs

@@ -167,6 +167,7 @@ make_counters! {
         compile_hir_time_ns,
         compile_hir_build_time_ns,
         compile_hir_strength_reduce_time_ns,
+        compile_hir_optimize_load_store_time_ns,
         compile_hir_fold_constants_time_ns,
         compile_hir_clean_cfg_time_ns,
         compile_hir_remove_redundant_patch_points_time_ns,