JIT: Use successor edges instead of block targets for BBJ_EHFINALLYRET (#98563)

Part of #93020. This is a first step in a broader effort to replace BasicBlock's jump targets (bbTarget, bbFalseTarget, etc.) with successor flow edges. This PR also adds overloaded implementations for commonly used edge maintenance methods (fgReplacePred, fgRemoveRefPred, etc.) that can take advantage of the cheap access to successor edges (i.e. without calling fgGetPredForBlock -- hopefully we'll be able to get rid of that at some point).

Author: amanasifkhalid

src/coreclr/jit/block.cpp

@@ -658,12 +658,12 @@ void BasicBlock::dspKind() const
             }
             else
             {
-                const unsigned     jumpCnt = bbEhfTargets->bbeCount;
-                BasicBlock** const jumpTab = bbEhfTargets->bbeSuccs;
+                const unsigned   jumpCnt = bbEhfTargets->bbeCount;
+                FlowEdge** const jumpTab = bbEhfTargets->bbeSuccs;
 
                 for (unsigned i = 0; i < jumpCnt; i++)
                 {
-                    printf("%c%s", (i == 0) ? ' ' : ',', dspBlockNum(jumpTab[i]));
+                    printf("%c%s", (i == 0) ? ' ' : ',', dspBlockNum(jumpTab[i]->getDestinationBlock()));
                 }
             }
 
@@ -1214,7 +1214,7 @@ BasicBlock* BasicBlock::GetSucc(unsigned i) const
             }
 
         case BBJ_EHFINALLYRET:
-            return bbEhfTargets->bbeSuccs[i];
+            return bbEhfTargets->bbeSuccs[i]->getDestinationBlock();
 
         case BBJ_SWITCH:
             return bbSwtTargets->bbsDstTab[i];
@@ -1315,7 +1315,7 @@ BasicBlock* BasicBlock::GetSucc(unsigned i, Compiler* comp)
         case BBJ_EHFINALLYRET:
             assert(bbEhfTargets != nullptr);
             assert(i < bbEhfTargets->bbeCount);
-            return bbEhfTargets->bbeSuccs[i];
+            return bbEhfTargets->bbeSuccs[i]->getDestinationBlock();
 
         case BBJ_CALLFINALLY:
         case BBJ_CALLFINALLYRET:
