From 8fbba30d42eb58ad17cec03b90f5c79390444ea9 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i@tianshilei.me>
Date: Wed, 22 Jan 2025 21:54:29 -0500
Subject: [PATCH] [StructurizeCFG] Fix a crash caused by not updating
 `Predicates` properly

A very simple IR could crash the pass at `assert(BB != Parent);`.

```
define void @foo() {
entry:
  br i1 false, label %cond.true, label %cond.false

cond.true:                                        ; preds = %entry
  br label %cond.end

cond.false:                                       ; preds = %entry
  br label %cond.end

cond.end:                                         ; preds = %cond.false, %cond.true
  ret void
}
```

TBH I'm not sure if my fix is right. I don't know why we don't update
`Predicates` after we replace a conditional branch with a branch, given the one
not taken will no longer be a sucessor of the current BB. There are many test
cases need to be updated but I'd like to update them after I fix this properly.
Any help would be appreciated.
---
 llvm/lib/Transforms/Scalar/StructurizeCFG.cpp |  4 +++
 .../simple-structurizecfg-crash.ll            | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 llvm/test/Transforms/StructurizeCFG/simple-structurizecfg-crash.ll

diff --git a/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp b/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
index b1f742b838f2a..eaeb6ca4f3e04 100644
--- a/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
+++ b/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
@@ -971,6 +971,10 @@ void StructurizeCFG::changeExit(RegionNode *Node, BasicBlock *NewExit,
     SubRegion->replaceExit(NewExit);
   } else {
     BasicBlock *BB = Node->getNodeAs<BasicBlock>();
+    for (BasicBlock *Succ : successors(BB)) {
+      if (Succ != NewExit)
+        Predicates[Succ].erase(BB);
+    }
     killTerminator(BB);
     BranchInst *Br = BranchInst::Create(NewExit, BB);
     Br->setDebugLoc(TermDL[BB]);
diff --git a/llvm/test/Transforms/StructurizeCFG/simple-structurizecfg-crash.ll b/llvm/test/Transforms/StructurizeCFG/simple-structurizecfg-crash.ll
new file mode 100644
index 0000000000000..15983c448588b
--- /dev/null
+++ b/llvm/test/Transforms/StructurizeCFG/simple-structurizecfg-crash.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=structurizecfg %s -o - | FileCheck %s
+
+define void @foo() {
+; CHECK-LABEL: define void @foo() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[COND_FALSE:.*]]
+; CHECK:       [[COND_TRUE:.*]]:
+; CHECK-NEXT:    br label %[[COND_END:.*]]
+; CHECK:       [[COND_FALSE]]:
+; CHECK-NEXT:    br i1 false, label %[[COND_TRUE]], label %[[COND_END]]
+; CHECK:       [[COND_END]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 false, label %cond.true, label %cond.false
+
+cond.true:                                        ; preds = %entry
+  br label %cond.end
+
+cond.false:                                       ; preds = %entry
+  br label %cond.end
+
+cond.end:                                         ; preds = %cond.false, %cond.true
+  ret void
+}