From dda5fe03744d765831896b3e1f011a68581bd40c Mon Sep 17 00:00:00 2001
From: Serge Pavlov <sepavloff@gmail.com>
Date: Tue, 10 Jan 2017 02:50:47 +0000
Subject: [PATCH] [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         | 37 +++++++++++++++++--
 lib/Transforms/Scalar/StructurizeCFG.cpp      |  1 +
 .../StructurizeCFG/no-branch-to-entry.ll      |  2 +-
 unittests/IR/DominatorTreeTest.cpp            | 10 +++++
 4 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/include/llvm/Support/GenericDomTree.h b/include/llvm/Support/GenericDomTree.h
index 07a53438085..6e6ee400164 100644
--- a/include/llvm/Support/GenericDomTree.h
+++ b/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.
   ///
diff --git a/lib/Transforms/Scalar/StructurizeCFG.cpp b/lib/Transforms/Scalar/StructurizeCFG.cpp
index fa2235e8439..49ce0262c97 100644
--- a/lib/Transforms/Scalar/StructurizeCFG.cpp
+++ b/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
diff --git a/test/Transforms/StructurizeCFG/no-branch-to-entry.ll b/test/Transforms/StructurizeCFG/no-branch-to-entry.ll
index 2e22c871534..1db1060ca82 100644
--- a/test/Transforms/StructurizeCFG/no-branch-to-entry.ll
+++ b/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:
diff --git a/unittests/IR/DominatorTreeTest.cpp b/unittests/IR/DominatorTreeTest.cpp
index 6c49deb32d9..ae9c2684212 100644
--- a/unittests/IR/DominatorTreeTest.cpp
+++ b/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 {