@@ -1792,7 +1792,7 @@ BBehfDesc::BBehfDesc(Compiler* comp, const BBehfDesc* other) : bbeCount(other->b
 {
     // Allocate and fill in a new dst tab
     //
-    bbeSuccs = new (comp, CMK_BasicBlock) BasicBlock*[bbeCount];
+    bbeSuccs = new (comp, CMK_FlowEdge) FlowEdge*[bbeCount];
     for (unsigned i = 0; i < bbeCount; i++)
     {
         bbeSuccs[i] = other->bbeSuccs[i];

src/coreclr/jit/block.h

@@ -314,21 +314,27 @@ class PredBlockList
 //
 class BBArrayIterator
 {
-    BasicBlock* const* m_bbEntry;
+    // Quirk: Some BasicBlock kinds refer to their successors with BasicBlock pointers,
+    // while others use FlowEdge pointers. Eventually, every type will use FlowEdge pointers.
+    // For now, support iterating with both types.
+    union {
+        BasicBlock* const* m_bbEntry;
+        FlowEdge* const*   m_edgeEntry;
+    };
+
+    bool iterateEdges;
 
 public:
-    BBArrayIterator(BasicBlock* const* bbEntry) : m_bbEntry(bbEntry)
+    BBArrayIterator(BasicBlock* const* bbEntry) : m_bbEntry(bbEntry), iterateEdges(false)
     {
     }
 
-    BasicBlock* operator*() const
+    BBArrayIterator(FlowEdge* const* edgeEntry) : m_edgeEntry(edgeEntry), iterateEdges(true)
     {
-        assert(m_bbEntry != nullptr);
-        BasicBlock* bTarget = *m_bbEntry;
-        assert(bTarget != nullptr);
-        return bTarget;
     }
 
+    BasicBlock* operator*() const;
+
     BBArrayIterator& operator++()
     {
         assert(m_bbEntry != nullptr);
@@ -1570,9 +1576,22 @@ struct BasicBlock : private LIR::Range
         // need to call a function or execute another `switch` to get them. Also, pre-compute the begin and end
         // points of the iteration, for use by BBArrayIterator. `m_begin` and `m_end` will either point at
         // `m_succs` or at the switch table successor array.
-        BasicBlock*        m_succs[2];
-        BasicBlock* const* m_begin;
-        BasicBlock* const* m_end;
+        BasicBlock* m_succs[2];
+
+        // Quirk: Some BasicBlock kinds refer to their successors with BasicBlock pointers,
+        // while others use FlowEdge pointers. Eventually, every type will use FlowEdge pointers.
+        // For now, support iterating with both types.
+        union {
+            BasicBlock* const* m_begin;
+            FlowEdge* const*   m_beginEdge;
+        };
+
+        union {
+            BasicBlock* const* m_end;
+            FlowEdge* const*   m_endEdge;
+        };
+
+        bool iterateEdges;
 
     public:
         BBSuccList(const BasicBlock* block);
@@ -1879,8 +1898,8 @@ inline BBArrayIterator BBSwitchTargetList::end() const
 //
 struct BBehfDesc
 {
-    BasicBlock** bbeSuccs; // array of `BasicBlock*` pointing to BBJ_EHFINALLYRET block successors
-    unsigned     bbeCount; // size of `bbeSuccs` array
+    FlowEdge** bbeSuccs; // array of `FlowEdge*` pointing to BBJ_EHFINALLYRET block successors
+    unsigned   bbeCount; // size of `bbeSuccs` array
 
     BBehfDesc() : bbeSuccs(nullptr), bbeCount(0)
     {
@@ -1913,6 +1932,8 @@ inline BBArrayIterator BBEhfSuccList::end() const
 inline BasicBlock::BBSuccList::BBSuccList(const BasicBlock* block)
 {
     assert(block != nullptr);
+    iterateEdges = false;
+
     switch (block->bbKind)
     {
         case BBJ_THROW:
@@ -1957,14 +1978,16 @@ inline BasicBlock::BBSuccList::BBSuccList(const BasicBlock* block)
             // been computed.
             if (block->GetEhfTargets() == nullptr)
             {
-                m_begin = nullptr;
-                m_end   = nullptr;
+                m_beginEdge = nullptr;
+                m_endEdge   = nullptr;
             }
             else
             {
-                m_begin = block->GetEhfTargets()->bbeSuccs;
-                m_end   = block->GetEhfTargets()->bbeSuccs + block->GetEhfTargets()->bbeCount;
+                m_beginEdge = block->GetEhfTargets()->bbeSuccs;
+                m_endEdge   = block->GetEhfTargets()->bbeSuccs + block->GetEhfTargets()->bbeCount;
             }
+
+            iterateEdges = true;
             break;
 
         case BBJ_SWITCH:
@@ -1984,12 +2007,12 @@ inline BasicBlock::BBSuccList::BBSuccList(const BasicBlock* block)
 
 inline BBArrayIterator BasicBlock::BBSuccList::begin() const
 {
-    return BBArrayIterator(m_begin);
+    return (iterateEdges ? BBArrayIterator(m_beginEdge) : BBArrayIterator(m_begin));
 }
 
 inline BBArrayIterator BasicBlock::BBSuccList::end() const
 {
-    return BBArrayIterator(m_end);
+    return (iterateEdges ? BBArrayIterator(m_endEdge) : BBArrayIterator(m_end));
 }
 
 // We have a simpler struct, BasicBlockList, which is simply a singly-linked
@@ -2192,6 +2215,25 @@ struct FlowEdge
     }
 };
 
+// BasicBlock iterator implementations (that are required to be defined after the declaration of FlowEdge)
+
+inline BasicBlock* BBArrayIterator::operator*() const
+{
+    if (iterateEdges)
+    {
+        assert(m_edgeEntry != nullptr);
+        FlowEdge* edgeTarget = *m_edgeEntry;
+        assert(edgeTarget != nullptr);
+        assert(edgeTarget->getDestinationBlock() != nullptr);
+        return edgeTarget->getDestinationBlock();
+    }
+
+    assert(m_bbEntry != nullptr);
+    BasicBlock* bTarget = *m_bbEntry;
+    assert(bTarget != nullptr);
+    return bTarget;
+}
+
 // Pred list iterator implementations (that are required to be defined after the declaration of BasicBlock and FlowEdge)
 
 inline PredEdgeList::iterator::iterator(FlowEdge* pred) : m_pred(pred)

src/coreclr/jit/compiler.h

@@ -5883,6 +5883,8 @@ class Compiler
 
     FlowEdge* fgRemoveRefPred(BasicBlock* block, BasicBlock* blockPred);
 
+    void fgRemoveRefPred(FlowEdge* edge);
+
     FlowEdge* fgRemoveAllRefPreds(BasicBlock* block, BasicBlock* blockPred);
 
     void fgRemoveBlockAsPred(BasicBlock* block);
@@ -5893,12 +5895,16 @@ class Compiler
 
     void fgReplaceEhfSuccessor(BasicBlock* block, BasicBlock* oldSucc, BasicBlock* newSucc);
 
-    void fgRemoveEhfSuccessor(BasicBlock* block, BasicBlock* succ);
+    void fgRemoveEhfSuccessor(BasicBlock* block, const unsigned succIndex);
+    
+    void fgRemoveEhfSuccessor(FlowEdge* succEdge);
 
     void fgReplaceJumpTarget(BasicBlock* block, BasicBlock* oldTarget, BasicBlock* newTarget);
 
     void fgReplacePred(BasicBlock* block, BasicBlock* oldPred, BasicBlock* newPred);
 
+    void fgReplacePred(FlowEdge* edge, BasicBlock* const newPred);
+
     // initializingPreds is only 'true' when we are computing preds in fgLinkBasicBlocks()
     template <bool initializingPreds = false>
     FlowEdge* fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowEdge* oldEdge = nullptr);

src/coreclr/jit/compiler.hpp

@@ -657,7 +657,7 @@ BasicBlockVisit BasicBlock::VisitAllSuccs(Compiler* comp, TFunc func)
             {
                 for (unsigned i = 0; i < bbEhfTargets->bbeCount; i++)
                 {
-                    RETURN_ON_ABORT(func(bbEhfTargets->bbeSuccs[i]));
+                    RETURN_ON_ABORT(func(bbEhfTargets->bbeSuccs[i]->getDestinationBlock()));
                 }
             }
 
@@ -732,7 +732,7 @@ BasicBlockVisit BasicBlock::VisitRegularSuccs(Compiler* comp, TFunc func)
             {
                 for (unsigned i = 0; i < bbEhfTargets->bbeCount; i++)
                 {
-                    RETURN_ON_ABORT(func(bbEhfTargets->bbeSuccs[i]));
+                    RETURN_ON_ABORT(func(bbEhfTargets->bbeSuccs[i]->getDestinationBlock()));
                 }
             }