From 76e74d939509a71ae2416744a5659fc6a30fb2c5 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Sat, 19 Dec 2020 15:38:30 +0300
Subject: [PATCH] [SimplifyCFG] Teach removeEmptyCleanup() to preserve DomTree

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp          | 22 +++++++++++++++++-----
 llvm/test/Transforms/Coroutines/coro-heap-elide.ll |  2 +-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 1e4ad30..0e7a8a7 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -4205,7 +4205,7 @@ bool SimplifyCFGOpt::simplifySingleResume(ResumeInst *RI) {
   return true;
 }
 
-static bool removeEmptyCleanup(CleanupReturnInst *RI) {
+static bool removeEmptyCleanup(CleanupReturnInst *RI, DomTreeUpdater *DTU) {
   // If this is a trivial cleanup pad that executes no instructions, it can be
   // eliminated.  If the cleanup pad continues to the caller, any predecessor
   // that is an EH pad will be updated to continue to the caller and any
@@ -4312,20 +4312,32 @@ static bool removeEmptyCleanup(CleanupReturnInst *RI) {
     }
   }
 
+  std::vector<DominatorTree::UpdateType> Updates;
+
   for (pred_iterator PI = pred_begin(BB), PE = pred_end(BB); PI != PE;) {
     // The iterator must be updated here because we are removing this pred.
     BasicBlock *PredBB = *PI++;
     if (UnwindDest == nullptr) {
-      removeUnwindEdge(PredBB);
+      if (DTU)
+        DTU->applyUpdatesPermissive(Updates);
+      Updates.clear();
+      removeUnwindEdge(PredBB, DTU);
       ++NumInvokes;
     } else {
       Instruction *TI = PredBB->getTerminator();
       TI->replaceUsesOfWith(BB, UnwindDest);
+      Updates.push_back({DominatorTree::Delete, PredBB, BB});
+      Updates.push_back({DominatorTree::Insert, PredBB, UnwindDest});
     }
   }
 
-  // The cleanup pad is now unreachable.  Zap it.
-  BB->eraseFromParent();
+  if (DTU) {
+    DTU->applyUpdatesPermissive(Updates);
+    DTU->deleteBB(BB);
+  } else
+    // The cleanup pad is now unreachable.  Zap it.
+    BB->eraseFromParent();
+
   return true;
 }
 
@@ -4372,7 +4384,7 @@ bool SimplifyCFGOpt::simplifyCleanupReturn(CleanupReturnInst *RI) {
   if (mergeCleanupPad(RI))
     return true;
 
-  if (removeEmptyCleanup(RI))
+  if (removeEmptyCleanup(RI, DTU))
     return true;
 
   return false;
diff --git a/llvm/test/Transforms/Coroutines/coro-heap-elide.ll b/llvm/test/Transforms/Coroutines/coro-heap-elide.ll
index a5315e0..4f59fa1 100644
--- a/llvm/test/Transforms/Coroutines/coro-heap-elide.ll
+++ b/llvm/test/Transforms/Coroutines/coro-heap-elide.ll
@@ -1,7 +1,7 @@
 ; Tests that the dynamic allocation and deallocation of the coroutine frame is
 ; elided and any tail calls referencing the coroutine frame has the tail 
 ; call attribute removed.
-; RUN: opt < %s -S -inline -coro-elide -instsimplify -simplifycfg | FileCheck %s
+; RUN: opt < %s -S -inline -coro-elide -instsimplify -simplifycfg -simplifycfg-require-and-preserve-domtree=1 | FileCheck %s
 ; RUN: opt < %s -S \
 ; RUN:   -passes='cgscc(inline,function(coro-elide,instsimplify,simplify-cfg))' \
 ; RUN:   -aa-pipeline='basic-aa' | FileCheck %s
-- 
2.7.4