3 files changed, 40 insertions, 7 deletions

diff --git a/include/llvm/Analysis/LazyCallGraph.h b/include/llvm/Analysis/LazyCallGraph.h
index 7f9a43aba3..df3925aba6 100644
--- a/include/llvm/Analysis/LazyCallGraph.h
+++ b/include/llvm/Analysis/LazyCallGraph.h
@@ -115,7 +115,7 @@ public:
   /// the graph.
   class iterator
       : public iterator_adaptor_base<iterator, NodeVectorImplT::iterator,
-                                     std::random_access_iterator_tag, Node> {
+                                     std::bidirectional_iterator_tag, Node> {
     friend class LazyCallGraph;
     friend class LazyCallGraph::Node;
 
@@ -124,11 +124,30 @@ public:
 
     // Build the iterator for a specific position in a node list.
     iterator(LazyCallGraph &G, NodeVectorImplT::iterator NI)
-        : iterator_adaptor_base(NI), G(&G) {}
+        : iterator_adaptor_base(NI), G(&G) {
+      while (I->isNull())
+        ++I;
+    }
 
   public:
     iterator() {}
 
+    using iterator_adaptor_base::operator++;
+    iterator &operator++() {
+      do {
+        ++I;
+      } while (I->isNull());
+      return *this;
+    }
+
+    using iterator_adaptor_base::operator--;
+    iterator &operator--() {
+      do {
+        --I;
+      } while (I->isNull());
+      return *this;
+    }
+
     reference operator*() const {
       if (I->is<Node *>())
         return *I->get<Node *>();
diff --git a/lib/Analysis/LazyCallGraph.cpp b/lib/Analysis/LazyCallGraph.cpp
index 05757e1e2f..dd940a98f3 100644
--- a/lib/Analysis/LazyCallGraph.cpp
+++ b/lib/Analysis/LazyCallGraph.cpp
@@ -88,7 +88,7 @@ void LazyCallGraph::Node::removeEdgeInternal(Function &Callee) {
   assert(IndexMapI != CalleeIndexMap.end() &&
          "Callee not in the callee set for this caller?");
 
-  Callees.erase(Callees.begin() + IndexMapI->second);
+  Callees[IndexMapI->second] = nullptr;
   CalleeIndexMap.erase(IndexMapI);
 }
 
@@ -115,11 +115,14 @@ LazyCallGraph::LazyCallGraph(Module &M) : NextDFSNumber(0) {
                   "entry set.\n");
   findCallees(Worklist, Visited, EntryNodes, EntryIndexMap);
 
-  for (auto &Entry : EntryNodes)
+  for (auto &Entry : EntryNodes) {
+    assert(!Entry.isNull() &&
+           "We can't have removed edges before we finish the constructor!");
     if (Function *F = Entry.dyn_cast<Function *>())
       SCCEntryNodes.push_back(F);
     else
       SCCEntryNodes.push_back(&Entry.get<Node *>()->getFunction());
+  }
 }
 
 LazyCallGraph::LazyCallGraph(LazyCallGraph &&G)
@@ -391,8 +394,9 @@ void LazyCallGraph::updateGraphPtrs() {
       Node *N = Worklist.pop_back_val();
       N->G = this;
       for (auto &Callee : N->Callees)
-        if (Node *CalleeN = Callee.dyn_cast<Node *>())
-          Worklist.push_back(CalleeN);
+        if (!Callee.isNull())
+          if (Node *CalleeN = Callee.dyn_cast<Node *>())
+            Worklist.push_back(CalleeN);
     }
   }
 
diff --git a/unittests/Analysis/LazyCallGraphTest.cpp b/unittests/Analysis/LazyCallGraphTest.cpp
index 92b9841879..4ef77f632e 100644
--- a/unittests/Analysis/LazyCallGraphTest.cpp
+++ b/unittests/Analysis/LazyCallGraphTest.cpp
@@ -290,7 +290,17 @@ TEST(LazyCallGraphTest, BasicGraphMutation) {
   CG.insertEdge(C, C.getFunction());
   EXPECT_EQ(2, std::distance(C.begin(), C.end()));
   EXPECT_EQ(&B, &*C.begin());
-  EXPECT_EQ(&C, &*(C.begin() + 1));
+  EXPECT_EQ(&C, &*std::next(C.begin()));
+
+  CG.removeEdge(C, B.getFunction());
+  EXPECT_EQ(1, std::distance(C.begin(), C.end()));
+  EXPECT_EQ(&C, &*C.begin());
+
+  CG.removeEdge(C, C.getFunction());
+  EXPECT_EQ(0, std::distance(C.begin(), C.end()));
+
+  CG.removeEdge(B, C.getFunction());
+  EXPECT_EQ(0, std::distance(B.begin(), B.end()));
 }
 
 TEST(LazyCallGraphTest, MultiArmSCC) {