[llvm-extract] Add option for recursive extraction

Summary:
Particularly, with --delete, this can be very useful for testing
new optimizations on some hotspots, without having to run it on the whole
application. E.g. as such:
```
llvm-extract app.bc --recursive --rfunc .*hotspot.* > hotspot.bc
llvm-extract app.bc --recursive --delete --rfunc .*hotspot.* > residual.bc
llc -filetype=obj residual.bc > residual.o
llc -filetype=obj hotspot.bc > hotspot.o
cc -o app residual.o hotspot.o
```

Reviewed By: davide
Differential Revision: https://reviews.llvm.org/D31722

llvm-svn: 299706

1 files changed, 34 insertions, 1 deletions

diff --git a/llvm/tools/llvm-extract/llvm-extract.cpp b/llvm/tools/llvm-extract/llvm-extract.cpp
index aa1eda2f094..d868db7f78a 100644
--- a/llvm/tools/llvm-extract/llvm-extract.cpp
+++ b/llvm/tools/llvm-extract/llvm-extract.cpp
@@ -17,10 +17,11 @@
 #include "llvm/Bitcode/BitcodeWriterPass.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IRPrintingPasses.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
-#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
@@ -50,6 +51,10 @@ Force("f", cl::desc("Enable binary output on terminals"));
 static cl::opt<bool>
 DeleteFn("delete", cl::desc("Delete specified Globals from Module"));
 
+static cl::opt<bool>
+    Recursive("recursive",
+              cl::desc("Recursively extract all called functions"));
+
 // ExtractFuncs - The functions to extract from the module.
 static cl::list<std::string>
 ExtractFuncs("func", cl::desc("Specify function to extract"),
@@ -226,6 +231,34 @@ int main(int argc, char **argv) {
   // Use *argv instead of argv[0] to work around a wrong GCC warning.
   ExitOnError ExitOnErr(std::string(*argv) + ": error reading input: ");
 
+  if (Recursive) {
+    std::vector<llvm::Function *> Workqueue;
+    for (GlobalValue *GV : GVs) {
+      if (auto *F = dyn_cast<Function>(GV)) {
+        Workqueue.push_back(F);
+      }
+    }
+    while (!Workqueue.empty()) {
+      Function *F = &*Workqueue.back();
+      Workqueue.pop_back();
+      ExitOnErr(F->materialize());
+      for (auto &BB : *F) {
+        for (auto &I : BB) {
+          auto *CI = dyn_cast<CallInst>(&I);
+          if (!CI)
+            continue;
+          Function *CF = CI->getCalledFunction();
+          if (!CF)
+            continue;
+          if (CF->isDeclaration() || GVs.count(CF))
+            continue;
+          GVs.insert(CF);
+          Workqueue.push_back(CF);
+        }
+      }
+    }
+  }
+
   auto Materialize = [&](GlobalValue &GV) { ExitOnErr(GV.materialize()); };
 
   // Materialize requisite global values.