JIT: update fgReplaceJumpTarget to maintain pred lists (dotnet#81246)

This is used some early phases; make it pred list aware.

Contributes to dotnet#80193.

src/coreclr/jit/fgbasic.cpp

@@ -540,14 +540,12 @@ void Compiler::fgReplaceSwitchJumpTarget(BasicBlock* blockSwitch, BasicBlock* ne
 //
 // Notes:
 // 1. Only branches are changed: BBJ_ALWAYS, the non-fallthrough path of BBJ_COND, BBJ_SWITCH, etc.
-//    We ignore other block types.
+//    We assert for other jump kinds.
 // 2. All branch targets found are updated. If there are multiple ways for a block
 //    to reach 'oldTarget' (e.g., multiple arms of a switch), all of them are changed.
-// 3. The predecessor lists are not changed.
+// 3. The predecessor lists are updated, if they've been built.
 // 4. If any switch table entry was updated, the switch table "unique successor" cache is invalidated.
 //
-// This function is most useful early, before the full predecessor lists have been computed.
-//
 void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* newTarget, BasicBlock* oldTarget)
 {
     assert(block != nullptr);
@@ -559,18 +557,18 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* newTarget, Bas
         case BBJ_ALWAYS:
         case BBJ_EHCATCHRET:
         case BBJ_EHFILTERRET:
-        case BBJ_LEAVE: // This function will be called before import, so we still have BBJ_LEAVE
+        case BBJ_LEAVE: // This function can be called before import, so we still have BBJ_LEAVE
 
             if (block->bbJumpDest == oldTarget)
             {
                 block->bbJumpDest = newTarget;
-            }
-            break;
 
-        case BBJ_NONE:
-        case BBJ_EHFINALLYRET:
-        case BBJ_THROW:
-        case BBJ_RETURN:
+                if (fgComputePredsDone)
+                {
+                    fgRemoveRefPred(oldTarget, block);
+                    fgAddRefPred(newTarget, block);
+                }
+            }
             break;
 
         case BBJ_SWITCH:
@@ -585,6 +583,12 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* newTarget, Bas
                 {
                     jumpTab[i] = newTarget;
                     changed    = true;
+
+                    if (fgComputePredsDone)
+                    {
+                        fgRemoveRefPred(oldTarget, block);
+                        fgAddRefPred(newTarget, block);
+                    }
                 }
             }
 
@@ -596,7 +600,7 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* newTarget, Bas
         }
 
         default:
-            assert(!"Block doesn't have a valid bbJumpKind!!!!");
+            assert(!"Block doesn't have a jump target!");
             unreached();
             break;
     }

src/coreclr/jit/fgopt.cpp

@@ -2010,38 +2010,36 @@ void Compiler::fgCompactBlocks(BasicBlock* block, BasicBlock* bNext)
     /* both or none must have an exception handler */
     noway_assert(block->hasTryIndex() == bNext->hasTryIndex());
 
-#ifdef DEBUG
-    if (verbose)
-    {
-        printf("\nCompacting blocks " FMT_BB " and " FMT_BB ":\n", block->bbNum, bNext->bbNum);
-    }
-#endif
+    JITDUMP("\nCompacting " FMT_BB " into " FMT_BB ":\n", bNext->bbNum, block->bbNum);
+    fgRemoveRefPred(bNext, block);
 
