1 files changed, 38 insertions, 11 deletions
diff --git a/tools/llvm-extract/llvm-extract.cpp b/tools/llvm-extract/llvm-extract.cpp
index 319737c..b5af6bc 100644
--- a/tools/llvm-extract/llvm-extract.cpp
+++ b/tools/llvm-extract/llvm-extract.cpp
@@ -23,9 +23,10 @@
 #include "llvm/Support/IRReader.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/PrettyStackTrace.h"
-#include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/SystemUtils.h"
 #include "llvm/System/Signals.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include <memory>
 using namespace llvm;
 
@@ -102,13 +103,39 @@ int main(int argc, char **argv) {
   }
 
   // Materialize requisite global values.
-  for (size_t i = 0, e = GVs.size(); i != e; ++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;
+  if (!DeleteFn)
+    for (size_t i = 0, e = GVs.size(); i != e; ++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;
+        }
+      }
+    }
+  else {
+    // Deleting. Materialize every GV that's *not* in GVs.
+    SmallPtrSet<GlobalValue *, 8> GVSet(GVs.begin(), GVs.end());
+    for (Module::global_iterator I = M->global_begin(), E = M->global_end();
+         I != E; ++I) {
+      GlobalVariable *G = I;
+      if (!GVSet.count(G) && G->isMaterializable()) {
+        std::string ErrInfo;
+        if (G->Materialize(&ErrInfo)) {
+          errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
+          return 1;
+        }
+      }
+    }
+    for (Module::iterator I = M->begin(), E = M->end(); I != E; ++I) {
+      Function *F = I;
+      if (!GVSet.count(F) && F->isMaterializable()) {
+        std::string ErrInfo;
+        if (F->Materialize(&ErrInfo)) {
+          errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
+          return 1;
+        }
       }
     }
   }
@@ -134,9 +161,9 @@ int main(int argc, char **argv) {
   }
 
   if (OutputAssembly)
-    Passes.add(createPrintModulePass(&Out));
-  else if (Force || !CheckBitcodeOutputToConsole(Out, true))
-    Passes.add(createBitcodeWriterPass(Out));
+    Passes.add(createPrintModulePass(&Out.os()));
+  else if (Force || !CheckBitcodeOutputToConsole(Out.os(), true))
+    Passes.add(createBitcodeWriterPass(Out.os()));
 
   Passes.run(*M.get());