diff --git a/src/passes/PrintCallGraph.cpp b/src/passes/PrintCallGraph.cpp index 6022c8703d7..79bc303f7ce 100644 --- a/src/passes/PrintCallGraph.cpp +++ b/src/passes/PrintCallGraph.cpp @@ -21,6 +21,7 @@ #include #include +#include #include "ir/element-utils.h" #include "ir/module-utils.h" @@ -34,72 +35,134 @@ struct PrintCallGraph : public Pass { bool modifiesBinaryenIR() override { return false; } void run(Module* module) override { - std::ostream& o = std::cout; - o << "digraph call {\n" - " rankdir = LR;\n" - " subgraph cluster_key {\n" - " node [shape=box, fontname=courier, fontsize=10];\n" - " edge [fontname=courier, fontsize=10];\n" - " label = \"Key\";\n" - " \"Import\" [style=\"filled\", fillcolor=\"turquoise\"];\n" - " \"Export\" [style=\"filled\", fillcolor=\"gray\"];\n" - " \"Indirect Target\" [style=\"filled, rounded\", " - "fillcolor=\"white\"];\n" - " \"A\" -> \"B\" [style=\"filled, rounded\", label = \"Direct " - "Call\"];\n" - " }\n\n" - " node [shape=box, fontname=courier, fontsize=10];\n"; - - // Defined functions - ModuleUtils::iterDefinedFunctions(*module, [&](Function* curr) { - std::cout << " \"" << curr->name - << "\" [style=\"filled\", fillcolor=\"white\"];\n"; - }); - - // Imported functions - ModuleUtils::iterImportedFunctions(*module, [&](Function* curr) { - o << " \"" << curr->name - << "\" [style=\"filled\", fillcolor=\"turquoise\"];\n"; - }); - - // Exports - for (auto& curr : module->exports) { - if (curr->kind == ExternalKind::Function) { - Function* func = module->getFunction(curr->value); - o << " \"" << func->name - << "\" [style=\"filled\", fillcolor=\"gray\"];\n"; - } - } + Name caller = getPassOptions().getArgumentOrDefault( + "callgraph-entry", + ""); - struct CallPrinter : public PostWalker { + struct GraphCreator : public PostWalker { Module* module; Function* currFunction; std::set visitedTargets; // Used to avoid printing duplicate edges. std::vector allIndirectTargets; - CallPrinter(Module* module) : module(module) { + std::multimap fullGraph; // First param is the caller/parent, second is the callee/child + std::stringstream callGraphStream; + + GraphCreator(Module* module, Name caller) : module(module) { // Walk function bodies. ModuleUtils::iterDefinedFunctions(*module, [&](Function* curr) { currFunction = curr; visitedTargets.clear(); walk(curr->body); }); + + constructCallgraph(caller); + } + + void constructCallgraph(Name caller) { + if (caller.toString().empty()) { + // construct callgraph for the entire module + for (const auto& graph : fullGraph) { + callGraphStream << " \"" << graph.first << "\" -> \"" << graph.second + << "\"; // call\n"; + } + } else { + // construct callgraph for a given entry point + std::set next{caller}; + visitedTargets.clear(); + + while (!next.empty()) { + std::set new_next; + for (auto curFunc : next) { + if (visitedTargets.count(curFunc)) { + continue; + } + visitedTargets.insert(curFunc); + + auto count = fullGraph.count(curFunc); + auto range = fullGraph.equal_range(curFunc); + for (auto iter = range.first; iter != range.second; ++iter) { + Name callee = iter->second; + new_next.insert(callee); + + callGraphStream << " \"" << curFunc.toString() << "\" -> \"" << callee.toString() + << "\"; // call\n"; + } + } + next = std::move(new_next); + } + } } + void visitCall(Call* curr) { auto* target = module->getFunction(curr->target); if (!visitedTargets.emplace(target->name).second) { return; } - std::cout << " \"" << currFunction->name << "\" -> \"" << target->name - << "\"; // call\n"; + fullGraph.emplace(currFunction->name, target->name); } }; - CallPrinter printer(module); + GraphCreator graphCreator(module, caller); + + std::ostream& o = std::cout; + o << "digraph call {\n" + " rankdir = LR;\n" + " subgraph cluster_key {\n" + " node [shape=box, fontname=courier, fontsize=10];\n" + " edge [fontname=courier, fontsize=10];\n" + " label = \"Key\";\n" + " \"Import\" [style=\"filled\", fillcolor=\"turquoise\"];\n" + " \"Export\" [style=\"filled\", fillcolor=\"gray\"];\n" + " \"Indirect Target\" [style=\"filled, rounded\", " + "fillcolor=\"white\"];\n" + " \"A\" -> \"B\" [style=\"filled, rounded\", label = \"Direct " + "Call\"];\n" + " }\n\n" + " node [shape=box, fontname=courier, fontsize=10];\n"; + + // Defined functions + ModuleUtils::iterDefinedFunctions(*module, [&](Function* curr) { + if (!caller.toString().empty() && + graphCreator.visitedTargets.count(curr->name) == 0) { + return; + } + std::cout << " \"" << curr->name + << "\" [style=\"filled\", fillcolor=\"white\"];\n"; + }); + + // Imported functions + ModuleUtils::iterImportedFunctions(*module, [&](Function* curr) { + if (!caller.toString().empty() && + graphCreator.visitedTargets.count(curr->name) == 0) { + return; + } + o << " \"" << curr->name + << "\" [style=\"filled\", fillcolor=\"turquoise\"];\n"; + }); + + // Exports + for (auto& curr : module->exports) { + if (curr->kind == ExternalKind::Function) { + Function* func = module->getFunction(curr->value); + if (!caller.toString().empty() && + graphCreator.visitedTargets.count(func->name) == 0) { + continue; + } + o << " \"" << func->name + << "\" [style=\"filled\", fillcolor=\"gray\"];\n"; + } + } // Indirect Targets - ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) { - auto* func = module->getFunction(name); - o << " \"" << func->name << "\" [style=\"filled, rounded\"];\n"; - }); + ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) { + auto* func = module->getFunction(name); + if (!caller.toString().empty() && + graphCreator.visitedTargets.count(func->name) == 0) { + return; + } + o << " \"" << func->name << "\" [style=\"filled, rounded\"];\n"; + }); + + o << graphCreator.callGraphStream.rdbuf(); o << "}\n"; }