[llvm] [CodeGen][CodeLayout] Fix segfault on access to deleted block in MBP. (PR #142357)

[Afanasyev Ivan via llvm-commits]

https://github.com/ivafanas created https://github.com/llvm/llvm-project/pull/142357

Problem 1: There is a typo which reassigns `BlockWorkList` to `EHPadWorkList` on attempt to remove `RemBB` from work lists.

Problem 2: `Chain->UnscheduledPredecessors == 0` is an incorrect way to check whether `RemBB` is enqueued or not. The root cause is a postponed deletion of `WorkList` from already scheduled blocks in `selectBestCandidateBlock`. Bug happens in the following scenario:

* `FunctionChain` is being processed with non-zero `UnscheduledPredecessors`
* Block `B'` is added to the `BlockWorkList`
* Block `B'` is chosen as the best successor (`selectBestSuccessor`) for some another block and added into `Chain`
* Block `B'` is removed by tail duplicator.

`RemovalCallback` erroneously won't erase `B'` from `BlockWorkList`, because `UnscheduledPredecessors` value of `FunctionChain` is not zero (and it is allowed to be non-zero).

Proposed solution is to always cleanup worklists on block deletion by tail duplicator.

>From 501954d48d349480a62459141ea4861df17867f0 Mon Sep 17 00:00:00 2001
From: Ivan Afanasyev <ivafanas at gmail.com>
Date: Mon, 2 Jun 2025 17:15:41 +0700
Subject: [PATCH] [CodeGen][CodeLayout] Fix segfault on access to deleted block
 in MBP.

---
 llvm/lib/CodeGen/MachineBlockPlacement.cpp | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineBlockPlacement.cpp b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
index 08fe3d47e2ff5..6d75fdee6b567 100644
--- a/llvm/lib/CodeGen/MachineBlockPlacement.cpp
+++ b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
@@ -3212,13 +3212,9 @@ bool MachineBlockPlacement::maybeTailDuplicateBlock(
     // Signal to outer function
     Removed = true;
 
-    // Conservative default.
-    bool InWorkList = true;
     // Remove from the Chain and Chain Map
     if (auto It = BlockToChain.find(RemBB); It != BlockToChain.end()) {
-      BlockChain *Chain = It->second;
-      InWorkList = Chain->UnscheduledPredecessors == 0;
-      Chain->remove(RemBB);
+      It->second->remove(RemBB);
       BlockToChain.erase(It);
     }
 
@@ -3228,11 +3224,10 @@ bool MachineBlockPlacement::maybeTailDuplicateBlock(
     }
 
     // Handle the Work Lists
-    if (InWorkList) {
-      SmallVectorImpl<MachineBasicBlock *> &RemoveList = BlockWorkList;
-      if (RemBB->isEHPad())
-        RemoveList = EHPadWorkList;
-      llvm::erase(RemoveList, RemBB);
+    if (RemBB->isEHPad()) {
+      llvm::erase(EHPadWorkList, RemBB);
+    } else {
+      llvm::erase(BlockWorkList, RemBB);
     }
 
     // Handle the filter set