[Dominators] Move IDoms out of DominatorTreeBase and put them in SNCAInfo

[kuhar]

Summary

The temporary IDoms map was used only during DomTree calculation. We can move it to SNCAInfo so that it's no longer a DominatorTreeBase member.

[• dberlin]

where else is getNodeForBlock used? Because requiring you pass a temporary thing that exists only during construction is ... weird.

Are we sure it actually belongs as part of dominatortreebase?

[kuhar]

In D34317#783154, @dberlin wrote:

where else is getNodeForBlock used? Because requiring you pass a temporary thing that exists only during construction is ... weird.

Are we sure it actually belongs as part of dominatortreebase?

It is used only in Calculate and I doesn't make much sense to keep it DominatorTreeBase.

My plan was to make a builder class that would run Semi-NCA and manage all of the temporaries. This would be probably to much to put in one review, so I tried to split it up into a series of small patches that move build-related temporaries out of the main DominatorTreeBase class.

This is the last commit that moves that data out of DTB and I wanted to structure whole thing better in a new patch -- so that DTB won't be responsible for the tree construction.

[• dberlin]

SGTM

[kuhar]

Update the diff to apply cleanly.

[grosser]

Cool, great to see this proceed.

include/llvm/Support/GenericDomTree.h

template <class NodeT> class DominatorTreeBase : public DominatorBase<NodeT> {
   }
 
   /// \brief Wipe this tree's state without releasing any resources.
   ///
   /// This is essentially a post-move helper only. It leaves the object in an
   /// assignable and destroyable state, but otherwise invalid.
   void wipe() {
     DomTreeNodes.clear();
-    IDoms.clear();
     RootNode = nullptr;
   }
 
 protected:
   using DomTreeNodeMapType =
      DenseMap<NodeT *, std::unique_ptr<DomTreeNodeBase<NodeT>>>;
   DomTreeNodeMapType DomTreeNodes;
   DomTreeNodeBase<NodeT> *RootNode;
 
   mutable bool DFSInfoValid = false;
   mutable unsigned int SlowQueries = 0;
 
-  DenseMap<NodeT *, NodeT *> IDoms;
-
   void reset() {
     DomTreeNodes.clear();
-    IDoms.clear();
     this->Roots.clear();
     RootNode = nullptr;
     DFSInfoValid = false;
     SlowQueries = 0;
   }
 
   // NewBB is split and now it has one successor. Update dominator tree to
   // reflect this change.
explicit DominatorTreeBase(bool isPostDom)
       : DominatorBase<NodeT>(isPostDom) {}
 
   DominatorTreeBase(DominatorTreeBase &&Arg)
       : DominatorBase<NodeT>(
             std::move(static_cast<DominatorBase<NodeT> &>(Arg))),
         DomTreeNodes(std::move(Arg.DomTreeNodes)),
         RootNode(std::move(Arg.RootNode)),
         DFSInfoValid(std::move(Arg.DFSInfoValid)),
