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