ZJIT: Fix hanging loop (ruby#16711)

ruby#16122 (c272297) worked for maximal SSA but
does not work for "normal" SSA. This is because it used information
propagating across block args/params as a proxy for tracking changes in
dependent blocks. To do this properly we need to move to something like
SCCP.

See example HIR diff from the repro in codegen test:

```diff
diff --git a/before b/after
index da00a9e..b1d2a58 100644
--- a/before
+++ b/after
@@ -64,10 +64,10 @@ bb7(v48:BasicObject):
   StoreField v46, :_shape_id@0x4, v105
   v79:HeapBasicObject = RefineType v46, HeapBasicObject
   Jump bb5(v79, v41)
-bb5(v81:HeapBasicObject, v82:Fixnum[0]):
+bb5(v81:HeapBasicObject, v82:Fixnum):
   PatchPoint NoEPEscape(set_value_loop)
   v89:Fixnum[1] = Const Value(1)
   PatchPoint MethodRedefined(Integer@ADDR, +@0x2b, cme:ADDR)
-  v114:Fixnum[1] = Const Value(1)
+  v114:Fixnum = FixnumAdd v82, v89
   Jump bb6(v81, v114)
```

(the `Fixnum[0]` is from type inference and the `Fixnum[1]`
is from constant folding having folded the `FixnumAdd`)

In the meantime, go back to looping over RPO repeatedly.

Fix Shopify#974

Author: tekknolagi

zjit/src/codegen_tests.rs

@@ -5610,3 +5610,25 @@ fn test_load_immediates_into_registers_before_masking() {
         test
     "#), @"true");
 }
+
+#[test]
+fn test_loop_terminates() {
+    set_call_threshold(3);
+    // Previous worklist-based type inference only worked for maximal SSA. This is a regression
+    // test for hanging.
+    assert_snapshot!(inspect(r#"
+        class TheClass
+          def set_value_loop
+            i = 0
+            while i < 10
+              @levar ||= i
+              i += 1
+            end
+          end
+        end
+
+        3.times do |i|
+          TheClass.new.set_value_loop
+        end
+    "#), @"3");
+}

zjit/src/hir.rs

@@ -3204,82 +3204,66 @@ impl Function {
         self.copy_param_types();
 
         let mut reachable = BlockSet::with_capacity(self.blocks.len());
-
-        // Maintain both a worklist and a fast membership check to avoid linear search
-        let mut worklist: VecDeque<BlockId> = VecDeque::with_capacity(self.blocks.len());
-        let mut in_worklist = BlockSet::with_capacity(self.blocks.len());
-        macro_rules! worklist_add {
-            ($block:expr) => {
-                if in_worklist.insert($block) {
-                    worklist.push_back($block);
-                }
-            };
-        }
-
         reachable.insert(self.entries_block);
-        worklist_add!(self.entries_block);
-
-        // Helper to propagate types along a branch edge and enqueue the target if anything changed
-        macro_rules! enqueue {
-            ($self:ident, $target:expr) => {
-                let newly_reachable = reachable.insert($target.target);
-                let mut target_changed = newly_reachable;
-                for (idx, arg) in $target.args.iter().enumerate() {
-                    let arg = self.union_find.borrow().find_const(*arg);
-                    let param = $self.blocks[$target.target.0].params[idx];
-                    let param = self.union_find.borrow().find_const(param);
-                    let new = self.insn_types[param.0].union(self.insn_types[arg.0]);
-                    if !self.insn_types[param.0].bit_equal(new) {
-                        self.insn_types[param.0] = new;
-                        target_changed = true;
-                    }
-                }
-                if target_changed {
-                    worklist_add!($target.target);
-                }
-            };
-        }
 
-        // Walk the graph, computing types until worklist is empty
-        while let Some(block) = worklist.pop_front() {
-            in_worklist.remove(block);
-            if !reachable.get(block) { continue; }
-            for insn_id in &self.blocks[block.0].insns {
-                let insn_id = self.union_find.borrow().find_const(*insn_id);
-                let insn_type = match &self.insns[insn_id.0] {
-                    &Insn::IfTrue { val, ref target } => {
-                        assert!(!self.type_of(val).bit_equal(types::Empty));
-                        if self.type_of(val).could_be(Type::from_cbool(true)) {
-                            enqueue!(self, target);
+        // Walk the graph, computing types until fixpoint
+        let rpo = self.rpo();
+        loop {
+            let mut changed = false;
+            for &block in &rpo {
+                if !reachable.get(block) { continue; }
+                for &insn_id in &self.blocks[block.0].insns {
+                    // Instructions without output, including branch instructions, can't be targets
+                    // of make_equal_to, so we don't need find() here.
+                    let insn_type = match &self.insns[insn_id.0] {
+                        &Insn::IfTrue { val, target: BranchEdge { target, ref args } } => {
+                            assert!(!self.type_of(val).bit_equal(types::Empty));
+                            if self.type_of(val).could_be(Type::from_cbool(true)) {
+                                reachable.insert(target);
+                                for (idx, arg) in args.iter().enumerate() {
+                                    let param = self.blocks[target.0].params[idx];
+                                    self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
+                                }
+                            }
+                            continue;
                         }
-                        continue;
-                    }
-                    &Insn::IfFalse { val, ref target } => {
-                        assert!(!self.type_of(val).bit_equal(types::Empty));
-                        if self.type_of(val).could_be(Type::from_cbool(false)) {
-                            enqueue!(self, target);
+                        &Insn::IfFalse { val, target: BranchEdge { target, ref args } } => {
+                            assert!(!self.type_of(val).bit_equal(types::Empty));
+                            if self.type_of(val).could_be(Type::from_cbool(false)) {
+                                reachable.insert(target);
+                                for (idx, arg) in args.iter().enumerate() {
+                                    let param = self.blocks[target.0].params[idx];
+                                    self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
+                                }
+                            }
+                            continue;
                         }
-                        continue;
-                    }
-                    &Insn::Jump(ref target) => {
-                        enqueue!(self, target);
-                        continue;
-                    }
-                    &Insn::Entries { ref targets } => {
-                        for target in targets {
-                            if reachable.insert(*target) {
-                                worklist_add!(*target);
+                        &Insn::Jump(BranchEdge { target, ref args }) => {
+                            reachable.insert(target);
+                            for (idx, arg) in args.iter().enumerate() {
+                                let param = self.blocks[target.0].params[idx];
+                                self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
                             }
+                            continue;
                         }
-                        continue;
+                        Insn::Entries { targets } => {
+                            for &target in targets {
+                                reachable.insert(target);
+                            }
+                            continue;
+                        }
+                        insn if insn.has_output() => self.infer_type(insn_id),
+                        _ => continue,
+                    };
+                    if !self.type_of(insn_id).bit_equal(insn_type) {
+                        self.insn_types[insn_id.0] = insn_type;
+                        changed = true;
                     }
-                    insn if insn.has_output() => self.infer_type(insn_id),
-                    _ => continue,
-                };
-                if !self.type_of(insn_id).bit_equal(insn_type) {
-                    self.insn_types[insn_id.0] = insn_type;
                 }
             }
+            if !changed {
+                break;
+            }
         }
     }
 

zjit/src/hir/opt_tests.rs

@@ -15544,4 +15544,97 @@ mod hir_opt_tests {
           Return v43
         ");
     }
+
+    #[test]
+    fn test_infer_types_across_non_maximal_basic_blocks() {
+        // Previous worklist-based type inference only worked for maximal SSA. This is a regression
+        // test for hanging.
+        eval("
+            class TheClass
+              def set_value_loop
+                i = 0
+                while i < 10
+                  @levar ||= i
+                  i += 1
+                end
+              end
+            end
+            3.times do |i|
+              TheClass.new.set_value_loop
+            end
+        ");
+        assert_snapshot!(hir_string_proc("TheClass.instance_method(:set_value_loop)"), @"
+        fn set_value_loop@<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[0] = Const Value(0)
+          CheckInterrupts
+          Jump bb6(v8, v13)
+        bb6(v19:BasicObject, v20:Fixnum):
+          v24:Fixnum[10] = Const Value(10)
+          PatchPoint MethodRedefined(Integer@0x1000, <@0x1008, cme:0x1010)
+          v110:BoolExact = FixnumLt v20, v24
+          CheckInterrupts
+          v30:CBool = Test v110
+          IfTrue v30, bb4(v19, v20)
+          v35:NilClass = Const Value(nil)
+          CheckInterrupts
+          Return v35
+        bb4(v40:BasicObject, v41:Fixnum):
+          PatchPoint SingleRactorMode
+          v46:HeapBasicObject = GuardType v40, HeapBasicObject
+          v47:CUInt64 = LoadField v46, :_rbasic_flags@0x1038
+          v49:CUInt64[0xffffffff0000001f] = Const CUInt64(0xffffffff0000001f)
+          v50:CPtr[CPtr(0x1039)] = Const CPtr(0x1039)
+          v51 = RefineType v50, CUInt64
+          v52:CInt64 = IntAnd v47, v49
+          v53:CBool = IsBitEqual v52, v51
+          IfTrue v53, bb8()
+          v57:CUInt64[0xffffffff0000001f] = Const CUInt64(0xffffffff0000001f)
+          v58:CPtr[CPtr(0x103a)] = Const CPtr(0x103a)
+          v59 = RefineType v58, CUInt64
+          v60:CInt64 = IntAnd v47, v57
+          v61:CBool = IsBitEqual v60, v59
+          IfTrue v61, bb9()
+          v97:CShape = LoadField v46, :_shape_id@0x103b
+          v98:CShape[0x103c] = GuardBitEquals v97, CShape(0x103c)
+          v99:BasicObject = LoadField v46, :@levar@0x103d
+          Jump bb7(v99)
+        bb8():
+          v55:BasicObject = LoadField v46, :@levar@0x103d
+          Jump bb7(v55)
+        bb9():
+          v63:NilClass = Const Value(nil)
+          Jump bb7(v63)
+        bb7(v48:BasicObject):
+          CheckInterrupts
+          v69:CBool = Test v48
+          IfTrue v69, bb5(v46, v41)
+          PatchPoint NoEPEscape(set_value_loop)
+          PatchPoint SingleRactorMode
+          v101:CShape = LoadField v46, :_shape_id@0x103b
+          v102:CShape[0x103e] = GuardBitEquals v101, CShape(0x103e)
+          StoreField v46, :@levar@0x103d, v41
+          WriteBarrier v46, v41
+          v105:CShape[0x103c] = Const CShape(0x103c)
+          StoreField v46, :_shape_id@0x103b, v105
+          v79:HeapBasicObject = RefineType v46, HeapBasicObject
+          Jump bb5(v79, v41)
+        bb5(v81:HeapBasicObject, v82:Fixnum):
+          PatchPoint NoEPEscape(set_value_loop)
+          v89:Fixnum[1] = Const Value(1)
+          PatchPoint MethodRedefined(Integer@0x1000, +@0x103f, cme:0x1040)
+          v114:Fixnum = FixnumAdd v82, v89
+          Jump bb6(v81, v114)
+        ");
+    }
 }