-        SlowQueries(std::move(Arg.SlowQueries)),
-        IDoms(std::move(Arg.IDoms)) {
+        SlowQueries(std::move(Arg.SlowQueries)) {
     Arg.wipe();
   }
 
   DominatorTreeBase &operator=(DominatorTreeBase &&RHS) {
     DominatorBase<NodeT>::operator=(
         std::move(static_cast<DominatorBase<NodeT> &>(RHS)));
     DomTreeNodes = std::move(RHS.DomTreeNodes);
     RootNode = std::move(RHS.RootNode);
     DFSInfoValid = std::move(RHS.DFSInfoValid);
     SlowQueries = std::move(RHS.SlowQueries);
-    IDoms = std::move(RHS.IDoms);
     RHS.wipe();
     return *this;
   }
 
   DominatorTreeBase(const DominatorTreeBase &) = delete;
   DominatorTreeBase &operator=(const DominatorTreeBase &) = delete;
 
   /// compare - Return false if the other dominator tree base matches this
protected:
   // Information record used by Semi-NCA during tree construction.
   struct SemiNCAInfo {
     using NodePtr = NodeT *;
     struct InfoRec {
       unsigned DFSNum = 0;
       unsigned Parent = 0;
       unsigned Semi = 0;
       NodePtr Label = nullptr;
+      NodePtr IDom = nullptr;
     };
 
     std::vector<NodePtr> NumToNode;
     DenseMap<NodePtr, InfoRec> NodeToInfo;
+
+    NodeT *getIDom(NodeT *BB) const {
+      auto InfoIt = NodeToInfo.find(BB);
+      if (InfoIt == NodeToInfo.end()) return nullptr;
+
+      return InfoIt->second.IDom;
+    }
   };
 
   template <class GraphT>
   friend typename GraphT::NodeRef Eval(
       DominatorTreeBaseByGraphTraits<GraphT> &DT, typename GraphT::NodeRef V,
       typename DominatorTreeBaseByGraphTraits<GraphT>::SemiNCAInfo &Info,
       unsigned LastLinked);
 
   template <class GraphT>
   friend unsigned ReverseDFSPass(
       DominatorTreeBaseByGraphTraits<GraphT> &DT, typename GraphT::NodeRef V,
       typename DominatorTreeBaseByGraphTraits<GraphT>::SemiNCAInfo &Info,
       unsigned N);
 
   template <class GraphT>
   friend unsigned DFSPass(
       DominatorTreeBaseByGraphTraits<GraphT> &DT, typename GraphT::NodeRef V,
       typename DominatorTreeBaseByGraphTraits<GraphT>::SemiNCAInfo &Info,
       unsigned N);
 
   template <class FuncT, class N>
   friend void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<N>> &DT,
   FuncT &F);
 
