[StructurizeCfg] Update dominator info.

[sepavloff]

Summary

In some cases StructurizeCfg updates root node, but dominator info
remains unchanges, it causes crash when expensive checks are enabled.

This change fixes PR27488.

[hfinkel]

include/llvm/Support/GenericDomTree.h

581 ↗

(On Diff #82492)

you should add that the old root node is assumed to be a successor of the new root node. I'm somewhat concerned that this is not really workable for post-dom trees where multiple roots are possible. Thoughts?

587 ↗

(On Diff #82492)

You're calling getNode on a null pointer? That seems wasteful, because you're doing a map lookup on a nullptr.

[• dberlin]

I'm fairly confused.
Forward dominators should not have multiple roots at once without multiple actual entry blocks (which i didn't think we supported).
Yet you seem to be creating such a thing.
That seems ... wrong.
What am i missing?

Post-dom can have multiple roots due to multiple exit blocks, but honestly, we kind of mess this up, and would likely be much better off having a single virtual root, and removing the support for multiple real roots
(which is what everyone else does)

[• dberlin]

In D28114#639445, @dberlin wrote:

I'm fairly confused.
Forward dominators should not have multiple roots at once without multiple actual entry blocks (which i didn't think we supported).
Yet you seem to be creating such a thing.
That seems ... wrong.
What am i missing?

Post-dom can have multiple roots due to multiple exit blocks, but honestly, we kind of mess this up, and would likely be much better off having a single virtual root, and removing the support for multiple real roots
(which is what everyone else does)

Oh, forget it, i see that you are swapping the root node.

This still won't work well with post-dom unless we move to a virtual root.
(and then you'd never have to swap the root node in either case)

[sepavloff]

Updated patch

As pointed out by @hfinkel and @dberlin, this code won't work with
post-dominator trees. It is possible to extend the fix to support
such trees also, but it is more complex. In the case of several exit
nodes we need to specify which root is replaced, and if the root is
virtual, this operation makes no sense. So the operation is now
confined to forward dominator trees by putting appropriate assertion
checks. There are some functions in DominatorTreeBase that impose
similar restriction.

Adding node as a new root is now implemented by a separate function
instead of reusing addNewBlock. It should prevent from accidental
root changes as a user have to explicitly specify the intent.

A small piece of unit test was added that is similar to the operation
made by StructurizeCfg pass.

[hfinkel]

Seems reasonable to me. @dberlin , what do you think?

[• dberlin]

I think this is fine for now. We can deal with the root issue itself later :)

llvm/trunk/include/llvm/Support/GenericDomTree.h

const NodeT *findNearestCommonDominator(const NodeT *A, const NodeT *B) {
     return findNearestCommonDominator(const_cast<NodeT *>(A),
                                       const_cast<NodeT *>(B));
   }
 
   //===--------------------------------------------------------------------===//
   // 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);
     assert(IDomNode && "Not immediate dominator specified for block!");
     DFSInfoValid = false;
     return (DomTreeNodes[BB] = IDomNode->addChild(
                 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.
   ///
   void changeImmediateDominator(DomTreeNodeBase<NodeT> *N,
                                 DomTreeNodeBase<NodeT> *NewIDom) {
     assert(N && NewIDom && "Cannot change null node pointers!");
     DFSInfoValid = false;
     N->setIDom(NewIDom);

llvm/trunk/lib/Transforms/Scalar/StructurizeCFG.cpp

if (LoopStart == &LoopFunc->getEntryBlock()) {
     LoopStart->setName("entry.orig");
 
     BasicBlock *NewEntry =
       BasicBlock::Create(LoopStart->getContext(),
                          "entry",
                          LoopFunc,
                          LoopStart);
     BranchInst::Create(LoopStart, NewEntry);
+    DT->setNewRoot(NewEntry);
   }
 
   // Create an extra loop end node
   LoopEnd = needPrefix(false);
   BasicBlock *Next = needPostfix(LoopEnd, ExitUseAllowed);
   LoopConds.push_back(BranchInst::Create(Next, LoopStart,
                                          BoolUndef, LoopEnd));
   addPhiValues(LoopEnd, LoopStart);

llvm/trunk/test/Transforms/StructurizeCFG/no-branch-to-entry.ll

-; 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:
 ; CHECK-NEXT: br label %entry.orig
 define void @no_branch_to_entry_undef(i32 addrspace(1)* %out) {
 entry:
   br i1 undef, label %for.end, label %for.body
 

llvm/trunk/unittests/IR/DominatorTreeTest.cpp

struct DPass : public FunctionPass {
         EXPECT_EQ(DT->getNode(BB1)->getDFSNumOut(), 4UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumIn(), 7UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumOut(), 8UL);
         EXPECT_EQ(DT->getNode(BB3)->getDFSNumIn(), 2UL);
         EXPECT_EQ(DT->getNode(BB3)->getDFSNumOut(), 3UL);
         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 {
         AU.addRequired<DominatorTreeWrapperPass>();
         AU.addRequired<PostDominatorTreeWrapperPass>();
       }
       DPass() : FunctionPass(ID) {
         initializeDPassPass(*PassRegistry::getPassRegistry());