[StructurizeCfg] Update dominator info.

In some cases StructurizeCfg updates root node, but dominator info
remains unchanges, it causes crash when expensive checks are enabled.
To cope with this problem a new method was added to DominatorTreeBase
that allows adding new root nodes, it is called in StructurizeCfg to
put dominator tree in sync.

This change fixes PR27488.

Differential Revision: https://reviews.llvm.org/D28114

llvm-svn: 291530

include/llvm/Support/GenericDomTree.h

@@ -571,9 +571,15 @@ public:
   // 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<NodeT> *addNewBlock(NodeT *BB, NodeT *DomBB) {
     assert(getNode(BB) == nullptr && "Block already in dominator tree!");
     DomTreeNodeBase<NodeT> *IDomNode = getNode(DomBB);
@@ -583,6 +589,31 @@ public:
                 llvm::make_unique<DomTreeNodeBase<NodeT>>(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<NodeT> *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<NodeT>::Roots;
+    DomTreeNodeBase<NodeT> *NewNode = (DomTreeNodes[BB] =
+      llvm::make_unique<DomTreeNodeBase<NodeT>>(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.
   ///

lib/Transforms/Scalar/StructurizeCFG.cpp

@@ -792,6 +792,7 @@ void StructurizeCFG::handleLoops(bool ExitUseAllowed,
                          LoopFunc,
                          LoopStart);
     BranchInst::Create(LoopStart, NewEntry);
+    DT->setNewRoot(NewEntry);
   }
 
   // Create an extra loop end node

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:

unittests/IR/DominatorTreeTest.cpp

@@ -203,6 +203,16 @@ namespace llvm {
         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 {