JIT: Add helper method for updating bbTargetEdge (#99362)

Part of #93020. Adds fgRedirectTargetEdge, which updates bbTargetEdge's block target while maintaining the rest of the edge's state to simplify edge profile maintenance (and avoid additional allocations).

Author: amanasifkhalid

src/coreclr/jit/compiler.h

@@ -5912,6 +5912,12 @@ class Compiler
     template <bool initializingPreds = false>
     FlowEdge* fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowEdge* oldEdge = nullptr);
 
+private:
+    FlowEdge** fgGetPredInsertPoint(BasicBlock* blockPred, BasicBlock* newTarget);
+
+public:
+    void fgRedirectTargetEdge(BasicBlock* block, BasicBlock* newTarget);
+
     void fgFindBasicBlocks();
 
     bool fgCheckEHCanInsertAfterBlock(BasicBlock* blk, unsigned regionIndex, bool putInTryRegion);

src/coreclr/jit/fgbasic.cpp

@@ -646,6 +646,7 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* oldTarget, Bas
         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());
@@ -4192,9 +4193,7 @@ void Compiler::fgFixEntryFlowForOSR()
     //
     fgEnsureFirstBBisScratch();
     assert(fgFirstBB->KindIs(BBJ_ALWAYS) && fgFirstBB->JumpsToNext());
-    fgRemoveRefPred(fgFirstBB->GetTargetEdge());
-    FlowEdge* const newEdge = fgAddRefPred(fgOSREntryBB, fgFirstBB);
-    fgFirstBB->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
+    fgRedirectTargetEdge(fgFirstBB, fgOSREntryBB);
 
     // We don't know the right weight for this block, since
     // execution of the method was interrupted within the

src/coreclr/jit/fgehopt.cpp

@@ -168,10 +168,8 @@ PhaseStatus Compiler::fgRemoveEmptyFinally()
                 fgRemoveBlock(leaveBlock, /* unreachable */ true);
 
                 // Ref count updates.
-                fgRemoveRefPred(currentBlock->GetTargetEdge());
-                FlowEdge* const newEdge = fgAddRefPred(postTryFinallyBlock, currentBlock);
-
-                currentBlock->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
+                fgRedirectTargetEdge(currentBlock, postTryFinallyBlock);
+                currentBlock->SetKind(BBJ_ALWAYS);
                 currentBlock->RemoveFlags(BBF_RETLESS_CALL); // no longer a BBJ_CALLFINALLY
 
                 // Cleanup the postTryFinallyBlock
@@ -1136,12 +1134,11 @@ PhaseStatus Compiler::fgCloneFinally()
                     fgRemoveBlock(leaveBlock, /* unreachable */ true);
 
                     // Ref count updates.
-                    fgRemoveRefPred(currentBlock->GetTargetEdge());
-                    FlowEdge* const newEdge = fgAddRefPred(firstCloneBlock, currentBlock);
+                    fgRedirectTargetEdge(currentBlock, firstCloneBlock);
 
                     // This call returns to the expected spot, so retarget it to branch to the clone.
                     currentBlock->RemoveFlags(BBF_RETLESS_CALL); // no longer a BBJ_CALLFINALLY
-                    currentBlock->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
+                    currentBlock->SetKind(BBJ_ALWAYS);
 
                     // Make sure iteration isn't going off the deep end.
                     assert(leaveBlock != endCallFinallyRangeBlock);
@@ -1758,9 +1755,7 @@ bool Compiler::fgRetargetBranchesToCanonicalCallFinally(BasicBlock*      block,
             canonicalCallFinally->bbNum);
 
     assert(callFinally->bbRefs > 0);
-    fgRemoveRefPred(block->GetTargetEdge());
-    FlowEdge* const newEdge = fgAddRefPred(canonicalCallFinally, block);
-    block->SetTargetEdge(newEdge);
+    fgRedirectTargetEdge(block, canonicalCallFinally);
 
     // Update profile counts
     //