-    if (bNext->countOfInEdges() > 1)
+    if (bNext->countOfInEdges() > 0)
     {
-        JITDUMP("Second block has multiple incoming edges\n");
+        JITDUMP("Second block has %u other incoming edges\n", bNext->countOfInEdges());
         assert(block->isEmpty());
 
+        // `block` can no longer be a loop pre-header (if it was before).
+        //
+        block->bbFlags &= ~BBF_LOOP_PREHEADER;
+
+        // Retarget all the other edges incident on bNext. Do this
+        // in two passes as we can't both walk and modify the pred list.
+        //
+        ArrayStack<BasicBlock*> preds(getAllocator(CMK_BasicBlock), bNext->countOfInEdges());
         for (BasicBlock* const predBlock : bNext->PredBlocks())
         {
+            preds.Push(predBlock);
+        }
+        while (preds.Height() > 0)
+        {
+            BasicBlock* const predBlock = preds.Pop();
             fgReplaceJumpTarget(predBlock, block, bNext);
-
-            if (predBlock != block)
-            {
-                fgAddRefPred(block, predBlock);
-            }
         }
-        bNext->bbPreds = nullptr;
-
-        // `block` can no longer be a loop pre-header (if it was before).
-        block->bbFlags &= ~BBF_LOOP_PREHEADER;
-    }
-    else
-    {
-        noway_assert(bNext->bbPreds->flNext == nullptr);
-        noway_assert(bNext->bbPreds->getBlock() == block);
     }
 
+    assert(bNext->countOfInEdges() == 0);
+    assert(bNext->bbPreds == nullptr);
+
     /* Start compacting - move all the statements in the second block to the first block */
 
     // First move any phi definitions of the second block after the phi defs of the first.

src/coreclr/jit/jiteh.cpp

@@ -2435,10 +2435,14 @@ bool Compiler::fgNormalizeEHCase2()
                                 fgAddCheapPred(newTryStart, predBlock);
                                 fgRemoveCheapPred(insertBeforeBlk, predBlock);
 
-                                // Now change the branch. If it was a BBJ_NONE fall-through to the top block, this will
-                                // do nothing. Since cheap preds contains dups (for switch duplicates), we will call
+                                // Now change pred branches.
+                                //
+                                // Since cheap preds contains dups (for switch duplicates), we will call
                                 // this once per dup.
-                                fgReplaceJumpTarget(predBlock, newTryStart, insertBeforeBlk);
+                                if (predBlock->bbJumpKind != BBJ_NONE)
+                                {
+                                    fgReplaceJumpTarget(predBlock, newTryStart, insertBeforeBlk);
+                                }
 
                                 // Need to adjust ref counts here since we're retargeting edges.
                                 newTryStart->bbRefs++;

src/coreclr/jit/redundantbranchopts.cpp

@@ -1553,33 +1553,15 @@ bool Compiler::optJumpThreadCore(JumpThreadInfo& jti)
                     " implies predicate true; we can safely redirect flow to be " FMT_BB " -> " FMT_BB "\n",
                     predBlock->bbNum, jti.m_block->bbNum, predBlock->bbNum, jti.m_trueTarget->bbNum);
 
-            if (predBlock->bbJumpKind == BBJ_SWITCH)
-            {
-                fgReplaceSwitchJumpTarget(predBlock, jti.m_trueTarget, jti.m_block);
-            }
-            else
-            {
-                fgRemoveRefPred(jti.m_block, predBlock);
-                fgReplaceJumpTarget(predBlock, jti.m_trueTarget, jti.m_block);
-                fgAddRefPred(jti.m_trueTarget, predBlock);
-            }
+            fgReplaceJumpTarget(predBlock, jti.m_trueTarget, jti.m_block);
         }
         else
         {
             JITDUMP("Jump flow from pred " FMT_BB " -> " FMT_BB
                     " implies predicate false; we can safely redirect flow to be " FMT_BB " -> " FMT_BB "\n",
                     predBlock->bbNum, jti.m_block->bbNum, predBlock->bbNum, jti.m_falseTarget->bbNum);
 
-            if (predBlock->bbJumpKind == BBJ_SWITCH)
-            {
-                fgReplaceSwitchJumpTarget(predBlock, jti.m_falseTarget, jti.m_block);
-            }
-            else
-            {
-                fgRemoveRefPred(jti.m_block, predBlock);
-                fgReplaceJumpTarget(predBlock, jti.m_falseTarget, jti.m_block);
-                fgAddRefPred(jti.m_falseTarget, predBlock);
-            }
+            fgReplaceJumpTarget(predBlock, jti.m_falseTarget, jti.m_block);
         }
     }