Index: include/llvm/Analysis/LazyCallGraph.h =================================================================== --- include/llvm/Analysis/LazyCallGraph.h +++ include/llvm/Analysis/LazyCallGraph.h @@ -202,7 +202,7 @@ friend class LazyCallGraph::RefSCC; LazyCallGraph *G; - Function &F; + Function *F; // We provide for the DFS numbering and Tarjan walk lowlink numbers to be // stored directly within the node. These are both '-1' when nodes are part @@ -229,6 +229,20 @@ /// Internal helper to remove the edge to the given function. void removeEdgeInternal(Function &ChildF); + /// Internal helper to directly replace the function with a new one. + /// + /// This is used to facilitate tranfsormations which need to replace the + /// formal Function object but directly move the body and users from one to + /// the other. + void replaceFunction(Function &NewF); + + /// Internal helper to replace an edge key with a new one. + /// + /// This should be used when the function for a particular node in the + /// graph gets replaced and we are updating all of the edges to that node + /// to use the new function as the key. + void replaceEdgeKey(Function &OldTarget, Function &NewTarget); + void clear() { Edges.clear(); EdgeIndexMap.clear(); @@ -236,7 +250,7 @@ /// Print the name of this node's function. friend raw_ostream &operator<<(raw_ostream &OS, const Node &N) { - return OS << N.F.getName(); + return OS << N.F->getName(); } /// Dump the name of this node's function to stderr. @@ -245,7 +259,7 @@ public: LazyCallGraph &getGraph() const { return *G; } - Function &getFunction() const { return F; } + Function &getFunction() const { return *F; } edge_iterator begin() const { return edge_iterator(Edges.begin(), Edges.end()); @@ -789,6 +803,17 @@ /// already existing edges. void insertTrivialRefEdge(Node &SourceN, Node &TargetN); + /// Directly replace a node's function with a new function. + /// + /// This should be used when moving the body and users of a function to + /// a new formal function object but not otherwise changing the call graph + /// structure in any way. + /// + /// It requires that the old function in the provided node have zero uses + /// and the new function must have calls and references to it establishing + /// an equivalent graph. + void replaceNodeFunction(Node &N, Function &NewF); + ///@} }; Index: lib/Analysis/LazyCallGraph.cpp =================================================================== --- lib/Analysis/LazyCallGraph.cpp +++ lib/Analysis/LazyCallGraph.cpp @@ -34,7 +34,7 @@ } LazyCallGraph::Node::Node(LazyCallGraph &G, Function &F) - : G(&G), F(F), DFSNumber(0), LowLink(0) { + : G(&G), F(&F), DFSNumber(0), LowLink(0) { DEBUG(dbgs() << " Adding functions called by '" << F.getName() << "' to the graph.\n"); @@ -108,6 +108,24 @@ EdgeIndexMap.erase(IndexMapI); } +void LazyCallGraph::Node::replaceFunction(Function &NewF) { + assert(F != &NewF && "Must not replace a function with itself!"); + F = &NewF; +} + +void LazyCallGraph::Node::replaceEdgeKey(Function &OldTarget, Function &NewTarget) { + auto IndexMapI = EdgeIndexMap.find(&OldTarget); + assert(IndexMapI != EdgeIndexMap.end() && "Old target not in the edge set!"); + int Index = IndexMapI->second; + assert(&Edges[Index].getFunction() == &NewTarget && + "Edge's target hasn't been updated yet!"); + // Ok, everything is in order, so remove the old key and insert a new one. + EdgeIndexMap.erase(IndexMapI); + bool Inserted = EdgeIndexMap.insert({&NewTarget, Index}).second; + (void)Inserted; + assert(Inserted && "Didn't actually insert a new key!"); +} + #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) LLVM_DUMP_METHOD void LazyCallGraph::Node::dump() const { dbgs() << *this << '\n'; @@ -1518,6 +1536,88 @@ handleTrivialEdgeInsertion(SourceN, TargetN); } +void LazyCallGraph::RefSCC::replaceNodeFunction(Node &N, Function &NewF) { +#ifndef NDEBUG + // Check that the RefSCC is still valid when we finish. + auto ExitVerifier = make_scope_exit([this] { verify(); }); + + assert(G->lookupRefSCC(N) == this && + "Cannot replace the function of a node outside this RefSCC."); + + assert(G->NodeMap.find(&NewF) == G->NodeMap.end() && + "Must not have already walked the new function!'"); + assert(G->EntryIndexMap.find(&NewF) == G->EntryIndexMap.end() && + "Must not have already walked the new function!'"); + + // It is important that this replacement not introduce graph changes so we + // insist that the caller has already removed every use of the original + // function and that all uses of the new function correspond to existing + // edges in the graph. The common and expected way to use this is when + // replacing the function itself in the IR without changing the call graph + // shape and just updating the analysis based on that. + Function &OldF = N.getFunction(); + assert(&OldF != &NewF && "Cannot replace a function with itself!"); + assert(OldF.use_empty() && + "Must have moved all uses from the old function to the new!"); + assert(llvm::find(G->RefSCCEntryNodes, &OldF) == G->RefSCCEntryNodes.end() && + "Must not replace a function which hasn't yet been processed into the " + "SCC graph!"); +#endif + + N.replaceFunction(NewF); + + // Update various call graph maps. + G->NodeMap.erase(&OldF); + G->NodeMap[&NewF] = &N; + auto EntryI = G->EntryIndexMap.find(&OldF); + if (EntryI != G->EntryIndexMap.end()) { + int Index = EntryI->second; + G->EntryIndexMap.erase(EntryI); + G->EntryIndexMap[&NewF] = Index; + } + + // Update all of the edges from other nodes in the call graph to this + // function. + SmallPtrSet Updated; + SmallPtrSet Visited; + SmallVector Worklist(NewF.user_begin(), NewF.user_end()); + for (User *U : Worklist) + Visited.insert(U); + while (!Worklist.empty()) { + User *U = Worklist.pop_back_val(); + // If the user is a constant, we just recurse down through its users. + if (auto *C = dyn_cast(U)) { + for (User *CU : C->users()) + if (Visited.insert(CU).second) + Worklist.push_back(CU); + continue; + } + + // We don't care what kind of instruction this is, but it must be an + // instruction at the bottom of the recursive walk through the users of the + // constants. + auto *I = cast(U); + + // Find the function and make sure we haven't processed it yet. + Function &UserF = *I->getParent()->getParent(); + if (!Updated.insert(&UserF).second) + // Already handled this node. + continue; + + Node *UserN = G->lookup(UserF); + if (!UserN) + // No need to update a function that hasn't been put into a node. + continue; + + assert(UserN->lookup(OldF) && + "No edge from a caller to the replacement node!"); + assert(UserN->lookup(OldF)->getNode() == &N && + "Edge from a caller doesn't point to the node!"); + + UserN->replaceEdgeKey(OldF, NewF); + } +} + void LazyCallGraph::insertEdge(Node &SourceN, Function &Target, Edge::Kind EK) { assert(SCCMap.empty() && DFSStack.empty() && "This method cannot be called after SCCs have been formed!"); Index: unittests/Analysis/LazyCallGraphTest.cpp =================================================================== --- unittests/Analysis/LazyCallGraphTest.cpp +++ unittests/Analysis/LazyCallGraphTest.cpp @@ -2134,4 +2134,83 @@ EXPECT_TRUE(GRC.isParentOf(FRC)); } +TEST(LazyCallGraphTest, ReplaceNodeFunction) { + LLVMContext Context; + // A graph with several different kinds of edges pointing at a particular + // function. + std::unique_ptr M = parseAssembly( + Context, "define void @a(i8** %ptr) {\n" + "entry:\n" + " store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n" + " ret void\n" + "}\n" + "define void @b(i8** %ptr) {\n" + "entry:\n" + " store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n" + " store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n" + " call void @d(i8** %ptr)" + " ret void\n" + "}\n" + "define void @c(i8** %ptr) {\n" + "entry:\n" + " call void @d(i8** %ptr)" + " call void @d(i8** %ptr)" + " store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n" + " ret void\n" + "}\n" + "define void @d(i8** %ptr) {\n" + "entry:\n" + " store i8* bitcast (void(i8**)* @b to i8*), i8** %ptr\n" + " call void @c(i8** %ptr)" + " call void @d(i8** %ptr)" + " store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n" + " ret void\n" + "}\n"); + LazyCallGraph CG(*M); + + // Force the graph to be fully expanded. + auto I = CG.postorder_ref_scc_begin(); + LazyCallGraph::RefSCC &RC1 = *I++; + LazyCallGraph::RefSCC &RC2 = *I++; + EXPECT_EQ(CG.postorder_ref_scc_end(), I); + + ASSERT_EQ(2u, RC1.size()); + LazyCallGraph::SCC &C1 = RC1[0]; + LazyCallGraph::SCC &C2 = RC1[1]; + + LazyCallGraph::Node &AN = *CG.lookup(lookupFunction(*M, "a")); + LazyCallGraph::Node &BN = *CG.lookup(lookupFunction(*M, "b")); + LazyCallGraph::Node &CN = *CG.lookup(lookupFunction(*M, "c")); + LazyCallGraph::Node &DN = *CG.lookup(lookupFunction(*M, "d")); + EXPECT_EQ(&C1, CG.lookupSCC(DN)); + EXPECT_EQ(&C1, CG.lookupSCC(CN)); + EXPECT_EQ(&C2, CG.lookupSCC(BN)); + EXPECT_EQ(&RC1, CG.lookupRefSCC(DN)); + EXPECT_EQ(&RC1, CG.lookupRefSCC(CN)); + EXPECT_EQ(&RC1, CG.lookupRefSCC(BN)); + EXPECT_EQ(&RC2, CG.lookupRefSCC(AN)); + + // Now we need to build a new function 'e' with the same signature as 'd'. + Function &D = DN.getFunction(); + Function &E = *Function::Create(D.getFunctionType(), D.getLinkage(), "e"); + D.getParent()->getFunctionList().insert(D.getIterator(), &E); + + // Change each use of 'd' to use 'e'. This is particularly easy as they have + // the same type. + D.replaceAllUsesWith(&E); + + // Splice the body of the old function into the new one. + E.getBasicBlockList().splice(E.begin(), D.getBasicBlockList()); + // And fix up the one argument. + D.arg_begin()->replaceAllUsesWith(&*E.arg_begin()); + E.arg_begin()->takeName(&*D.arg_begin()); + + // Now replace the function in the graph. + RC1.replaceNodeFunction(DN, E); + + EXPECT_EQ(&E, &DN.getFunction()); + EXPECT_EQ(&DN, CN[E].getNode()); + EXPECT_EQ(&DN, BN[E].getNode()); +} + }