-  DomTreeNodeBase<NodeT> *getNodeForBlock(NodeT *BB) {
+  DomTreeNodeBase<NodeT> *getNodeForBlock(NodeT *BB,
+                                          const SemiNCAInfo& SNCAInfo) {
     if (DomTreeNodeBase<NodeT> *Node = getNode(BB))
       return Node;
 
     // Haven't calculated this node yet?  Get or calculate the node for the
     // immediate dominator.
-    NodeT *IDom = getIDom(BB);
+    NodeT *IDom = SNCAInfo.getIDom(BB);
 
     assert(IDom || DomTreeNodes[nullptr]);
-    DomTreeNodeBase<NodeT> *IDomNode = getNodeForBlock(IDom);
+    DomTreeNodeBase<NodeT> *IDomNode = getNodeForBlock(IDom, SNCAInfo);
 
     // Add a new tree node for this NodeT, and link it as a child of
     // IDomNode
     return (DomTreeNodes[BB] = IDomNode->addChild(
                 llvm::make_unique<DomTreeNodeBase<NodeT>>(BB, IDomNode))).get();
   }
 
-  NodeT *getIDom(NodeT *BB) const { return IDoms.lookup(BB); }
-
   void addRoot(NodeT *BB) { this->Roots.push_back(BB); }
 
 public:
   /// updateDFSNumbers - Assign In and Out numbers to the nodes while walking
   /// dominator tree in dfs order.
   void updateDFSNumbers() const {
     if (DFSInfoValid) {
       SlowQueries = 0;

include/llvm/Support/GenericDomTreeConstruction.h

void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT,
 
   // It might be that some blocks did not get a DFS number (e.g., blocks of
   // infinite loops). In these cases an artificial exit node is required.
   MultipleRoots |= (DT.isPostDominator() && N != GraphTraits<FuncT*>::size(&F));
 
   // Initialize IDoms to spanning tree parents.
   for (unsigned i = 1; i <= N; ++i) {
     const NodePtr V = SNCA.NumToNode[i];
-    DT.IDoms[V] = SNCA.NumToNode[SNCA.NodeToInfo[V].Parent];
+    auto &VInfo = SNCA.NodeToInfo[V];
+    VInfo.IDom = SNCA.NumToNode[VInfo.Parent];
   }
 
   // Step #2: Calculate the semidominators of all vertices.
   for (unsigned i = N; i >= 2; --i) {
     NodePtr W = SNCA.NumToNode[i];
     auto &WInfo = SNCA.NodeToInfo[W];
 
     // Initialize the semi dominator to point to the parent node.
     WInfo.Semi = WInfo.Parent;
     for (const auto &N : inverse_children<NodeT>(W))
       if (SNCA.NodeToInfo.count(N)) {  // Only if this predecessor is reachable!
         unsigned SemiU = SNCA.NodeToInfo[Eval<GraphT>(DT, N, SNCA, i + 1)].Semi;
         if (SemiU < WInfo.Semi)
           WInfo.Semi = SemiU;
       }
   }
 
   // Step #3: Explicitly define the immediate dominator of each vertex.
   //          IDom[i] = NCA(SDom[i], SpanningTreeParent(i)).
   // Note that the parents were stored in IDoms and later got invalidated during
   // path compression in Eval.
   for (unsigned i = 2; i <= N; ++i) {
     const NodePtr W = SNCA.NumToNode[i];
-    const auto &WInfo = SNCA.NodeToInfo[W];
+    auto &WInfo = SNCA.NodeToInfo[W];
     const unsigned SDomNum = SNCA.NodeToInfo[SNCA.NumToNode[WInfo.Semi]].DFSNum;
-    NodePtr WIDomCandidte = DT.IDoms[W];
+    NodePtr WIDomCandidte = WInfo.IDom;
     while (SNCA.NodeToInfo[WIDomCandidte].DFSNum > SDomNum)
-      WIDomCandidte = DT.IDoms[WIDomCandidte];
+      WIDomCandidte = SNCA.NodeToInfo[WIDomCandidte].IDom;
 
-    DT.IDoms[W] = WIDomCandidte;
+    WInfo.IDom = WIDomCandidte;
   }
 
   if (DT.Roots.empty()) return;
 
   // Add a node for the root.  This node might be the actual root, if there is
   // one exit block, or it may be the virtual exit (denoted by (BasicBlock *)0)
   // which postdominates all real exits if there are multiple exit blocks, or
   // an infinite loop.
   NodePtr Root = !MultipleRoots ? DT.Roots[0] : nullptr;
 
   DT.RootNode =
       (DT.DomTreeNodes[Root] =
            llvm::make_unique<DomTreeNodeBase<NodeType>>(Root, nullptr))
           .get();
 
   // Loop over all of the reachable blocks in the function...
   for (unsigned i = 2; i <= N; ++i) {
     NodePtr W = SNCA.NumToNode[i];
 
     // Don't replace this with 'count', the insertion side effect is important
     if (DT.DomTreeNodes[W])
       continue; // Haven't calculated this node yet?
 
-    NodePtr ImmDom = DT.getIDom(W);
+    NodePtr ImmDom = SNCA.getIDom(W);
 
     assert(ImmDom || DT.DomTreeNodes[nullptr]);
 
     // Get or calculate the node for the immediate dominator
-    DomTreeNodeBase<NodeType> *IDomNode = DT.getNodeForBlock(ImmDom);
+    DomTreeNodeBase<NodeType> *IDomNode = DT.getNodeForBlock(ImmDom, SNCA);
 
     // Add a new tree node for this BasicBlock, and link it as a child of
     // IDomNode
     DT.DomTreeNodes[W] = IDomNode->addChild(
         llvm::make_unique<DomTreeNodeBase<NodeType>>(W, IDomNode));
   }
 
-  // Free temporary memory used to construct idom's
-  DT.IDoms.clear();
-
   DT.updateDFSNumbers();
 }
 }
 
 #endif