Fix a bug that prevented llvm-extract -delete from working.
[oota-llvm.git] / lib / Transforms / IPO / ExtractGV.cpp
1 //===-- ExtractGV.cpp - Global Value extraction pass ----------------------===//
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 pass extracts global values
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "llvm/Instructions.h"
15 #include "llvm/Module.h"
16 #include "llvm/Pass.h"
17 #include "llvm/Constants.h"
18 #include "llvm/Transforms/IPO.h"
19 #include "llvm/Support/Compiler.h"
20 #include <algorithm>
21 using namespace llvm;
22
23 namespace {
24   /// @brief A pass to extract specific functions and their dependencies.
25   class VISIBILITY_HIDDEN GVExtractorPass : public ModulePass {
26     std::vector<GlobalValue*> Named;
27     bool deleteStuff;
28     bool reLink;
29   public:
30     static char ID; // Pass identification, replacement for typeid
31
32     /// FunctionExtractorPass - If deleteFn is true, this pass deletes as the
33     /// specified function. Otherwise, it deletes as much of the module as
34     /// possible, except for the function specified.
35     ///
36     explicit GVExtractorPass(std::vector<GlobalValue*>& GVs, bool deleteS = true,
37                              bool relinkCallees = false)
38       : ModulePass(&ID), Named(GVs), deleteStuff(deleteS),
39         reLink(relinkCallees) {}
40
41     bool runOnModule(Module &M) {
42       if (Named.size() == 0) {
43         return false;  // Nothing to extract
44       }
45       
46       if (deleteStuff)
47         return deleteGV();
48       M.setModuleInlineAsm("");
49       return isolateGV(M);
50     }
51
52     bool deleteGV() {
53       for (std::vector<GlobalValue*>::iterator GI = Named.begin(), 
54              GE = Named.end(); GI != GE; ++GI) {
55         if (Function* NamedFunc = dyn_cast<Function>(*GI)) {
56          // If we're in relinking mode, set linkage of all internal callees to
57          // external. This will allow us extract function, and then - link
58          // everything together
59          if (reLink) {
60            for (Function::iterator B = NamedFunc->begin(), BE = NamedFunc->end();
61                 B != BE; ++B) {
62              for (BasicBlock::iterator I = B->begin(), E = B->end();
63                   I != E; ++I) {
64                if (CallInst* callInst = dyn_cast<CallInst>(&*I)) {
65                  Function* Callee = callInst->getCalledFunction();
66                  if (Callee && Callee->hasInternalLinkage())
67                    Callee->setLinkage(GlobalValue::ExternalLinkage);
68                }
69              }
70            }
71          }
72          
73          NamedFunc->setLinkage(GlobalValue::ExternalLinkage);
74          NamedFunc->deleteBody();
75          assert(NamedFunc->isDeclaration() && "This didn't make the function external!");
76        } else {
77           if (!(*GI)->isDeclaration()) {
78             cast<GlobalVariable>(*GI)->setInitializer(0);  //clear the initializer
79             (*GI)->setLinkage(GlobalValue::ExternalLinkage);
80           }
81         }
82       }
83       return true;
84     }
85
86     bool isolateGV(Module &M) {
87       // Mark all globals internal
88       for (Module::global_iterator I = M.global_begin(), E = M.global_end(); I != E; ++I)
89         if (!I->isDeclaration()) {
90           I->setLinkage(GlobalValue::InternalLinkage);
91         }
92       for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I)
93         if (!I->isDeclaration()) {
94           I->setLinkage(GlobalValue::InternalLinkage);
95         }
96
97       // Make sure our result is globally accessible...
98       // by putting them in the used array
99       {
100         std::vector<Constant *> AUGs;
101         const Type *SBP= PointerType::getUnqual(Type::Int8Ty);
102         for (std::vector<GlobalValue*>::iterator GI = Named.begin(), 
103                GE = Named.end(); GI != GE; ++GI) {
104           (*GI)->setLinkage(GlobalValue::ExternalLinkage);
105           AUGs.push_back(ConstantExpr::getBitCast(*GI, SBP));
106         }
107         ArrayType *AT = ArrayType::get(SBP, AUGs.size());
108         Constant *Init = ConstantArray::get(AT, AUGs);
109         GlobalValue *gv = new GlobalVariable(AT, false, 
110                                              GlobalValue::AppendingLinkage, 
111                                              Init, "llvm.used", &M);
112         gv->setSection("llvm.metadata");
113       }
114
115       // All of the functions may be used by global variables or the named
116       // globals.  Loop through them and create a new, external functions that
117       // can be "used", instead of ones with bodies.
118       std::vector<Function*> NewFunctions;
119
120       Function *Last = --M.end();  // Figure out where the last real fn is.
121
122       for (Module::iterator I = M.begin(); ; ++I) {
123         if (std::find(Named.begin(), Named.end(), &*I) == Named.end()) {
124           Function *New = Function::Create(I->getFunctionType(),
125                                            GlobalValue::ExternalLinkage);
126           New->copyAttributesFrom(I);
127
128           // If it's not the named function, delete the body of the function
129           I->dropAllReferences();
130
131           M.getFunctionList().push_back(New);
132           NewFunctions.push_back(New);
133           New->takeName(I);
134         }
135
136         if (&*I == Last) break;  // Stop after processing the last function
137       }
138
139       // Now that we have replacements all set up, loop through the module,
140       // deleting the old functions, replacing them with the newly created
141       // functions.
142       if (!NewFunctions.empty()) {
143         unsigned FuncNum = 0;
144         Module::iterator I = M.begin();
145         do {
146           if (std::find(Named.begin(), Named.end(), &*I) == Named.end()) {
147             // Make everything that uses the old function use the new dummy fn
148             I->replaceAllUsesWith(NewFunctions[FuncNum++]);
149
150             Function *Old = I;
151             ++I;  // Move the iterator to the new function
152
153             // Delete the old function!
154             M.getFunctionList().erase(Old);
155
156           } else {
157             ++I;  // Skip the function we are extracting
158           }
159         } while (&*I != NewFunctions[0]);
160       }
161
162       return true;
163     }
164   };
165
166   char GVExtractorPass::ID = 0;
167 }
168
169 ModulePass *llvm::createGVExtractionPass(std::vector<GlobalValue*>& GVs, 
170                                          bool deleteFn, bool relinkCallees) {
171   return new GVExtractorPass(GVs, deleteFn, relinkCallees);
172 }