From 1f6fa0f45ad50b9de70df7ddc046cc8ff8c76369 Mon Sep 17 00:00:00 2001
From: Keno Fischer <keno@alumni.harvard.edu>
Date: Thu, 6 Apr 2017 20:51:40 +0000
Subject: [PATCH] [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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@299706 91177308-0d34-0410-b5e6-96231b3b80d8
---
 test/tools/llvm-extract/recursive.ll | 32 +++++++++++++++++++++++++
 tools/llvm-extract/llvm-extract.cpp  | 35 +++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 test/tools/llvm-extract/recursive.ll

diff --git a/test/tools/llvm-extract/recursive.ll b/test/tools/llvm-extract/recursive.ll
new file mode 100644
index 00000000000..54813dba796
--- /dev/null
+++ b/test/tools/llvm-extract/recursive.ll
@@ -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
+}
diff --git a/tools/llvm-extract/llvm-extract.cpp b/tools/llvm-extract/llvm-extract.cpp
index aa1eda2f094..d868db7f78a 100644
--- a/tools/llvm-extract/llvm-extract.cpp
+++ b/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.