Index: llvm/trunk/include/llvm/Support/GenericDomTree.h =================================================================== --- llvm/trunk/include/llvm/Support/GenericDomTree.h +++ llvm/trunk/include/llvm/Support/GenericDomTree.h @@ -571,9 +571,15 @@ // API to update (Post)DominatorTree information based on modifications to // the CFG... - /// addNewBlock - Add a new node to the dominator tree information. This - /// creates a new node as a child of DomBB dominator node,linking it into - /// the children list of the immediate dominator. + /// Add a new node to the dominator tree information. + /// + /// This creates a new node as a child of DomBB dominator node, linking it + /// into the children list of the immediate dominator. + /// + /// \param BB New node in CFG. + /// \param DomBB CFG node that is dominator for BB. + /// \returns New dominator tree node that represents new CFG node. + /// DomTreeNodeBase *addNewBlock(NodeT *BB, NodeT *DomBB) { assert(getNode(BB) == nullptr && "Block already in dominator tree!"); DomTreeNodeBase *IDomNode = getNode(DomBB); @@ -583,6 +589,31 @@ llvm::make_unique>(BB, IDomNode))).get(); } + /// Add a new node to the forward dominator tree and make it a new root. + /// + /// \param BB New node in CFG. + /// \returns New dominator tree node that represents new CFG node. + /// + DomTreeNodeBase *setNewRoot(NodeT *BB) { + assert(getNode(BB) == nullptr && "Block already in dominator tree!"); + assert(!this->isPostDominator() && + "Cannot change root of post-dominator tree"); + DFSInfoValid = false; + auto &Roots = DominatorBase::Roots; + DomTreeNodeBase *NewNode = (DomTreeNodes[BB] = + llvm::make_unique>(BB, nullptr)).get(); + if (Roots.empty()) { + addRoot(BB); + } else { + assert(Roots.size() == 1); + NodeT *OldRoot = Roots.front(); + DomTreeNodes[OldRoot] = + NewNode->addChild(std::move(DomTreeNodes[OldRoot])); + Roots[0] = BB; + } + return RootNode = NewNode; + } + /// changeImmediateDominator - This method is used to update the dominator /// tree information when a node's immediate dominator changes. /// Index: llvm/trunk/lib/Transforms/Scalar/StructurizeCFG.cpp =================================================================== --- llvm/trunk/lib/Transforms/Scalar/StructurizeCFG.cpp +++ llvm/trunk/lib/Transforms/Scalar/StructurizeCFG.cpp @@ -792,6 +792,7 @@ LoopFunc, LoopStart); BranchInst::Create(LoopStart, NewEntry); + DT->setNewRoot(NewEntry); } // Create an extra loop end node Index: llvm/trunk/test/Transforms/StructurizeCFG/no-branch-to-entry.ll =================================================================== --- llvm/trunk/test/Transforms/StructurizeCFG/no-branch-to-entry.ll +++ llvm/trunk/test/Transforms/StructurizeCFG/no-branch-to-entry.ll @@ -1,4 +1,4 @@ -; RUN: opt -S -o - -structurizecfg < %s | FileCheck %s +; RUN: opt -S -o - -structurizecfg -verify-dom-info < %s | FileCheck %s ; CHECK-LABEL: @no_branch_to_entry_undef( ; CHECK: entry: Index: llvm/trunk/unittests/IR/DominatorTreeTest.cpp =================================================================== --- llvm/trunk/unittests/IR/DominatorTreeTest.cpp +++ llvm/trunk/unittests/IR/DominatorTreeTest.cpp @@ -203,6 +203,16 @@ EXPECT_EQ(DT->getNode(BB4)->getDFSNumIn(), 5UL); EXPECT_EQ(DT->getNode(BB4)->getDFSNumOut(), 6UL); + // Change root node + DT->verifyDomTree(); + BasicBlock *NewEntry = BasicBlock::Create(F.getContext(), "new_entry", + &F, BB0); + BranchInst::Create(BB0, NewEntry); + EXPECT_EQ(F.begin()->getName(), NewEntry->getName()); + EXPECT_TRUE(&F.getEntryBlock() == NewEntry); + DT->setNewRoot(NewEntry); + DT->verifyDomTree(); + return false; } void getAnalysisUsage(AnalysisUsage &AU) const override {