llvm/tools/llvm-extract/llvm-extract.cpp

   1 //===- llvm-extract.cpp - LLVM function extraction utility ----------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This utility changes the input module to only contain a single function,
  10 // which is primarily used for debugging transformations.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/ADT/SetVector.h"
  15 #include "llvm/ADT/SmallPtrSet.h"
  16 #include "llvm/Bitcode/BitcodeWriterPass.h"
  17 #include "llvm/IR/DataLayout.h"
  18 #include "llvm/IR/IRPrintingPasses.h"
  19 #include "llvm/IR/Instructions.h"
  20 #include "llvm/IR/LLVMContext.h"
  21 #include "llvm/IR/Module.h"
  22 #include "llvm/IRPrinter/IRPrintingPasses.h"
  23 #include "llvm/IRReader/IRReader.h"
  24 #include "llvm/Passes/PassBuilder.h"
  25 #include "llvm/Support/CommandLine.h"
  26 #include "llvm/Support/Error.h"
  27 #include "llvm/Support/FileSystem.h"
  28 #include "llvm/Support/InitLLVM.h"
  29 #include "llvm/Support/Regex.h"
  30 #include "llvm/Support/SourceMgr.h"
  31 #include "llvm/Support/SystemUtils.h"
  32 #include "llvm/Support/ToolOutputFile.h"
  33 #include "llvm/Transforms/IPO.h"
  34 #include "llvm/Transforms/IPO/BlockExtractor.h"
  35 #include "llvm/Transforms/IPO/ExtractGV.h"
  36 #include "llvm/Transforms/IPO/GlobalDCE.h"
  37 #include "llvm/Transforms/IPO/StripDeadPrototypes.h"
  38 #include "llvm/Transforms/IPO/StripSymbols.h"
  39 #include <memory>
  40 #include <utility>
  41
  42 using namespace llvm;
  43
  44 cl::OptionCategory ExtractCat("llvm-extract Options");
  45
  46 // InputFilename - The filename to read from.
  47 static cl::opt<std::string> InputFilename(cl::Positional,
  48                                           cl::desc("<input bitcode file>"),
  49                                           cl::init("-"),
  50                                           cl::value_desc("filename"));
  51
  52 static cl::opt<std::string> OutputFilename("o",
  53                                            cl::desc("Specify output filename"),
  54                                            cl::value_desc("filename"),
  55                                            cl::init("-"), cl::cat(ExtractCat));
  56
  57 static cl::opt<bool> Force("f", cl::desc("Enable binary output on terminals"),
  58                            cl::cat(ExtractCat));
  59
  60 static cl::opt<bool> DeleteFn("delete",
  61                               cl::desc("Delete specified Globals from Module"),
  62                               cl::cat(ExtractCat));
  63
  64 static cl::opt<bool> KeepConstInit("keep-const-init",
  65                               cl::desc("Keep initializers of constants"),
  66                               cl::cat(ExtractCat));
  67
  68 static cl::opt<bool>
  69     Recursive("recursive", cl::desc("Recursively extract all called functions"),
  70               cl::cat(ExtractCat));
  71
  72 // ExtractFuncs - The functions to extract from the module.
  73 static cl::list<std::string>
  74     ExtractFuncs("func", cl::desc("Specify function to extract"),
  75                  cl::value_desc("function"), cl::cat(ExtractCat));
  76
  77 // ExtractRegExpFuncs - The functions, matched via regular expression, to
  78 // extract from the module.
  79 static cl::list<std::string>
  80     ExtractRegExpFuncs("rfunc",
  81                        cl::desc("Specify function(s) to extract using a "
  82                                 "regular expression"),
  83                        cl::value_desc("rfunction"), cl::cat(ExtractCat));
  84
  85 // ExtractBlocks - The blocks to extract from the module.
  86 static cl::list<std::string> ExtractBlocks(
  87     "bb",
  88     cl::desc(
  89         "Specify <function, basic block1[;basic block2...]> pairs to extract.\n"
  90         "Each pair will create a function.\n"
  91         "If multiple basic blocks are specified in one pair,\n"
  92         "the first block in the sequence should dominate the rest.\n"
  93         "eg:\n"
  94         "  --bb=f:bb1;bb2 will extract one function with both bb1 and bb2;\n"
  95         "  --bb=f:bb1 --bb=f:bb2 will extract two functions, one with bb1, one "
  96         "with bb2."),
  97     cl::value_desc("function:bb1[;bb2...]"), cl::cat(ExtractCat));
  98
  99 // ExtractAlias - The alias to extract from the module.
 100 static cl::list<std::string>
 101     ExtractAliases("alias", cl::desc("Specify alias to extract"),
 102                    cl::value_desc("alias"), cl::cat(ExtractCat));
 103
 104 // ExtractRegExpAliases - The aliases, matched via regular expression, to
 105 // extract from the module.
 106 static cl::list<std::string>
 107     ExtractRegExpAliases("ralias",
 108                          cl::desc("Specify alias(es) to extract using a "
 109                                   "regular expression"),
 110                          cl::value_desc("ralias"), cl::cat(ExtractCat));
 111
 112 // ExtractGlobals - The globals to extract from the module.
 113 static cl::list<std::string>
 114     ExtractGlobals("glob", cl::desc("Specify global to extract"),
 115                    cl::value_desc("global"), cl::cat(ExtractCat));
 116
 117 // ExtractRegExpGlobals - The globals, matched via regular expression, to
 118 // extract from the module...
 119 static cl::list<std::string>
 120     ExtractRegExpGlobals("rglob",
 121                          cl::desc("Specify global(s) to extract using a "
 122                                   "regular expression"),
 123                          cl::value_desc("rglobal"), cl::cat(ExtractCat));
 124
 125 static cl::opt<bool> OutputAssembly("S",
 126                                     cl::desc("Write output as LLVM assembly"),
 127                                     cl::Hidden, cl::cat(ExtractCat));
 128
 129 static cl::opt<bool> PreserveBitcodeUseListOrder(
 130     "preserve-bc-uselistorder",
 131     cl::desc("Preserve use-list order when writing LLVM bitcode."),
 132     cl::init(true), cl::Hidden, cl::cat(ExtractCat));
 133
 134 static cl::opt<bool> PreserveAssemblyUseListOrder(
 135     "preserve-ll-uselistorder",
 136     cl::desc("Preserve use-list order when writing LLVM assembly."),
 137     cl::init(false), cl::Hidden, cl::cat(ExtractCat));
 138
 139 int main(int argc, char **argv) {
 140   InitLLVM X(argc, argv);
 141
 142   LLVMContext Context;
 143   cl::HideUnrelatedOptions(ExtractCat);
 144   cl::ParseCommandLineOptions(argc, argv, "llvm extractor\n");
 145
 146   // Use lazy loading, since we only care about selected global values.
 147   SMDiagnostic Err;
 148   std::unique_ptr<Module> M = getLazyIRFileModule(InputFilename, Err, Context);
 149
 150   if (!M) {
 151     Err.print(argv[0], errs());
 152     return 1;
 153   }
 154
 155   // Use SetVector to avoid duplicates.
 156   SetVector<GlobalValue *> GVs;
 157
 158   // Figure out which aliases we should extract.
 159   for (size_t i = 0, e = ExtractAliases.size(); i != e; ++i) {
 160     GlobalAlias *GA = M->getNamedAlias(ExtractAliases[i]);
 161     if (!GA) {
 162       errs() << argv[0] << ": program doesn't contain alias named '"
 163              << ExtractAliases[i] << "'!\n";
 164       return 1;
 165     }
 166     GVs.insert(GA);
 167   }
 168
 169   // Extract aliases via regular expression matching.
 170   for (size_t i = 0, e = ExtractRegExpAliases.size(); i != e; ++i) {
 171     std::string Error;
 172     Regex RegEx(ExtractRegExpAliases[i]);
 173     if (!RegEx.isValid(Error)) {
 174       errs() << argv[0] << ": '" << ExtractRegExpAliases[i] << "' "
 175         "invalid regex: " << Error;
 176     }
 177     bool match = false;
 178     for (Module::alias_iterator GA = M->alias_begin(), E = M->alias_end();
 179          GA != E; GA++) {
 180       if (RegEx.match(GA->getName())) {
 181         GVs.insert(&*GA);
 182         match = true;
 183       }
 184     }
 185     if (!match) {
 186       errs() << argv[0] << ": program doesn't contain global named '"
 187              << ExtractRegExpAliases[i] << "'!\n";
 188       return 1;
 189     }
 190   }
 191
 192   // Figure out which globals we should extract.
 193   for (size_t i = 0, e = ExtractGlobals.size(); i != e; ++i) {
 194     GlobalValue *GV = M->getNamedGlobal(ExtractGlobals[i]);
 195     if (!GV) {
 196       errs() << argv[0] << ": program doesn't contain global named '"
 197              << ExtractGlobals[i] << "'!\n";
 198       return 1;
 199     }
 200     GVs.insert(GV);
 201   }
 202
 203   // Extract globals via regular expression matching.
 204   for (size_t i = 0, e = ExtractRegExpGlobals.size(); i != e; ++i) {
 205     std::string Error;
 206     Regex RegEx(ExtractRegExpGlobals[i]);
 207     if (!RegEx.isValid(Error)) {
 208       errs() << argv[0] << ": '" << ExtractRegExpGlobals[i] << "' "
 209         "invalid regex: " << Error;
 210     }
 211     bool match = false;
 212     for (auto &GV : M->globals()) {
 213       if (RegEx.match(GV.getName())) {
 214         GVs.insert(&GV);
 215         match = true;
 216       }
 217     }
 218     if (!match) {
 219       errs() << argv[0] << ": program doesn't contain global named '"
 220              << ExtractRegExpGlobals[i] << "'!\n";
 221       return 1;
 222     }
 223   }
 224
 225   // Figure out which functions we should extract.
 226   for (size_t i = 0, e = ExtractFuncs.size(); i != e; ++i) {
 227     GlobalValue *GV = M->getFunction(ExtractFuncs[i]);
 228     if (!GV) {
 229       errs() << argv[0] << ": program doesn't contain function named '"
 230              << ExtractFuncs[i] << "'!\n";
 231       return 1;
 232     }
 233     GVs.insert(GV);
 234   }
 235   // Extract functions via regular expression matching.
 236   for (size_t i = 0, e = ExtractRegExpFuncs.size(); i != e; ++i) {
 237     std::string Error;
 238     StringRef RegExStr = ExtractRegExpFuncs[i];
 239     Regex RegEx(RegExStr);
 240     if (!RegEx.isValid(Error)) {
 241       errs() << argv[0] << ": '" << ExtractRegExpFuncs[i] << "' "
 242         "invalid regex: " << Error;
 243     }
 244     bool match = false;
 245     for (Module::iterator F = M->begin(), E = M->end(); F != E;
 246          F++) {
 247       if (RegEx.match(F->getName())) {
 248         GVs.insert(&*F);
 249         match = true;
 250       }
 251     }
 252     if (!match) {
 253       errs() << argv[0] << ": program doesn't contain global named '"
 254              << ExtractRegExpFuncs[i] << "'!\n";
 255       return 1;
 256     }
 257   }
 258
 259   // Figure out which BasicBlocks we should extract.
 260   SmallVector<std::pair<Function *, SmallVector<StringRef, 16>>, 2> BBMap;
 261   for (StringRef StrPair : ExtractBlocks) {
 262     SmallVector<StringRef, 16> BBNames;
 263     auto BBInfo = StrPair.split(':');
 264     // Get the function.
 265     Function *F = M->getFunction(BBInfo.first);
 266     if (!F) {
 267       errs() << argv[0] << ": program doesn't contain a function named '"
 268              << BBInfo.first << "'!\n";
 269       return 1;
 270     }
 271     // Add the function to the materialize list, and store the basic block names
 272     // to check after materialization.
 273     GVs.insert(F);
 274     BBInfo.second.split(BBNames, ';', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
 275     BBMap.push_back({F, std::move(BBNames)});
 276   }
 277
 278   // Use *argv instead of argv[0] to work around a wrong GCC warning.
 279   ExitOnError ExitOnErr(std::string(*argv) + ": error reading input: ");
 280
 281   if (Recursive) {
 282     std::vector<llvm::Function *> Workqueue;
 283     for (GlobalValue *GV : GVs) {
 284       if (auto *F = dyn_cast<Function>(GV)) {
 285         Workqueue.push_back(F);
 286       }
 287     }
 288     while (!Workqueue.empty()) {
 289       Function *F = &*Workqueue.back();
 290       Workqueue.pop_back();
 291       ExitOnErr(F->materialize());
 292       for (auto &BB : *F) {
 293         for (auto &I : BB) {
 294           CallBase *CB = dyn_cast<CallBase>(&I);
 295           if (!CB)
 296             continue;
 297           Function *CF = CB->getCalledFunction();
 298           if (!CF)
 299             continue;
 300           if (CF->isDeclaration() || !GVs.insert(CF))
 301             continue;
 302           Workqueue.push_back(CF);
 303         }
 304       }
 305     }
 306   }
 307
 308   auto Materialize = [&](GlobalValue &GV) { ExitOnErr(GV.materialize()); };
 309
 310   // Materialize requisite global values.
 311   if (!DeleteFn) {
 312     for (size_t i = 0, e = GVs.size(); i != e; ++i)
 313       Materialize(*GVs[i]);
 314   } else {
 315     // Deleting. Materialize every GV that's *not* in GVs.
 316     SmallPtrSet<GlobalValue *, 8> GVSet(GVs.begin(), GVs.end());
 317     for (auto &F : *M) {
 318       if (!GVSet.count(&F))
 319         Materialize(F);
 320     }
 321   }
 322
 323   {
 324     std::vector<GlobalValue *> Gvs(GVs.begin(), GVs.end());
 325     LoopAnalysisManager LAM;
 326     FunctionAnalysisManager FAM;
 327     CGSCCAnalysisManager CGAM;
 328     ModuleAnalysisManager MAM;
 329
 330     PassBuilder PB;
 331
 332     PB.registerModuleAnalyses(MAM);
 333     PB.registerCGSCCAnalyses(CGAM);
 334     PB.registerFunctionAnalyses(FAM);
 335     PB.registerLoopAnalyses(LAM);
 336     PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
 337
 338     ModulePassManager PM;
 339     PM.addPass(ExtractGVPass(Gvs, DeleteFn, KeepConstInit));
 340     PM.run(*M, MAM);
 341
 342     // Now that we have all the GVs we want, mark the module as fully
 343     // materialized.
 344     // FIXME: should the GVExtractionPass handle this?
 345     ExitOnErr(M->materializeAll());
 346   }
 347
 348   // Extract the specified basic blocks from the module and erase the existing
 349   // functions.
 350   if (!ExtractBlocks.empty()) {
 351     // Figure out which BasicBlocks we should extract.
 352     std::vector<std::vector<BasicBlock *>> GroupOfBBs;
 353     for (auto &P : BBMap) {
 354       std::vector<BasicBlock *> BBs;
 355       for (StringRef BBName : P.second) {
 356         // The function has been materialized, so add its matching basic blocks
 357         // to the block extractor list, or fail if a name is not found.
 358         auto Res = llvm::find_if(*P.first, [&](const BasicBlock &BB) {
 359           return BB.getName() == BBName;
 360         });
 361         if (Res == P.first->end()) {
 362           errs() << argv[0] << ": function " << P.first->getName()
 363                  << " doesn't contain a basic block named '" << BBName
 364                  << "'!\n";
 365           return 1;
 366         }
 367         BBs.push_back(&*Res);
 368       }
 369       GroupOfBBs.push_back(BBs);
 370     }
 371
 372     LoopAnalysisManager LAM;
 373     FunctionAnalysisManager FAM;
 374     CGSCCAnalysisManager CGAM;
 375     ModuleAnalysisManager MAM;
 376
 377     PassBuilder PB;
 378
 379     PB.registerModuleAnalyses(MAM);
 380     PB.registerCGSCCAnalyses(CGAM);
 381     PB.registerFunctionAnalyses(FAM);
 382     PB.registerLoopAnalyses(LAM);
 383     PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
 384
 385     ModulePassManager PM;
 386     PM.addPass(BlockExtractorPass(std::move(GroupOfBBs), true));
 387     PM.run(*M, MAM);
 388   }
 389
 390   // In addition to deleting all other functions, we also want to spiff it
 391   // up a little bit.  Do this now.
 392
 393   LoopAnalysisManager LAM;
 394   FunctionAnalysisManager FAM;
 395   CGSCCAnalysisManager CGAM;
 396   ModuleAnalysisManager MAM;
 397
 398   PassBuilder PB;
 399
 400   PB.registerModuleAnalyses(MAM);
 401   PB.registerCGSCCAnalyses(CGAM);
 402   PB.registerFunctionAnalyses(FAM);
 403   PB.registerLoopAnalyses(LAM);
 404   PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
 405
 406   ModulePassManager PM;
 407   if (!DeleteFn)
 408     PM.addPass(GlobalDCEPass());
 409   PM.addPass(StripDeadDebugInfoPass());
 410   PM.addPass(StripDeadPrototypesPass());
 411
 412   std::error_code EC;
 413   ToolOutputFile Out(OutputFilename, EC, sys::fs::OF_None);
 414   if (EC) {
 415     errs() << EC.message() << '\n';
 416     return 1;
 417   }
 418
 419   if (OutputAssembly)
 420     PM.addPass(PrintModulePass(Out.os(), "", PreserveAssemblyUseListOrder));
 421   else if (Force || !CheckBitcodeOutputToConsole(Out.os()))
 422     PM.addPass(BitcodeWriterPass(Out.os(), PreserveBitcodeUseListOrder));
 423
 424   PM.run(*M, MAM);
 425
 426   // Declare success.
 427   Out.keep();
 428
 429   return 0;
 430 }