[PATCH] [llvm-extract] Add -dag option.

[Michael Spencer]

Hi chandlerc,

The -dag option informs llvm-extract to keep all global values that were
explicitly listed and that are reachable from those listed. This is useful for
extracting sub-graphs of a program.

http://reviews.llvm.org/D3884

Files:
  tools/llvm-extract/llvm-extract.cpp

Index: tools/llvm-extract/llvm-extract.cpp
===================================================================
--- tools/llvm-extract/llvm-extract.cpp
+++ tools/llvm-extract/llvm-extract.cpp
@@ -14,6 +14,7 @@
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Bitcode/BitcodeWriterPass.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IRPrintingPasses.h"
@@ -90,6 +91,63 @@
 OutputAssembly("S",
                cl::desc("Write output as LLVM assembly"), cl::Hidden);
 
+static cl::opt<bool>
+ExtractDAG("dag", cl::desc("Extract all globals and functions referenced by "
+                           "the selected global values."));
+
+SmallPtrSet<Constant *, 8> SeenConstants;
+
+void MarkUsed(Constant *C, SetVector<GlobalValue *> &GVs);
+
+/// \brief Mark the specific global value as used, and recursively mark anything
+///   that it uses as also used.
+void MarkUsed(GlobalValue *G, SetVector<GlobalValue *> &GVs) {
+  // If the global is already in the set, no need to reprocess it.
+  if (!GVs.insert(G))
+    return;
+  
+  if (GlobalVariable *GV = dyn_cast<GlobalVariable>(G)) {
+    // If this is a global variable, we must make sure to add any global values
+    // referenced by the initializer to the alive set.
+    if (GV->hasInitializer())
+      MarkUsed(GV->getInitializer(), GVs);
+  } else if (GlobalAlias *GA = dyn_cast<GlobalAlias>(G)) {
+    // The target of a global alias is needed.
+    MarkUsed(GA->getAliasee(), GVs);
+  } else {
+    // Otherwise this must be a function object.  We have to scan the body of
+    // the function looking for constants and global values which are used as
+    // operands.  Any operands of these types must be processed to ensure that
+    // any globals used will be marked as needed.
+    Function *F = cast<Function>(G);
+
+    if (F->hasPrefixData())
+      MarkUsed(F->getPrefixData(), GVs);
+
+    for (Function::iterator BB = F->begin(), E = F->end(); BB != E; ++BB)
+      for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
+        for (User::op_iterator U = I->op_begin(), E = I->op_end(); U != E; ++U)
+          if (GlobalValue *GV = dyn_cast<GlobalValue>(*U))
+            MarkUsed(GV, GVs);
+          else if (Constant *C = dyn_cast<Constant>(*U))
+            MarkUsed(C, GVs);
+  }
+}
+
+void MarkUsed(Constant *C, SetVector<GlobalValue *> &GVs) {
+  if (GlobalValue *GV = dyn_cast<GlobalValue>(C))
+    return MarkUsed(GV, GVs);
+
+  // Loop over all of the operands of the constant, adding any globals they
+  // use to the list of needed globals.
+  for (User::op_iterator I = C->op_begin(), E = C->op_end(); I != E; ++I) {
+    // If we've already processed this constant there's no need to do it again.
+    Constant *Op = dyn_cast<Constant>(*I);
+    if (Op && SeenConstants.insert(Op))
+      MarkUsed(Op, GVs);
+  }
+}
+
 int main(int argc, char **argv) {
   // Print a stack trace if we signal out.
   sys::PrintStackTraceOnErrorSignal();
@@ -213,6 +271,23 @@
     }
   }
 
+  // Add all referenced GVs.
+  if (ExtractDAG) {
+    SetVector<GlobalValue *> NewGVs;
+    for (size_t i = 0; i < GVs.size(); ++i) {
+      GlobalValue *GV = GVs[i];
+      if (GV->isMaterializable()) {
+        std::string ErrInfo;
+        if (GV->Materialize(&ErrInfo)) {
+          errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
+          return 1;
+        }
+      }
+      MarkUsed(GV, NewGVs);
+    }
+    GVs = NewGVs;
+  }
+
   // Materialize requisite global values.
   if (!DeleteFn)
     for (size_t i = 0, e = GVs.size(); i != e; ++i) {
-------------- next part --------------
A non-text attachment was scrubbed...
Name: D3884.9722.patch
Type: text/x-patch
Size: 3629 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20140522/e967f1c3/attachment.bin>