JIT: Use successor edges instead of block targets for BBJ_EHFINALLYRET

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()));
                 }
             }