2 files changed, 24 insertions, 7 deletions
diff --git a/lib/Analysis/LoopInfo.cpp b/lib/Analysis/LoopInfo.cpp
index cf6e20a909..20b470f3eb 100644
--- a/lib/Analysis/LoopInfo.cpp
+++ b/lib/Analysis/LoopInfo.cpp
@@ -104,7 +104,7 @@ void LoopInfo::releaseMemory() {
 }
 
 
-void LoopInfo::Calculate(const ETForest &EF) {
+void LoopInfo::Calculate(ETForest &EF) {
   BasicBlock *RootNode = EF.getRoot();
 
   for (df_iterator<BasicBlock*> NI = df_begin(RootNode),
@@ -135,7 +135,7 @@ static bool isNotAlreadyContainedIn(Loop *SubLoop, Loop *ParentLoop) {
   return isNotAlreadyContainedIn(SubLoop->getParentLoop(), ParentLoop);
 }
 
-Loop *LoopInfo::ConsiderForLoop(BasicBlock *BB, const ETForest &EF) {
+Loop *LoopInfo::ConsiderForLoop(BasicBlock *BB, ETForest &EF) {
   if (BBMap.find(BB) != BBMap.end()) return 0;   // Haven't processed this node?
 
   std::vector<BasicBlock *> TodoStack;
diff --git a/lib/VMCore/Dominators.cpp b/lib/VMCore/Dominators.cpp
index 6021b33941..951da55589 100644
--- a/lib/VMCore/Dominators.cpp
+++ b/lib/VMCore/Dominators.cpp
@@ -812,6 +812,21 @@ void ETForestBase::reset() {
   Nodes.clear();
 }
 
+void ETForestBase::updateDFSNumbers()
+{
+  int dfsnum = 0;
+  // Iterate over all nodes in depth first order.
+  for (unsigned i = 0, e = Roots.size(); i != e; ++i)
+    for (df_iterator<BasicBlock*> I = df_begin(Roots[i]),
+           E = df_end(Roots[i]); I != E; ++I) {
+      BasicBlock *BB = *I;
+      if (!getNode(BB)->hasFather())
+        getNode(BB)->assignDFSNumber(dfsnum);    
+  }
+  SlowQueries = 0;
+  DFSInfoValid = true;
+}
+
 ETNode *ETForest::getNodeForBlock(BasicBlock *BB) {
   ETNode *&BBNode = Nodes[BB];
   if (BBNode) return BBNode;
@@ -855,12 +870,14 @@ void ETForest::calculate(const ImmediateDominators &ID) {
       }
     }
 
-  int dfsnum = 0;
-  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) { 
-   if (!getNodeForBlock(I)->hasFather())
-     getNodeForBlock(I)->assignDFSNumber(dfsnum);
+  // Make sure we've got nodes around for every block
+  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
+    ETNode *&BBNode = Nodes[I];
+    if (!BBNode)
+      BBNode = new ETNode(I);
   }
-  DFSInfoValid = true;
+
+  updateDFSNumbers ();
 }
 
 //===----------------------------------------------------------------------===//