@@ -2132,16 +2127,11 @@ void Compiler::fgTailMergeThrowsJumpToHelper(BasicBlock* predBlock,
 {
     JITDUMP("*** " FMT_BB " now branching to " FMT_BB "\n", predBlock->bbNum, canonicalBlock->bbNum);
 
-    FlowEdge* const newEdge = fgAddRefPred(canonicalBlock, predBlock, predEdge);
-
     if (predBlock->KindIs(BBJ_ALWAYS))
     {
-        // Remove the old flow
+        // Update flow to new target
         assert(predBlock->TargetIs(nonCanonicalBlock));
-        fgRemoveRefPred(predBlock->GetTargetEdge());
-
-        // Wire up the new flow
-        predBlock->SetTargetEdge(newEdge);
+        fgRedirectTargetEdge(predBlock, canonicalBlock);
     }
     else
     {
@@ -2151,6 +2141,7 @@ void Compiler::fgTailMergeThrowsJumpToHelper(BasicBlock* predBlock,
         fgRemoveRefPred(predBlock->GetTrueEdge());
 
         // Wire up the new flow
-        predBlock->SetTrueEdge(newEdge);
+        FlowEdge* const trueEdge = fgAddRefPred(canonicalBlock, predBlock, predEdge);
+        predBlock->SetTrueEdge(trueEdge);
     }
 }

src/coreclr/jit/fgflow.cpp

@@ -117,15 +117,16 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
     // dependency on this order. Note also that we don't allow duplicates in the list; we maintain a DupCount
     // count of duplication. This also necessitates walking the flow list for every edge we add.
     //
-    FlowEdge*  flow  = nullptr;
-    FlowEdge** listp = &block->bbPreds;
+    FlowEdge*  flow = nullptr;
+    FlowEdge** listp;
 
     if (initializingPreds)
     {
         // List is sorted order and we're adding references in
         // increasing blockPred->bbNum order. The only possible
         // dup list entry is the last one.
         //
+        listp              = &block->bbPreds;
         FlowEdge* flowLast = block->bbLastPred;
         if (flowLast != nullptr)
         {
@@ -143,10 +144,7 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
     {
         // References are added randomly, so we have to search.
         //
-        while ((*listp != nullptr) && ((*listp)->getSourceBlock()->bbNum < blockPred->bbNum))
-        {
-            listp = (*listp)->getNextPredEdgeRef();
-        }
+        listp = fgGetPredInsertPoint(blockPred, block);
 
         if ((*listp != nullptr) && ((*listp)->getSourceBlock() == blockPred))
         {
@@ -381,6 +379,76 @@ void Compiler::fgRemoveBlockAsPred(BasicBlock* block)
     }
 }
 
+//------------------------------------------------------------------------
+// fgGetPredInsertPoint: Searches newTarget->bbPreds for where to insert an edge from blockPred.
+//
+// Arguments:
+//    blockPred -- The block we want to make a predecessor of newTarget (it could already be one).
+//    newTarget -- The block whose pred list we will search.
+//
+// Return Value:
+//    Returns a pointer to the next pointer of an edge in newTarget's pred list.
+//    A new edge from blockPred to newTarget can be inserted here
+//    without disrupting bbPreds' sorting invariant.
+//
+FlowEdge** Compiler::fgGetPredInsertPoint(BasicBlock* blockPred, BasicBlock* newTarget)
+{
+    assert(blockPred != nullptr);
+    assert(newTarget != nullptr);
+    assert(fgPredsComputed);
+
+    FlowEdge** listp = &newTarget->bbPreds;
+
+    // Search pred list for insertion point
+    //
+    while ((*listp != nullptr) && ((*listp)->getSourceBlock()->bbNum < blockPred->bbNum))
+    {
+        listp = (*listp)->getNextPredEdgeRef();
+    }
+
+    return listp;
+}
+
+//------------------------------------------------------------------------
+// fgRedirectTargetEdge: Sets block->bbTargetEdge's target block to newTarget,
+// updating pred lists as necessary.
+//
+// Arguments:
+//    block -- The block we want to make a predecessor of newTarget.
+//             It could be one already, in which case nothing changes.
+//    newTarget -- The new successor of block.
+//
+void Compiler::fgRedirectTargetEdge(BasicBlock* block, BasicBlock* newTarget)
+{
+    assert(block != nullptr);
+    assert(newTarget != nullptr);
+
+    FlowEdge* edge = block->GetTargetEdge();
+    assert(edge->getDupCount() == 1);
+
+    // Update oldTarget's pred list.
+    // We could call fgRemoveRefPred, but since we're removing the one and only ref from block to oldTarget,
+    // fgRemoveAllRefPreds is slightly more efficient (one fewer branch, doesn't update edge's dup count, etc).
+    //
+    BasicBlock* oldTarget = edge->getDestinationBlock();
+    fgRemoveAllRefPreds(oldTarget, block);
+
+    // Splice edge into new target block's pred list
+    //
+    FlowEdge** predListPtr = fgGetPredInsertPoint(block, newTarget);
+    edge->setNextPredEdge(*predListPtr);
+    edge->setDestinationBlock(newTarget);
+    *predListPtr = edge;
+
+    // Pred list of target should (still) be ordered
+    //
+    assert(newTarget->checkPredListOrder());
+
+    // Edge should still have only one ref
+    assert(edge->getDupCount() == 1);
+    newTarget->bbRefs++;
+}
+
 Compiler::SwitchUniqueSuccSet Compiler::GetDescriptorForSwitch(BasicBlock* switchBlk)
 {
     assert(switchBlk->KindIs(BBJ_SWITCH));

src/coreclr/jit/fginline.cpp

@@ -1550,11 +1550,9 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo)
         // Insert inlinee's blocks into inliner's block list.
         assert(topBlock->KindIs(BBJ_ALWAYS));
         assert(topBlock->TargetIs(bottomBlock));
-        fgRemoveRefPred(topBlock->GetTargetEdge());
-        FlowEdge* const newEdge = fgAddRefPred(InlineeCompiler->fgFirstBB, topBlock, topBlock->GetTargetEdge());
+        fgRedirectTargetEdge(topBlock, InlineeCompiler->fgFirstBB);
 
         topBlock->SetNext(InlineeCompiler->fgFirstBB);
-        topBlock->SetTargetEdge(newEdge);
         topBlock->SetFlags(BBF_NONE_QUIRK);
         InlineeCompiler->fgLastBB->SetNext(bottomBlock);
 

src/coreclr/jit/fgopt.cpp

@@ -834,9 +834,7 @@ PhaseStatus Compiler::fgPostImportationCleanup()
 
                 if (entryJumpTarget != osrEntry)
                 {
-                    fgRemoveRefPred(fgFirstBB->GetTargetEdge());
-                    FlowEdge* const newEdge = fgAddRefPred(entryJumpTarget, fgFirstBB, fgFirstBB->GetTargetEdge());
-                    fgFirstBB->SetTargetEdge(newEdge);
+                    fgRedirectTargetEdge(fgFirstBB, entryJumpTarget);
 
                     JITDUMP("OSR: redirecting flow from method entry " FMT_BB " to OSR entry " FMT_BB
                             " via step blocks.\n",
@@ -1570,9 +1568,7 @@ bool Compiler::fgOptimizeBranchToEmptyUnconditional(BasicBlock* block, BasicBloc
             case BBJ_ALWAYS:
             case BBJ_CALLFINALLYRET:
             {
-                fgRemoveRefPred(block->GetTargetEdge());
-                FlowEdge* const newEdge = fgAddRefPred(bDest->GetTarget(), block, block->GetTargetEdge());
-                block->SetTargetEdge(newEdge);
+                fgRedirectTargetEdge(block, bDest->GetTarget());
                 break;
             }
 

src/coreclr/jit/helperexpansion.cpp

@@ -372,7 +372,6 @@ bool Compiler::fgExpandRuntimeLookupsForCall(BasicBlock** pBlock, Statement* stm
     // Update preds in all new blocks
     //
     assert(prevBb->KindIs(BBJ_ALWAYS));
-    fgRemoveRefPred(prevBb->GetTargetEdge());
 
     {
         FlowEdge* const newEdge = fgAddRefPred(block, fastPathBb);
@@ -389,8 +388,7 @@ bool Compiler::fgExpandRuntimeLookupsForCall(BasicBlock** pBlock, Statement* stm
     if (needsSizeCheck)
     {
         // sizeCheckBb is the first block after prevBb
-        FlowEdge* const newEdge = fgAddRefPred(sizeCheckBb, prevBb);
-        prevBb->SetTargetEdge(newEdge);
+        fgRedirectTargetEdge(prevBb, sizeCheckBb);
 
         // sizeCheckBb flows into nullcheckBb in case if the size check passes
         {
@@ -406,8 +404,7 @@ bool Compiler::fgExpandRuntimeLookupsForCall(BasicBlock** pBlock, Statement* stm
     else
     {
         // nullcheckBb is the first block after prevBb
-        FlowEdge* const newEdge = fgAddRefPred(nullcheckBb, prevBb);
-        prevBb->SetTargetEdge(newEdge);
+        fgRedirectTargetEdge(prevBb, nullcheckBb);
 
         // No size check, nullcheckBb jumps to fast path
         // fallbackBb is only reachable from nullcheckBb (jump destination)
@@ -742,9 +739,7 @@ bool Compiler::fgExpandThreadLocalAccessForCallNativeAOT(BasicBlock** pBlock, St
     // fallback will just execute first time
     fallbackBb->bbSetRunRarely();
 
-    fgRemoveRefPred(prevBb->GetTargetEdge());
-    FlowEdge* const newEdge = fgAddRefPred(tlsRootNullCondBB, prevBb);
-    prevBb->SetTargetEdge(newEdge);
+    fgRedirectTargetEdge(prevBb, tlsRootNullCondBB);
 
     // All blocks are expected to be in the same EH region
     assert(BasicBlock::sameEHRegion(prevBb, block));
@@ -1089,12 +1084,7 @@ bool Compiler::fgExpandThreadLocalAccessForCall(BasicBlock** pBlock, Statement*
     // Update preds in all new blocks
     //
     assert(prevBb->KindIs(BBJ_ALWAYS));
-    fgRemoveRefPred(prevBb->GetTargetEdge());
-
-    {
-        FlowEdge* const newEdge = fgAddRefPred(maxThreadStaticBlocksCondBB, prevBb);
-        prevBb->SetTargetEdge(newEdge);
-    }
+    fgRedirectTargetEdge(prevBb, maxThreadStaticBlocksCondBB);
 
     {
         FlowEdge* const trueEdge  = fgAddRefPred(fallbackBb, maxThreadStaticBlocksCondBB);
@@ -1462,8 +1452,10 @@ bool Compiler::fgExpandStaticInitForCall(BasicBlock** pBlock, Statement* stmt, G
     // Update preds in all new blocks
     //
 
-    // Unlink block and prevBb
-    fgRemoveRefPred(prevBb->GetTargetEdge());
+    // Redirect prevBb from block to isInitedBb
+    fgRedirectTargetEdge(prevBb, isInitedBb);
+    prevBb->SetFlags(BBF_NONE_QUIRK);
+    assert(prevBb->JumpsToNext());
 
     {
         // Block has two preds now: either isInitedBb or helperCallBb
@@ -1473,15 +1465,6 @@ bool Compiler::fgExpandStaticInitForCall(BasicBlock** pBlock, Statement* stmt, G
         helperCallBb->SetFlags(BBF_NONE_QUIRK);
     }
 
-    {
-        // prevBb always flows into isInitedBb
-        assert(prevBb->KindIs(BBJ_ALWAYS));
-        FlowEdge* const newEdge = fgAddRefPred(isInitedBb, prevBb);
-        prevBb->SetTargetEdge(newEdge);
-        prevBb->SetFlags(BBF_NONE_QUIRK);
-        assert(prevBb->JumpsToNext());
-    }
-
     {
         // Both fastPathBb and helperCallBb have a single common pred - isInitedBb
         FlowEdge* const trueEdge  = fgAddRefPred(block, isInitedBb);
@@ -1792,17 +1775,10 @@ bool Compiler::fgVNBasedIntrinsicExpansionForCall_ReadUtf8(BasicBlock** pBlock,
     //
     // Update preds in all new blocks
     //
-    // block is no longer a predecessor of prevBb
-    fgRemoveRefPred(prevBb->GetTargetEdge());
-
-    {
-        // prevBb flows into lengthCheckBb
-        assert(prevBb->KindIs(BBJ_ALWAYS));
-        FlowEdge* const newEdge = fgAddRefPred(lengthCheckBb, prevBb);
-        prevBb->SetTargetEdge(newEdge);
-        prevBb->SetFlags(BBF_NONE_QUIRK);
-        assert(prevBb->JumpsToNext());
-    }
+    // Redirect prevBb to lengthCheckBb
+    fgRedirectTargetEdge(prevBb, lengthCheckBb);
+    prevBb->SetFlags(BBF_NONE_QUIRK);
+    assert(prevBb->JumpsToNext());
 
     {
         // lengthCheckBb has two successors: block and fastpathBb
@@ -2511,12 +2487,7 @@ bool Compiler::fgLateCastExpansionForCall(BasicBlock** pBlock, Statement* stmt,
         }
     }
 
-    fgRemoveRefPred(firstBb->GetTargetEdge());
-
-    {
-        FlowEdge* const newEdge = fgAddRefPred(nullcheckBb, firstBb);
-        firstBb->SetTargetEdge(newEdge);
-    }
+    fgRedirectTargetEdge(firstBb, nullcheckBb);
 
     {
         FlowEdge* const trueEdge = fgAddRefPred(lastBb, nullcheckBb);