From 9f17dab1f48eed788d29e4c6f045e64b0679b3a9 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Thu, 31 Dec 2020 22:46:42 +0300
Subject: [PATCH] [SimplifyCFG] Teach simplifyIndirectBr() to preserve DomTree

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp      | 7 +++++++
 llvm/test/Transforms/SimplifyCFG/indirectbr.ll | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index d081908c49e9..1fd2956fac51 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6148,10 +6148,13 @@ bool SimplifyCFGOpt::simplifyIndirectBr(IndirectBrInst *IBI) {
   bool Changed = false;
 
   // Eliminate redundant destinations.
+  std::vector<DominatorTree::UpdateType> Updates;
   SmallPtrSet<Value *, 8> Succs;
   for (unsigned i = 0, e = IBI->getNumDestinations(); i != e; ++i) {
     BasicBlock *Dest = IBI->getDestination(i);
     if (!Dest->hasAddressTaken() || !Succs.insert(Dest).second) {
+      if (!Dest->hasAddressTaken())
+        Updates.push_back({DominatorTree::Delete, BB, Dest});
       Dest->removePredecessor(BB);
       IBI->removeDestination(i);
       --i;
@@ -6160,6 +6163,10 @@ bool SimplifyCFGOpt::simplifyIndirectBr(IndirectBrInst *IBI) {
     }
   }
 
+  if (DTU)
+    DTU->applyUpdatesPermissive(Updates);
+  Updates.clear();
+
   if (IBI->getNumDestinations() == 0) {
     // If the indirectbr has no successors, change it to unreachable.
     new UnreachableInst(IBI->getContext(), IBI);
diff --git a/llvm/test/Transforms/SimplifyCFG/indirectbr.ll b/llvm/test/Transforms/SimplifyCFG/indirectbr.ll
index 67e23d2fe935..52ef6b77166c 100644
--- a/llvm/test/Transforms/SimplifyCFG/indirectbr.ll
+++ b/llvm/test/Transforms/SimplifyCFG/indirectbr.ll
@@ -1,4 +1,4 @@
-; RUN: opt -S -simplifycfg < %s | FileCheck %s
+; RUN: opt -S -simplifycfg -simplifycfg-require-and-preserve-domtree=1 < %s | FileCheck %s
 
 ; SimplifyCFG should eliminate redundant indirectbr edges.