From c1f2611e434d98f4b0732888b8e5b6e55974761b Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Tue, 30 Nov 2021 21:20:21 +0900
Subject: [PATCH] optimizer: fix #43254, avoid infinite CFG traversal in SROA's
 domination analysis

Since CFG can be cyclic, the previous implementation of `has_safe_def`
that walks predecessors recursively was just wrong.

This commit fixes it by making it `has_safe_def` maintain a single set that
keeps the identities of basic blocks that have been examined already.

Since we still examine each predecessor recursively, this approach doesn't
handle cycles very nicely (this commit doesn't come with any sort of fancy
tags that indicates a basic block is "being examined", as opposed to inference).
But I don't see much motivation to make `has_safe_def` more accurate while
keeping it correct, since it's called only in rare cases.
---
 base/compiler/ssair/passes.jl | 53 ++++++++++++++++++++++++-----------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index d1be326529112..2ad5ae0daf248 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -97,20 +97,39 @@ end
 # if this load at `idx` have any "safe" `setfield!` calls that define the field
 function has_safe_def(
     ir::IRCode, domtree::DomTree, allblocks::Vector{Int}, du::SSADefUse,
-    newidx::Int, idx::Int, inclusive::Bool = false)
-    def = first(find_def_for_use(ir, domtree, allblocks, du, idx, inclusive))
-
-    # this field is supposed to be defined at the `:new` site (but it's not and thus this load will throw)
+    newidx::Int, idx::Int)
+    def, _, _ = find_def_for_use(ir, domtree, allblocks, du, idx)
+    # will throw since we already checked this `:new` site doesn't define this field
     def == newidx && return false
+    # found a "safe" definition
+    def ≠ 0 && return true
+    # we may still be able to replace this load with `PhiNode`
+    block = block_for_inst(ir, idx)
+    seen = BitSet(block)
+    return has_safe_def(ir, domtree, allblocks, du, newidx, block, seen)
+end
 
-    def ≠ 0 && return true # found a "safe" definition
-
-    # we may be able to replace this load with `PhiNode` if all the predecessors have "safe" definitions
-    idxblock = block_for_inst(ir, idx)
-    for pred in ir.cfg.blocks[idxblock].preds
-        lastidx = last(ir.cfg.blocks[pred].stmts)
-        # NOTE `lastidx` isn't a load, thus we can use inclusive coondition within the `find_def_for_use`
-        has_safe_def(ir, domtree, allblocks, du, newidx, lastidx, true) || return false
+# recursively examine if all predecessors of `block` have any "safe" definition
+# XXX this implementation avoids infinite cycles by maintaining a sigle set of seen basic blocks,
+# and obviously it doesn't handle cycles very well
+function has_safe_def(
+    ir::IRCode, domtree::DomTree, allblocks::Vector{Int}, du::SSADefUse,
+    newidx::Int, block::Int, seen::BitSet)
+    preds = ir.cfg.blocks[block].preds
+    isempty(preds) && return false
+    for pred in preds
+        # if this block has already been examined, we should bail out to avoid infinite cycles
+        pred in seen && return false
+        idx = last(ir.cfg.blocks[pred].stmts)
+        # NOTE `idx` isn't a load, thus we can use inclusive coondition within the `find_def_for_use`
+        def, _, _ = find_def_for_use(ir, domtree, allblocks, du, idx, true)
+        # will throw since we already checked this `:new` site doesn't define this field
+        def == newidx && return false
+        # found a "safe" definition for this predecessor
+        def ≠ 0 && continue
+        # check for the predecessors of this predecessor
+        push!(seen, pred)
+        has_safe_def(ir, domtree, allblocks, du, newidx, pred, seen) || return false
     end
     return true
 end
@@ -599,8 +618,8 @@ function perform_lifting!(compact::IncrementalCompact,
     return stmt_val # N.B. should never happen
 end
 
-# NOTE we use `IdSet{Int}` instead of `BitSet` for `sroa_pass!` since it works on IR after inlining,
-# which can be very large sometimes, and analyzed program counters are often very sparse
+# NOTE we use `IdSet{Int}` instead of `BitSet` for in these passes since they work on IR after inlining,
+# which can be very large sometimes, and program counters in question are often very sparse
 const SPCSet = IdSet{Int}
 
 """
@@ -897,8 +916,8 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse
                     end
                 end
                 for b in phiblocks
+                    n = ir[phinodes[b]]::PhiNode
                     for p in ir.cfg.blocks[b].preds
-                        n = ir[phinodes[b]]::PhiNode
                         push!(n.edges, p)
                         push!(n.values, compute_value_for_block(ir, domtree,
                             allblocks, du, phinodes, fidx, p))
@@ -967,7 +986,7 @@ function count_uses(@nospecialize(stmt), uses::Vector{Int})
     end
 end
 
-function mark_phi_cycles!(compact::IncrementalCompact, safe_phis::BitSet, phi::Int)
+function mark_phi_cycles!(compact::IncrementalCompact, safe_phis::SPCSet, phi::Int)
     worklist = Int[]
     push!(worklist, phi)
     while !isempty(worklist)
@@ -1037,7 +1056,7 @@ function adce_pass!(ir::IRCode)
     changed = true
     while changed
         changed = false
-        safe_phis = BitSet()
+        safe_phis = SPCSet()
         for phi in all_phis
             # Save any phi cycles that have non-phi uses
             if compact.used_ssas[phi] - phi_uses[phi] != 0