diff --git a/include/llvm/Analysis/LazyCallGraph.h b/include/llvm/Analysis/LazyCallGraph.h index 95d3ed61e92..e4069819524 100644 --- a/include/llvm/Analysis/LazyCallGraph.h +++ b/include/llvm/Analysis/LazyCallGraph.h @@ -163,6 +163,9 @@ public: /// CalleeIndexMap. Node(LazyCallGraph &G, Function &F); + /// \brief Internal helper to insert a callee. + void insertEdgeInternal(Function &Callee); + /// \brief Internal helper to remove a callee from this node. void removeEdgeInternal(Function &Callee); @@ -371,6 +374,14 @@ public: /// Once you begin manipulating a call graph's SCCs, you must perform all /// mutation of the graph via the SCC methods. + /// \brief Update the call graph after inserting a new edge. + void insertEdge(Node &Caller, Function &Callee); + + /// \brief Update the call graph after inserting a new edge. + void insertEdge(Function &Caller, Function &Callee) { + return insertEdge(get(Caller), Callee); + } + /// \brief Update the call graph after deleting an edge. void removeEdge(Node &Caller, Function &Callee); diff --git a/lib/Analysis/LazyCallGraph.cpp b/lib/Analysis/LazyCallGraph.cpp index b593069e664..05757e1e2f5 100644 --- a/lib/Analysis/LazyCallGraph.cpp +++ b/lib/Analysis/LazyCallGraph.cpp @@ -75,6 +75,14 @@ LazyCallGraph::Node::Node(LazyCallGraph &G, Function &F) findCallees(Worklist, Visited, Callees, CalleeIndexMap); } +void LazyCallGraph::Node::insertEdgeInternal(Function &Callee) { + CalleeIndexMap.insert(std::make_pair(&Callee, Callees.size())); + if (Node *N = G->lookup(Callee)) + Callees.push_back(N); + else + Callees.push_back(&Callee); +} + void LazyCallGraph::Node::removeEdgeInternal(Function &Callee) { auto IndexMapI = CalleeIndexMap.find(&Callee); assert(IndexMapI != CalleeIndexMap.end() && @@ -353,6 +361,13 @@ LazyCallGraph::SCC::removeIntraSCCEdge(Node &CallerN, return ResultSCCs; } +void LazyCallGraph::insertEdge(Node &CallerN, Function &Callee) { + assert(SCCMap.empty() && DFSStack.empty() && + "This method cannot be called after SCCs have been formed!"); + + return CallerN.insertEdgeInternal(Callee); +} + void LazyCallGraph::removeEdge(Node &CallerN, Function &Callee) { assert(SCCMap.empty() && DFSStack.empty() && "This method cannot be called after SCCs have been formed!"); diff --git a/unittests/Analysis/LazyCallGraphTest.cpp b/unittests/Analysis/LazyCallGraphTest.cpp index 3fbd3ec20f1..92b98418794 100644 --- a/unittests/Analysis/LazyCallGraphTest.cpp +++ b/unittests/Analysis/LazyCallGraphTest.cpp @@ -255,6 +255,44 @@ static Function &lookupFunction(Module &M, StringRef Name) { report_fatal_error("Couldn't find function!"); } +TEST(LazyCallGraphTest, BasicGraphMutation) { + std::unique_ptr M = parseAssembly( + "define void @a() {\n" + "entry:\n" + " call void @b()\n" + " call void @c()\n" + " ret void\n" + "}\n" + "define void @b() {\n" + "entry:\n" + " ret void\n" + "}\n" + "define void @c() {\n" + "entry:\n" + " ret void\n" + "}\n"); + LazyCallGraph CG(*M); + + LazyCallGraph::Node &A = CG.get(lookupFunction(*M, "a")); + LazyCallGraph::Node &B = CG.get(lookupFunction(*M, "b")); + EXPECT_EQ(2, std::distance(A.begin(), A.end())); + EXPECT_EQ(0, std::distance(B.begin(), B.end())); + + CG.insertEdge(B, lookupFunction(*M, "c")); + EXPECT_EQ(1, std::distance(B.begin(), B.end())); + LazyCallGraph::Node &C = *B.begin(); + EXPECT_EQ(0, std::distance(C.begin(), C.end())); + + CG.insertEdge(C, B.getFunction()); + EXPECT_EQ(1, std::distance(C.begin(), C.end())); + EXPECT_EQ(&B, &*C.begin()); + + 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)); +} + TEST(LazyCallGraphTest, MultiArmSCC) { // Two interlocking cycles. The really useful thing about this SCC is that it // will require Tarjan's DFS to backtrack and finish processing all of the