diff --git a/llvm/include/llvm/Analysis/LazyCallGraph.h b/llvm/include/llvm/Analysis/LazyCallGraph.h index 7f9a43aba36f..df3925aba633 100644 --- a/llvm/include/llvm/Analysis/LazyCallGraph.h +++ b/llvm/include/llvm/Analysis/LazyCallGraph.h @@ -115,7 +115,7 @@ public: /// the graph. class iterator : public iterator_adaptor_base { + 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()) return *I->get(); diff --git a/llvm/lib/Analysis/LazyCallGraph.cpp b/llvm/lib/Analysis/LazyCallGraph.cpp index 05757e1e2f58..dd940a98f3b0 100644 --- a/llvm/lib/Analysis/LazyCallGraph.cpp +++ b/llvm/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()) SCCEntryNodes.push_back(F); else SCCEntryNodes.push_back(&Entry.get()->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()) - Worklist.push_back(CalleeN); + if (!Callee.isNull()) + if (Node *CalleeN = Callee.dyn_cast()) + Worklist.push_back(CalleeN); } } diff --git a/llvm/unittests/Analysis/LazyCallGraphTest.cpp b/llvm/unittests/Analysis/LazyCallGraphTest.cpp index 92b98418794b..4ef77f632e79 100644 --- a/llvm/unittests/Analysis/LazyCallGraphTest.cpp +++ b/llvm/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) {