2 files changed, 21 insertions, 22 deletions

diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index 2ceb78ee917..49fcd9e87d5 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -228,7 +228,6 @@ template <class NodeT> class DominatorTreeBase : public DominatorBase<NodeT> {
   /// assignable and destroyable state, but otherwise invalid.
   void wipe() {
     DomTreeNodes.clear();
-    IDoms.clear();
     RootNode = nullptr;
   }
 
@@ -241,11 +240,8 @@ protected:
   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;
@@ -320,8 +316,7 @@ public:
         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();
   }
 
@@ -332,7 +327,6 @@ public:
     RootNode = std::move(RHS.RootNode);
     DFSInfoValid = std::move(RHS.DFSInfoValid);
     SlowQueries = std::move(RHS.SlowQueries);
-    IDoms = std::move(RHS.IDoms);
     RHS.wipe();
     return *this;
   }
@@ -662,10 +656,18 @@ protected:
       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>
@@ -690,16 +692,17 @@ protected:
   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
@@ -707,8 +710,6 @@ protected:
                 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:
diff --git a/llvm/include/llvm/Support/GenericDomTreeConstruction.h b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
index 1868c043788..9fe07e0946c 100644
--- a/llvm/include/llvm/Support/GenericDomTreeConstruction.h
+++ b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
@@ -186,7 +186,8 @@ void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT,
   // 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.
@@ -211,13 +212,13 @@ void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT,
   // 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 WIDomCandidate = DT.IDoms[W];
+    NodePtr WIDomCandidate = WInfo.IDom;
     while (SNCA.NodeToInfo[WIDomCandidate].DFSNum > SDomNum)
-      WIDomCandidate = DT.IDoms[WIDomCandidate];
+      WIDomCandidate = SNCA.NodeToInfo[WIDomCandidate].IDom;
 
-    DT.IDoms[W] = WIDomCandidate;
+    WInfo.IDom = WIDomCandidate;
   }
 
   if (DT.Roots.empty()) return;
@@ -241,12 +242,12 @@ void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT,
     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
@@ -254,9 +255,6 @@ void Calculate(DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT,
         llvm::make_unique<DomTreeNodeBase<NodeType>>(W, IDomNode));
   }
 
-  // Free temporary memory used to construct idom's
-  DT.IDoms.clear();
-
   DT.updateDFSNumbers();
 }
 }