Index: include/llvm/Transforms/Utils/Local.h =================================================================== --- include/llvm/Transforms/Utils/Local.h +++ include/llvm/Transforms/Utils/Local.h @@ -117,7 +117,8 @@ /// conditions and indirectbr addresses this might make dead if /// DeleteDeadConditions is true. bool ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions = false, - const TargetLibraryInfo *TLI = nullptr); + const TargetLibraryInfo *TLI = nullptr, + DominatorTree *DT = nullptr); //===----------------------------------------------------------------------===// // Local dead code elimination. Index: lib/Transforms/Utils/Local.cpp =================================================================== --- lib/Transforms/Utils/Local.cpp +++ lib/Transforms/Utils/Local.cpp @@ -100,7 +100,8 @@ /// conditions and indirectbr addresses this might make dead if /// DeleteDeadConditions is true. bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions, - const TargetLibraryInfo *TLI) { + const TargetLibraryInfo *TLI, + DominatorTree *DT) { TerminatorInst *T = BB->getTerminator(); IRBuilder<> Builder(T); @@ -123,6 +124,9 @@ // Replace the conditional branch with an unconditional one. Builder.CreateBr(Destination); BI->eraseFromParent(); + + if (DT && OldDest != Destination) + DT->deleteEdge(BB, OldDest); return true; } @@ -159,6 +163,8 @@ TheOnlyDest = SI->case_begin()->getCaseSuccessor(); } + SmallVector Updates; + // Figure out which case it goes to. for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) { // Found case matching a constant operand? @@ -194,6 +200,7 @@ } // Remove this entry. DefaultDest->removePredecessor(SI->getParent()); + i = SI->removeCase(i); e = SI->case_end(); continue; @@ -221,19 +228,28 @@ // Insert the new branch. Builder.CreateBr(TheOnlyDest); BasicBlock *BB = SI->getParent(); + SmallPtrSet Dests; + Dests.insert(TheOnlyDest); // Remove entries from PHI nodes which we no longer branch to... for (BasicBlock *Succ : SI->successors()) { // Found case matching a constant operand? if (Succ == TheOnlyDest) TheOnlyDest = nullptr; // Don't modify the first branch to TheOnlyDest - else + else { Succ->removePredecessor(BB); + + if (DT && Dests.insert(Succ).second) + Updates.push_back({DominatorTree::Delete, BB, Succ}); + } } // Delete the old switch. Value *Cond = SI->getCondition(); SI->eraseFromParent(); + + if (DT) + DT->applyUpdates(Updates); if (DeleteDeadConditions) RecursivelyDeleteTriviallyDeadInstructions(Cond, TLI); return true; @@ -284,14 +300,26 @@ // Insert the new branch. Builder.CreateBr(TheOnlyDest); + SmallVector Updates; + SmallPtrSet Dests; + Dests.insert(TheOnlyDest); + for (unsigned i = 0, e = IBI->getNumDestinations(); i != e; ++i) { - if (IBI->getDestination(i) == TheOnlyDest) + BasicBlock *Dest = IBI->getDestination(i); + if (Dest == TheOnlyDest) TheOnlyDest = nullptr; - else - IBI->getDestination(i)->removePredecessor(IBI->getParent()); + else { + Dest->removePredecessor(IBI->getParent()); + if (DT && Dests.insert(Dest).second) + Updates.push_back({ DominatorTree::Delete, IBI->getParent(), Dest }); + } } Value *Address = IBI->getAddress(); IBI->eraseFromParent(); + + if (DT) + DT->applyUpdates(Updates); + if (DeleteDeadConditions) RecursivelyDeleteTriviallyDeadInstructions(Address, TLI); Index: unittests/Transforms/Utils/Local.cpp =================================================================== --- unittests/Transforms/Utils/Local.cpp +++ unittests/Transforms/Utils/Local.cpp @@ -212,3 +212,127 @@ EXPECT_TRUE(DT->verify()); }); } + +TEST(Local, ConstantFoldTerminator) { + LLVMContext C; + + std::unique_ptr M = parseIR( + C, + "define void @br_same_dest() {\n" + "entry:\n" + " br i1 false, label %bb0, label %bb0\n" + "bb0:\n" + " ret void\n" + "}\n" + "\n" + "define void @br_different_dest() {\n" + "entry:\n" + " br i1 true, label %bb0, label %bb1\n" + "bb0:\n" + " br label %exit\n" + "bb1:\n" + " br label %exit\n" + "exit:\n" + " ret void\n" + "}\n" + "\n" + "define void @switch_2_different_dest() {\n" + "entry:\n" + " switch i32 0, label %default [ i32 0, label %bb0 ]\n" + "default:\n" + " ret void\n" + "bb0:\n" + " ret void\n" + "}\n" + "define void @switch_2_different_dest_default() {\n" + "entry:\n" + " switch i32 1, label %default [ i32 0, label %bb0 ]\n" + "default:\n" + " ret void\n" + "bb0:\n" + " ret void\n" + "}\n" + "define void @switch_3_different_dest() {\n" + "entry:\n" + " switch i32 0, label %default [ i32 0, label %bb0\n" + " i32 1, label %bb1 ]\n" + "default:\n" + " ret void\n" + "bb0:\n" + " ret void\n" + "bb1:\n" + " ret void\n" + "}\n" + "\n" + "define void @switch_variable_2_default_dest(i32 %arg) {\n" + "entry:\n" + " switch i32 %arg, label %default [ i32 0, label %default ]\n" + "default:\n" + " ret void\n" + "}\n" + "\n" + "define void @switch_constant_2_default_dest() {\n" + "entry:\n" + " switch i32 1, label %default [ i32 0, label %default ]\n" + "default:\n" + " ret void\n" + "}\n" + "\n" + "define void @switch_constant_3_repeated_dest() {\n" + "entry:\n" + " switch i32 0, label %default [ i32 0, label %bb0\n" + " i32 1, label %bb0 ]\n" + " bb0:\n" + " ret void\n" + "default:\n" + " ret void\n" + "}\n" + "\n" + "define void @indirectbr() {\n" + "entry:\n" + " indirectbr i8* blockaddress(@indirectbr, %bb0), [label %bb0, label %bb1]\n" + "bb0:\n" + " ret void\n" + "bb1:\n" + " ret void\n" + "}\n" + "\n" + "define void @indirectbr_repeated() {\n" + "entry:\n" + " indirectbr i8* blockaddress(@indirectbr_repeated, %bb0), [label %bb0, label %bb0]\n" + "bb0:\n" + " ret void\n" + "}\n" + "\n" + "define void @indirectbr_unreachable() {\n" + "entry:\n" + " indirectbr i8* blockaddress(@indirectbr_unreachable, %bb0), [label %bb1]\n" + "bb0:\n" + " ret void\n" + "bb1:\n" + " ret void\n" + "}\n" + "\n" + ); + + auto CFAllTerminators = [&](Function &F, DominatorTree *DT) { + for (Function::iterator I = F.begin(), E = F.end(); I != E;) { + BasicBlock *BB = &*I++; + ConstantFoldTerminator(BB, true, nullptr, DT); + } + + EXPECT_TRUE(DT->verify()); + }; + + runWithDomTree(*M, "br_same_dest", CFAllTerminators); + runWithDomTree(*M, "br_different_dest", CFAllTerminators); + runWithDomTree(*M, "switch_2_different_dest", CFAllTerminators); + runWithDomTree(*M, "switch_2_different_dest_default", CFAllTerminators); + runWithDomTree(*M, "switch_3_different_dest", CFAllTerminators); + runWithDomTree(*M, "switch_variable_2_default_dest", CFAllTerminators); + runWithDomTree(*M, "switch_constant_2_default_dest", CFAllTerminators); + runWithDomTree(*M, "switch_constant_3_repeated_dest", CFAllTerminators); + runWithDomTree(*M, "indirectbr", CFAllTerminators); + runWithDomTree(*M, "indirectbr_repeated", CFAllTerminators); + runWithDomTree(*M, "indirectbr_unreachable", CFAllTerminators); +}