[llvm] r299706 - [llvm-extract] Add option for recursive extraction

[Keno Fischer via llvm-commits]

Author: kfischer
Date: Thu Apr  6 15:51:40 2017
New Revision: 299706

URL: http://llvm.org/viewvc/llvm-project?rev=299706&view=rev
Log:
[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

Added:
    llvm/trunk/test/tools/llvm-extract/
    llvm/trunk/test/tools/llvm-extract/recursive.ll
Modified:
    llvm/trunk/tools/llvm-extract/llvm-extract.cpp

Added: llvm/trunk/test/tools/llvm-extract/recursive.ll
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/test/tools/llvm-extract/recursive.ll?rev=299706&view=auto
==============================================================================
--- llvm/trunk/test/tools/llvm-extract/recursive.ll (added)
+++ llvm/trunk/test/tools/llvm-extract/recursive.ll Thu Apr  6 15:51:40 2017
@@ -0,0 +1,32 @@
+; RUN: llvm-extract -func=a --recursive %s -S | FileCheck --check-prefix=CHECK-AB %s
+; RUN: llvm-extract -func=a --recursive --delete %s -S | FileCheck --check-prefix=CHECK-CD %s
+; RUN: llvm-extract -func=d --recursive %s -S | FileCheck --check-prefix=CHECK-CD %s
+
+; CHECK-AB: define void @a
+; CHECK-AB: define void @b
+; CHECK-AB-NOT: define void @c
+; CHECK-AB-NOT: define void @d
+
+; CHECK-CD-NOT: define void @a
+; CHECK-CD-NOT: define void @b
+; CHECK-CD: define void @c
+; CHECK-CD: define void @d
+
+define void @a() {
+  call void @b()
+  ret void
+}
+
+define void @b() {
+  ret void
+}
+
+define void @c() {
+  call void @d()
+  ret void
+}
+
+define void @d() {
+  call void @c()
+  ret void
+}

Modified: llvm/trunk/tools/llvm-extract/llvm-extract.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/tools/llvm-extract/llvm-extract.cpp?rev=299706&r1=299705&r2=299706&view=diff
==============================================================================
--- llvm/trunk/tools/llvm-extract/llvm-extract.cpp (original)
+++ llvm/trunk/tools/llvm-extract/llvm-extract.cpp Thu Apr  6 15:51:40 2017
@@ -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 outpu
 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.