From 4810051a828c4e0bc6a7d11ee5da1a40b3fc32f9 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Fri, 28 Jan 2022 12:37:28 +0100
Subject: [PATCH] [Inline][Cloning] Reliably remove unreachable blocks during
 cloning (PR53206)

The pruning cloner already tries to remove unreachable blocks. The
original cloning process will simplify instructions and constant
terminators, and only clone blocks that are reachable at that point.
However, phi nodes can only be simplified after everything has been
cloned. For that reason, additional blocks may become unreachable
after phi simplification.

The code does try to handle this as well, but only removes blocks
that don't have predecessors. It misses unreachable cycles. This
can cause issues if SEH exception handling code is part of an
unreachable cycle, as the inliner is not prepared to deal with that.

This patch instead performs an explicit scan for reachable blocks,
and drops everything else.

Fixes https://github.com/llvm/llvm-project/issues/53206.

Differential Revision: https://reviews.llvm.org/D118449
---
 llvm/lib/Transforms/Utils/CloneFunction.cpp | 53 +++++++++++----------
 llvm/test/Transforms/Inline/pr53206.ll      | 51 ++++++++++++++++++++
 2 files changed, 78 insertions(+), 26 deletions(-)
 create mode 100644 llvm/test/Transforms/Inline/pr53206.ll

diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 048e691e33cf1b..86413df664a0dc 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -694,38 +694,39 @@ void llvm::CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
       VMap[OrigV] = I;
   }
 
+  // Simplify conditional branches and switches with a constant operand. We try
+  // to prune these out when cloning, but if the simplification required
+  // looking through PHI nodes, those are only available after forming the full
+  // basic block. That may leave some here, and we still want to prune the dead
+  // code as early as possible.
+  Function::iterator Begin = cast<BasicBlock>(VMap[StartingBB])->getIterator();
+  for (BasicBlock &BB : make_range(Begin, NewFunc->end()))
+    ConstantFoldTerminator(&BB);
+
+  // Some blocks may have become unreachable as a result. Find and delete them.
+  {
+    SmallPtrSet<BasicBlock *, 16> ReachableBlocks;
+    SmallVector<BasicBlock *, 16> Worklist;
+    Worklist.push_back(&*Begin);
+    while (!Worklist.empty()) {
+      BasicBlock *BB = Worklist.pop_back_val();
+      if (ReachableBlocks.insert(BB).second)
+        append_range(Worklist, successors(BB));
+    }
+
+    SmallVector<BasicBlock *, 16> UnreachableBlocks;
+    for (BasicBlock &BB : make_range(Begin, NewFunc->end()))
+      if (!ReachableBlocks.contains(&BB))
+        UnreachableBlocks.push_back(&BB);
+    DeleteDeadBlocks(UnreachableBlocks);
+  }
+
   // Now that the inlined function body has been fully constructed, go through
   // and zap unconditional fall-through branches. This happens all the time when
   // specializing code: code specialization turns conditional branches into
   // uncond branches, and this code folds them.
-  Function::iterator Begin = cast<BasicBlock>(VMap[StartingBB])->getIterator();
   Function::iterator I = Begin;
   while (I != NewFunc->end()) {
-    // We need to simplify conditional branches and switches with a constant
-    // operand. We try to prune these out when cloning, but if the
-    // simplification required looking through PHI nodes, those are only
-    // available after forming the full basic block. That may leave some here,
-    // and we still want to prune the dead code as early as possible.
-    //
-    // Do the folding before we check if the block is dead since we want code
-    // like
-    //  bb:
-    //    br i1 undef, label %bb, label %bb
-    // to be simplified to
-    //  bb:
-    //    br label %bb
-    // before we call I->getSinglePredecessor().
-    ConstantFoldTerminator(&*I);
-
-    // Check if this block has become dead during inlining or other
-    // simplifications. Note that the first block will appear dead, as it has
-    // not yet been wired up properly.
-    if (I != Begin && (pred_empty(&*I) || I->getSinglePredecessor() == &*I)) {
-      BasicBlock *DeadBB = &*I++;
-      DeleteDeadBlock(DeadBB);
-      continue;
-    }
-
     BranchInst *BI = dyn_cast<BranchInst>(I->getTerminator());
     if (!BI || BI->isConditional()) {
       ++I;
diff --git a/llvm/test/Transforms/Inline/pr53206.ll b/llvm/test/Transforms/Inline/pr53206.ll
new file mode 100644
index 00000000000000..b7790e6e9f2c9f
--- /dev/null
+++ b/llvm/test/Transforms/Inline/pr53206.ll
@@ -0,0 +1,51 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -inline < %s | FileCheck %s
+
+; Check that the exception handling code is fully pruned, and does not
+; leave behind invalid IR.
+
+define internal void @foo() personality i32 (...)* undef {
+entry:
+  br i1 false, label %join, label %split
+
+split:
+  br label %join
+
+join:
+  %phi = phi i64 [ 1, %split ], [ 0, %entry ]
+  %cmp = icmp ugt i64 1, %phi
+  br i1 %cmp, label %invoke1, label %exit
+
+invoke1:
+  invoke void undef()
+  to label %exit unwind label %cleanup1
+
+cleanup1:
+  %pad1 = cleanuppad within none []
+  br label %cleanup1.cont
+
+cleanup1.cont:
+  br i1 undef, label %cleanupret, label %invoke2
+
+invoke2:
+  invoke void undef() [ "funclet"(token %pad1) ]
+  to label %cleanup1.cont unwind label %cleanup2
+
+cleanup2:
+  %pad2 = cleanuppad within %pad1 []
+  unreachable
+
+cleanupret:
+  unreachable
+
+exit:
+  ret void
+}
+
+define void @test() personality i32 (...)* undef {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    ret void
+;
+  call void @foo()
+  ret void
+}