tools/llvm-extract/llvm-extract.cpp

   1 //===- llvm-extract.cpp - LLVM function extraction utility ----------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This utility changes the input module to only contain a single function,
  11 // which is primarily used for debugging transformations.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "llvm/IR/LLVMContext.h"
  16 #include "llvm/ADT/SetVector.h"
  17 #include "llvm/ADT/SmallPtrSet.h"
  18 #include "llvm/Assembly/PrintModulePass.h"
  19 #include "llvm/Bitcode/ReaderWriter.h"
  20 #include "llvm/IR/DataLayout.h"
  21 #include "llvm/IR/Module.h"
  22 #include "llvm/IRReader/IRReader.h"
  23 #include "llvm/PassManager.h"
  24 #include "llvm/Support/CommandLine.h"
  25 #include "llvm/Support/ManagedStatic.h"
  26 #include "llvm/Support/PrettyStackTrace.h"
  27 #include "llvm/Support/Regex.h"
  28 #include "llvm/Support/Signals.h"
  29 #include "llvm/Support/SourceMgr.h"
  30 #include "llvm/Support/SystemUtils.h"
  31 #include "llvm/Support/ToolOutputFile.h"
  32 #include "llvm/Transforms/IPO.h"
  33 #include <memory>
  34 using namespace llvm;
  35
  36 // InputFilename - The filename to read from.
  37 static cl::opt<std::string>
  38 InputFilename(cl::Positional, cl::desc("<input bitcode file>"),
  39               cl::init("-"), cl::value_desc("filename"));
  40
  41 static cl::opt<std::string>
  42 OutputFilename("o", cl::desc("Specify output filename"),
  43                cl::value_desc("filename"), cl::init("-"));
  44
  45 static cl::opt<bool>
  46 Force("f", cl::desc("Enable binary output on terminals"));
  47
  48 static cl::opt<bool>
  49 DeleteFn("delete", cl::desc("Delete specified Globals from Module"));
  50
  51 // ExtractFuncs - The functions to extract from the module.
  52 static cl::list<std::string>
  53 ExtractFuncs("func", cl::desc("Specify function to extract"),
  54              cl::ZeroOrMore, cl::value_desc("function"));
  55
  56 // ExtractRegExpFuncs - The functions, matched via regular expression, to
  57 // extract from the module.
  58 static cl::list<std::string>
  59 ExtractRegExpFuncs("rfunc", cl::desc("Specify function(s) to extract using a "
  60                                      "regular expression"),
  61                    cl::ZeroOrMore, cl::value_desc("rfunction"));
  62
  63 // ExtractAlias - The alias to extract from the module.
  64 static cl::list<std::string>
  65 ExtractAliases("alias", cl::desc("Specify alias to extract"),
  66                cl::ZeroOrMore, cl::value_desc("alias"));
  67
  68
  69 // ExtractRegExpAliases - The aliases, matched via regular expression, to
  70 // extract from the module.
  71 static cl::list<std::string>
  72 ExtractRegExpAliases("ralias", cl::desc("Specify alias(es) to extract using a "
  73                                         "regular expression"),
  74                      cl::ZeroOrMore, cl::value_desc("ralias"));
  75
  76 // ExtractGlobals - The globals to extract from the module.
  77 static cl::list<std::string>
  78 ExtractGlobals("glob", cl::desc("Specify global to extract"),
  79                cl::ZeroOrMore, cl::value_desc("global"));
  80
  81 // ExtractRegExpGlobals - The globals, matched via regular expression, to
  82 // extract from the module...
  83 static cl::list<std::string>
  84 ExtractRegExpGlobals("rglob", cl::desc("Specify global(s) to extract using a "
  85                                        "regular expression"),
  86                      cl::ZeroOrMore, cl::value_desc("rglobal"));
  87
  88 static cl::opt<bool>
  89 OutputAssembly("S",
  90                cl::desc("Write output as LLVM assembly"), cl::Hidden);
  91
  92 int main(int argc, char **argv) {
  93   // Print a stack trace if we signal out.
  94   sys::PrintStackTraceOnErrorSignal();
  95   PrettyStackTraceProgram X(argc, argv);
  96
  97   LLVMContext &Context = getGlobalContext();
  98   llvm_shutdown_obj Y;  // Call llvm_shutdown() on exit.
  99
 100   // Initialize PassManager for -time-passes support.
 101   initializePassManager();
 102
 103   cl::ParseCommandLineOptions(argc, argv, "llvm extractor\n");
 104
 105   // Use lazy loading, since we only care about selected global values.
 106   SMDiagnostic Err;
 107   OwningPtr<Module> M;
 108   M.reset(getLazyIRFileModule(InputFilename, Err, Context));
 109
 110   if (M.get() == 0) {
 111     Err.print(argv[0], errs());
 112     return 1;
 113   }
 114
 115   // Use SetVector to avoid duplicates.
 116   SetVector<GlobalValue *> GVs;
 117
 118   // Figure out which aliases we should extract.
 119   for (size_t i = 0, e = ExtractAliases.size(); i != e; ++i) {
 120     GlobalAlias *GA = M->getNamedAlias(ExtractAliases[i]);
 121     if (!GA) {
 122       errs() << argv[0] << ": program doesn't contain alias named '"
 123              << ExtractAliases[i] << "'!\n";
 124       return 1;
 125     }
 126     GVs.insert(GA);
 127   }
 128
 129   // Extract aliases via regular expression matching.
 130   for (size_t i = 0, e = ExtractRegExpAliases.size(); i != e; ++i) {
 131     std::string Error;
 132     Regex RegEx(ExtractRegExpAliases[i]);
 133     if (!RegEx.isValid(Error)) {
 134       errs() << argv[0] << ": '" << ExtractRegExpAliases[i] << "' "
 135         "invalid regex: " << Error;
 136     }
 137     bool match = false;
 138     for (Module::alias_iterator GA = M->alias_begin(), E = M->alias_end();
 139          GA != E; GA++) {
 140       if (RegEx.match(GA->getName())) {
 141         GVs.insert(&*GA);
 142         match = true;
 143       }
 144     }
 145     if (!match) {
 146       errs() << argv[0] << ": program doesn't contain global named '"
 147              << ExtractRegExpAliases[i] << "'!\n";
 148       return 1;
 149     }
 150   }
 151
 152   // Figure out which globals we should extract.
 153   for (size_t i = 0, e = ExtractGlobals.size(); i != e; ++i) {
 154     GlobalValue *GV = M->getNamedGlobal(ExtractGlobals[i]);
 155     if (!GV) {
 156       errs() << argv[0] << ": program doesn't contain global named '"
 157              << ExtractGlobals[i] << "'!\n";
 158       return 1;
 159     }
 160     GVs.insert(GV);
 161   }
 162
 163   // Extract globals via regular expression matching.
 164   for (size_t i = 0, e = ExtractRegExpGlobals.size(); i != e; ++i) {
 165     std::string Error;
 166     Regex RegEx(ExtractRegExpGlobals[i]);
 167     if (!RegEx.isValid(Error)) {
 168       errs() << argv[0] << ": '" << ExtractRegExpGlobals[i] << "' "
 169         "invalid regex: " << Error;
 170     }
 171     bool match = false;
 172     for (Module::global_iterator GV = M->global_begin(),
 173            E = M->global_end(); GV != E; GV++) {
 174       if (RegEx.match(GV->getName())) {
 175         GVs.insert(&*GV);
 176         match = true;
 177       }
 178     }
 179     if (!match) {
 180       errs() << argv[0] << ": program doesn't contain global named '"
 181              << ExtractRegExpGlobals[i] << "'!\n";
 182       return 1;
 183     }
 184   }
 185
 186   // Figure out which functions we should extract.
 187   for (size_t i = 0, e = ExtractFuncs.size(); i != e; ++i) {
 188     GlobalValue *GV = M->getFunction(ExtractFuncs[i]);
 189     if (!GV) {
 190       errs() << argv[0] << ": program doesn't contain function named '"
 191              << ExtractFuncs[i] << "'!\n";
 192       return 1;
 193     }
 194     GVs.insert(GV);
 195   }
 196   // Extract functions via regular expression matching.
 197   for (size_t i = 0, e = ExtractRegExpFuncs.size(); i != e; ++i) {
 198     std::string Error;
 199     StringRef RegExStr = ExtractRegExpFuncs[i];
 200     Regex RegEx(RegExStr);
 201     if (!RegEx.isValid(Error)) {
 202       errs() << argv[0] << ": '" << ExtractRegExpFuncs[i] << "' "
 203         "invalid regex: " << Error;
 204     }
 205     bool match = false;
 206     for (Module::iterator F = M->begin(), E = M->end(); F != E;
 207          F++) {
 208       if (RegEx.match(F->getName())) {
 209         GVs.insert(&*F);
 210         match = true;
 211       }
 212     }
 213     if (!match) {
 214       errs() << argv[0] << ": program doesn't contain global named '"
 215              << ExtractRegExpFuncs[i] << "'!\n";
 216       return 1;
 217     }
 218   }
 219
 220   // Materialize requisite global values.
 221   if (!DeleteFn)
 222     for (size_t i = 0, e = GVs.size(); i != e; ++i) {
 223       GlobalValue *GV = GVs[i];
 224       if (GV->isMaterializable()) {
 225         std::string ErrInfo;
 226         if (GV->Materialize(&ErrInfo)) {
 227           errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
 228           return 1;
 229         }
 230       }
 231     }
 232   else {
 233     // Deleting. Materialize every GV that's *not* in GVs.
 234     SmallPtrSet<GlobalValue *, 8> GVSet(GVs.begin(), GVs.end());
 235     for (Module::global_iterator I = M->global_begin(), E = M->global_end();
 236          I != E; ++I) {
 237       GlobalVariable *G = I;
 238       if (!GVSet.count(G) && G->isMaterializable()) {
 239         std::string ErrInfo;
 240         if (G->Materialize(&ErrInfo)) {
 241           errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
 242           return 1;
 243         }
 244       }
 245     }
 246     for (Module::iterator I = M->begin(), E = M->end(); I != E; ++I) {
 247       Function *F = I;
 248       if (!GVSet.count(F) && F->isMaterializable()) {
 249         std::string ErrInfo;
 250         if (F->Materialize(&ErrInfo)) {
 251           errs() << argv[0] << ": error reading input: " << ErrInfo << "\n";
 252           return 1;
 253         }
 254       }
 255     }
 256   }
 257
 258   // In addition to deleting all other functions, we also want to spiff it
 259   // up a little bit.  Do this now.
 260   PassManager Passes;
 261   Passes.add(new DataLayout(M.get())); // Use correct DataLayout
 262
 263   std::vector<GlobalValue*> Gvs(GVs.begin(), GVs.end());
 264
 265   Passes.add(createGVExtractionPass(Gvs, DeleteFn));
 266   if (!DeleteFn)
 267     Passes.add(createGlobalDCEPass());           // Delete unreachable globals
 268   Passes.add(createStripDeadDebugInfoPass());    // Remove dead debug info
 269   Passes.add(createStripDeadPrototypesPass());   // Remove dead func decls
 270
 271   std::string ErrorInfo;
 272   tool_output_file Out(OutputFilename.c_str(), ErrorInfo, sys::fs::F_Binary);
 273   if (!ErrorInfo.empty()) {
 274     errs() << ErrorInfo << '\n';
 275     return 1;
 276   }
 277
 278   if (OutputAssembly)
 279     Passes.add(createPrintModulePass(&Out.os()));
 280   else if (Force || !CheckBitcodeOutputToConsole(Out.os(), true))
 281     Passes.add(createBitcodeWriterPass(Out.os()));
 282
 283   Passes.run(*M.get());
 284
 285   // Declare success.
 286   Out.keep();
 287
 288   return 0;
 289 }