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