[Attributor] Fold terminators before changing instructions to unreachable

It is possible that an instruction to be changed to unreachable is
in the same block with a terminator that can be constant-folded.
In this case, as of now, the instruction will be changed to
unreachable before the terminator is folded. But, then the
whole BB becomes invalidated and so when we go ahead to fold
the terminator, we trap.

Change the order of these two.

Differential Revision: https://reviews.llvm.org/D75780

2 files changed, 28 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index f08d7f8..3aef6e9f 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -7798,15 +7798,15 @@ ChangeStatus Attributor::run() {
           ToBeChangedToUnreachableInsts.insert(&NormalDestBB->front());
         }
       }
+    for (Instruction *I : TerminatorsToFold) {
+      CGModifiedFunctions.insert(I->getFunction());
+      ConstantFoldTerminator(I->getParent());
+    }
     for (auto &V : ToBeChangedToUnreachableInsts)
       if (Instruction *I = dyn_cast_or_null<Instruction>(V)) {
         CGModifiedFunctions.insert(I->getFunction());
         changeToUnreachable(I, /* UseLLVMTrap */ false);
       }
-    for (Instruction *I : TerminatorsToFold) {
-      CGModifiedFunctions.insert(I->getFunction());
-      ConstantFoldTerminator(I->getParent());
-    }
 
     for (auto &V : ToBeDeletedInsts) {
       if (Instruction *I = dyn_cast_or_null<Instruction>(V)) {
diff --git a/llvm/test/Transforms/Attributor/undefined_behavior.ll b/llvm/test/Transforms/Attributor/undefined_behavior.ll
index cf72161..492dcd1 100644
--- a/llvm/test/Transforms/Attributor/undefined_behavior.ll
+++ b/llvm/test/Transforms/Attributor/undefined_behavior.ll
@@ -298,3 +298,27 @@ t:
 e:
   ret i32 2
 }
+
+; Note that the `load` has UB (so it will be changed to unreachable)
+; and the branch is a terminator that can be constant-folded.
+; We want to test that doing both won't cause a segfault.
+define internal i32 @callee(i1 %C, i32* %A) {
+; ATTRIBUTOR-NOT: @callee(
+;
+entry:
+  %A.0 = load i32, i32* null
+  br i1 %C, label %T, label %F
+
+T:
+  ret i32 %A.0
+
+F:
+  ret i32 1
+}
+
+define i32 @foo() {
+; ATTRIBUTOR-LABEL: @foo()
+; ATTRIBUTOR-NEXT:    ret i32 1
+  %X = call i32 @callee(i1 false, i32* null)
+  ret i32 %X
+}