JIT: Add pred block iterator that tolerates pred list modifications (#99466)

Follow-up to #99362. fgRedirectTargetEdge modifies the predecessor lists of the old and new successor blocks, so if we want to be able to use it in fgReplaceJumpTarget, we need a pred block iterator that is resilient to these changes, as we frequently call the latter method while iterating predecessors.

Author: amanasifkhalid

src/coreclr/jit/block.h

@@ -258,7 +258,9 @@ class PredEdgeList
 // PredBlockList: adapter class for forward iteration of the predecessor edge linked list yielding
 // predecessor blocks, using range-based `for`, normally used via BasicBlock::PredBlocks(), e.g.:
 //    for (BasicBlock* const predBlock : block->PredBlocks()) ...
+// allowEdits controls whether the iterator should be resilient to changes to the predecessor list.
 //
+template <bool allowEdits>
 class PredBlockList
 {
     FlowEdge* m_begin;
@@ -270,13 +272,12 @@ class PredBlockList
     {
         FlowEdge* m_pred;
 
-#ifdef DEBUG
-        // Try to guard against the user of the iterator from making changes to the IR that would invalidate
-        // the iterator: cache the edge we think should be next, then check it when we actually do the `++`
+        // When allowEdits=false, try to guard against the user of the iterator from modifying the predecessor list
+        // being traversed: cache the edge we think should be next, then check it when we actually do the `++`
         // operation. This is a bit conservative, but attempts to protect against callers assuming too much about
         // this iterator implementation.
+        // When allowEdits=true, m_next is always used to update m_pred, so changes to m_pred don't break the iterator.
         FlowEdge* m_next;
-#endif
 
     public:
         iterator(FlowEdge* pred);
@@ -1530,9 +1531,18 @@ struct BasicBlock : private LIR::Range
     // PredBlocks: convenience method for enabling range-based `for` iteration over predecessor blocks, e.g.:
     //    for (BasicBlock* const predBlock : block->PredBlocks()) ...
     //
-    PredBlockList PredBlocks() const
+    PredBlockList<false> PredBlocks() const
+    {
+        return PredBlockList<false>(bbPreds);
+    }
+
+    // PredBlocksEditing: convenience method for enabling range-based `for` iteration over predecessor blocks, e.g.:
+    //    for (BasicBlock* const predBlock : block->PredBlocksList()) ...
+    // This iterator tolerates modifications to bbPreds.
+    //
+    PredBlockList<true> PredBlocksEditing() const
     {
-        return PredBlockList(bbPreds);
+        return PredBlockList<true>(bbPreds);
     }
 
     // Pred list maintenance
@@ -2399,29 +2409,45 @@ inline PredEdgeList::iterator& PredEdgeList::iterator::operator++()
     return *this;
 }
 
-inline PredBlockList::iterator::iterator(FlowEdge* pred) : m_pred(pred)
+template <bool allowEdits>
+inline PredBlockList<allowEdits>::iterator::iterator(FlowEdge* pred) : m_pred(pred)
 {
-#ifdef DEBUG
-    m_next = (m_pred == nullptr) ? nullptr : m_pred->getNextPredEdge();
-#endif
+    bool initNextPointer = allowEdits;
+    INDEBUG(initNextPointer = true);
+    if (initNextPointer)
+    {
+        m_next = (m_pred == nullptr) ? nullptr : m_pred->getNextPredEdge();
+    }
 }
 
-inline BasicBlock* PredBlockList::iterator::operator*() const
+template <bool     allowEdits>
+inline BasicBlock* PredBlockList<allowEdits>::iterator::operator*() const
 {
     return m_pred->getSourceBlock();
 }
 
-inline PredBlockList::iterator& PredBlockList::iterator::operator++()
+template <bool                                       allowEdits>
+inline typename PredBlockList<allowEdits>::iterator& PredBlockList<allowEdits>::iterator::operator++()
 {
-    FlowEdge* next = m_pred->getNextPredEdge();
+    if (allowEdits)
+    {
+        // For editing iterators, m_next is always used and maintained
+        m_pred = m_next;
+        m_next = (m_next == nullptr) ? nullptr : m_next->getNextPredEdge();
+    }
+    else
+    {
+        FlowEdge* next = m_pred->getNextPredEdge();
 
 #ifdef DEBUG
-    // Check that the next block is the one we expect to see.
-    assert(next == m_next);
-    m_next = (next == nullptr) ? nullptr : next->getNextPredEdge();
+        // If allowEdits=false, check that the next block is the one we expect to see.
+        assert(next == m_next);
+        m_next = (m_next == nullptr) ? nullptr : m_next->getNextPredEdge();
 #endif // DEBUG
 
-    m_pred = next;
+        m_pred = next;
+    }
+
     return *this;
 }
 

src/coreclr/jit/fgbasic.cpp

@@ -645,17 +645,11 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* oldTarget, Bas
         case BBJ_EHCATCHRET:
         case BBJ_EHFILTERRET:
         case BBJ_LEAVE: // This function can be called before import, so we still have BBJ_LEAVE
-        {
-            // TODO: Use fgRedirectTargetEdge once pred edge iterators are resilient to bbPreds being modified.
             assert(block->TargetIs(oldTarget));
-            fgRemoveRefPred(block->GetTargetEdge());
-            FlowEdge* const newEdge = fgAddRefPred(newTarget, block, block->GetTargetEdge());
-            block->SetTargetEdge(newEdge);
+            fgRedirectTargetEdge(block, newTarget);
             break;
-        }
 
         case BBJ_COND:
-
             if (block->TrueTargetIs(oldTarget))
             {
                 FlowEdge* const oldEdge = block->GetTrueEdge();
@@ -5297,7 +5291,7 @@ BasicBlock* Compiler::fgRemoveBlock(BasicBlock* block, bool unreachable)
 
         fgRemoveRefPred(block->GetTargetEdge());
 
-        for (BasicBlock* const predBlock : block->PredBlocks())
+        for (BasicBlock* const predBlock : block->PredBlocksEditing())
         {
             /* change all jumps/refs to the removed block */
             switch (predBlock->GetKind())

src/coreclr/jit/fgehopt.cpp

@@ -2002,36 +2002,30 @@ PhaseStatus Compiler::fgTailMergeThrows()
 
         // Walk pred list of the non canonical block, updating flow to target
         // the canonical block instead.
-        for (FlowEdge* predEdge = nonCanonicalBlock->bbPreds; predEdge != nullptr; predEdge = nextPredEdge)
+        for (BasicBlock* const predBlock : nonCanonicalBlock->PredBlocksEditing())
         {
-            BasicBlock* const predBlock = predEdge->getSourceBlock();
-            nextPredEdge                = predEdge->getNextPredEdge();
-
             switch (predBlock->GetKind())
             {
-                case BBJ_ALWAYS:
-                {
-                    fgTailMergeThrowsJumpToHelper(predBlock, nonCanonicalBlock, canonicalBlock, predEdge);
-                    updated = true;
-                }
-                break;
-
+                // TODO: Use fgReplaceJumpTarget once it preserves edge weights for BBJ_COND blocks
                 case BBJ_COND:
                 {
-                    // Flow to non canonical block could be via fall through or jump or both.
+                    // Flow to non canonical block could be via true or false target.
                     if (predBlock->FalseTargetIs(nonCanonicalBlock))
                     {
-                        fgTailMergeThrowsFallThroughHelper(predBlock, nonCanonicalBlock, canonicalBlock, predEdge);
+                        fgTailMergeThrowsFallThroughHelper(predBlock, nonCanonicalBlock, canonicalBlock,
+                                                           predBlock->GetFalseEdge());
                     }
 
                     if (predBlock->TrueTargetIs(nonCanonicalBlock))
                     {
-                        fgTailMergeThrowsJumpToHelper(predBlock, nonCanonicalBlock, canonicalBlock, predEdge);
+                        fgTailMergeThrowsJumpToHelper(predBlock, nonCanonicalBlock, canonicalBlock,
+                                                      predBlock->GetTrueEdge());
                     }
                     updated = true;
                 }
                 break;
 
+                case BBJ_ALWAYS:
                 case BBJ_SWITCH:
                 {
                     JITDUMP("*** " FMT_BB " now branching to " FMT_BB "\n", predBlock->bbNum, canonicalBlock->bbNum);
@@ -2127,21 +2121,12 @@ void Compiler::fgTailMergeThrowsJumpToHelper(BasicBlock* predBlock,
 {
     JITDUMP("*** " FMT_BB " now branching to " FMT_BB "\n", predBlock->bbNum, canonicalBlock->bbNum);
 
-    if (predBlock->KindIs(BBJ_ALWAYS))
-    {
-        // Update flow to new target
-        assert(predBlock->TargetIs(nonCanonicalBlock));
-        fgRedirectTargetEdge(predBlock, canonicalBlock);
-    }
-    else
-    {
-        // Remove the old flow
-        assert(predBlock->KindIs(BBJ_COND));
-        assert(predBlock->TrueTargetIs(nonCanonicalBlock));
-        fgRemoveRefPred(predBlock->GetTrueEdge());
-
-        // Wire up the new flow
-        FlowEdge* const trueEdge = fgAddRefPred(canonicalBlock, predBlock, predEdge);
-        predBlock->SetTrueEdge(trueEdge);
-    }
+    // Remove the old flow
+    assert(predBlock->KindIs(BBJ_COND));
+    assert(predBlock->TrueTargetIs(nonCanonicalBlock));
+    fgRemoveRefPred(predBlock->GetTrueEdge());
+
+    // Wire up the new flow
+    FlowEdge* const trueEdge = fgAddRefPred(canonicalBlock, predBlock, predEdge);
+    predBlock->SetTrueEdge(trueEdge);
 }

src/coreclr/jit/optimizer.cpp

@@ -2247,7 +2247,7 @@ bool Compiler::optInvertWhileLoop(BasicBlock* block)
     //
     unsigned const loopFirstNum  = bTop->bbNum;
     unsigned const loopBottomNum = bTest->bbNum;
-    for (BasicBlock* const predBlock : bTest->PredBlocks())
+    for (BasicBlock* const predBlock : bTest->PredBlocksEditing())
     {
         unsigned const bNum = predBlock->bbNum;
         if ((loopFirstNum <= bNum) && (bNum <= loopBottomNum))
@@ -3154,7 +3154,7 @@ bool Compiler::optCanonicalizeExit(FlowGraphNaturalLoop* loop, BasicBlock* exit)
     JITDUMP("Created new exit " FMT_BB " to replace " FMT_BB " for " FMT_LP "\n", newExit->bbNum, exit->bbNum,
             loop->GetIndex());
 
-    for (BasicBlock* pred : exit->PredBlocks())
+    for (BasicBlock* pred : exit->PredBlocksEditing())
     {
         if (loop->ContainsBlock(pred))
         {

src/coreclr/jit/redundantbranchopts.cpp

@@ -1600,7 +1600,7 @@ bool Compiler::optJumpThreadCore(JumpThreadInfo& jti)
     // If this pred is in the set that will reuse block, do nothing.
     // Else revise pred to branch directly to the appropriate successor of block.
     //
-    for (BasicBlock* const predBlock : jti.m_block->PredBlocks())
+    for (BasicBlock* const predBlock : jti.m_block->PredBlocksEditing())
     {
         // If this was an ambiguous pred, skip.
         //