Reoganize Function/BasicBlock creation
[frida/frida.git] / src / disassembler / llvm / LLVMDisassembler.cxx
index b6f8dae75ae2d6cc33b2586bc2b45df68b899eec..773a5189a27a33d35085da9f698864970b3ac9ec 100644 (file)
@@ -5,6 +5,7 @@
 
 #include <stack>
 #include <algorithm>
+#include <cassert>
 
 using namespace llvm;
 using namespace llvm::object;
@@ -161,19 +162,11 @@ void LLVMDisassembler<ELFT>::start() {
 }
 
 template <typename ELFT>
-LLVMDisassembler<ELFT>::~LLVMDisassembler() {
-       // std::for_each(functions.begin(), functions.end(),
-       //               [](std::pair<uint64_t,LLVMFunction*> it) {
-       //                    delete it.second;
-       //               });
-       // std::for_each(blocks.begin(), blocks.end(),
-       //               [](std::pair<uint64_t, LLVMBasicBlock*> it) {
-       //                    delete it.second;
-       //               });
-}
+LLVMDisassembler<ELFT>::~LLVMDisassembler() {}
 
 template <typename ELFT>
 Function* LLVMDisassembler<ELFT>::disassembleFunctionAt(uint64_t address, const std::string& name) {
+       Function * function;
        SectionRef text_section = sections[".text"];
        uint64_t base_address, size;
        text_section.getAddress(base_address);
@@ -184,21 +177,20 @@ Function* LLVMDisassembler<ELFT>::disassembleFunctionAt(uint64_t address, const
                return NULL;
        }
 
-       if (functions.find(address) != functions.end()) {
-               return functions[address];
-       }
+       if (NULL == (function = manager->getFunction(address))) {
 
-       Function * function;
-       if (name == "") {
-               std::stringstream s;
-               s << "<Unnamed 0x" << std::hex << address << ">";
-               function = new Function(s.str(), address, manager);
-       } else {
-               function = new Function(name, address, manager);
+               if (name == "") {
+                       std::stringstream s;
+                       s << "<Unnamed 0x" << std::hex << address << ">";
+                       function = manager->newFunction(address);
+                       function->setName(s.str());
+               } else {
+                       function = manager->newFunction(address);
+                       function->setName(name);
+               }
+               disassembleFunction(function);
+               manager->finishFunction(function);
        }
-       functions.insert(std::make_pair(address, function));
-
-       disassembleFunction(function);
 
        return function;
 }
@@ -206,6 +198,11 @@ Function* LLVMDisassembler<ELFT>::disassembleFunctionAt(uint64_t address, const
 template <typename ELFT>
 void LLVMDisassembler<ELFT>::disassembleFunction(Function* function) {
        std::stack<BasicBlock*> remaining_blocks;
+       /* TODO:
+        * Do all blocks get added properly? We should take care to remove
+        * the other ones at the end of the function!
+        */
+       std::map<uint64_t, BasicBlock*> new_blocks;
        SectionRef text_section = sections[".text"];
        StringRef bytes;
        text_section.getContents(bytes);
@@ -213,16 +210,17 @@ void LLVMDisassembler<ELFT>::disassembleFunction(Function* function) {
 
        LOG4CXX_DEBUG(logger, "Handling function " << function->getName());
 
-       BasicBlock * block = new BasicBlock(function->getStartAddress(), this, manager);
+       BasicBlock * block = manager->newBasicBlock(function->getStartAddress());
        remaining_blocks.push(block);
-       blocks.insert(std::make_pair(block->getStartAddress(), block));
+       new_blocks.insert(std::make_pair(block->getStartAddress(), block));
        function->addBasicBlock(block);
 
        while (remaining_blocks.size()) {
                BasicBlock * current_block = remaining_blocks.top();
                remaining_blocks.pop();
 
-               LOG4CXX_DEBUG(logger, "Handling Block starting at " << std::hex << current_block->getStartAddress());
+               LOG4CXX_DEBUG(logger, "Handling Block starting at " << std::hex
+                             << current_block->getStartAddress());
 
                uint64_t inst_size;
                uint64_t base_address;
@@ -241,31 +239,34 @@ void LLVMDisassembler<ELFT>::disassembleFunction(Function* function) {
                                        jmptarget += base_address;
                                        if (!MIA->isIndirectBranch(inst)) {
                                                if (MIA->isCall(inst)) {
-                                                       if (functions.find(jmptarget) == functions.end()) {
+                                                       if (NULL == manager->getFunction(jmptarget))
                                                                disassembleFunctionAt(jmptarget);
-                                                       }
                                                } else {
                                                        current_block->setNextBlock(0, jmptarget);
-                                                       if (blocks.find(jmptarget) == blocks.end()) {
-                                                               BasicBlock * block = new BasicBlock(jmptarget, this, manager);
-                                                               blocks.insert(std::make_pair(block->getStartAddress(), block));
+                                                       if (new_blocks.find(jmptarget) == new_blocks.end()) {
+                                                               BasicBlock * block = manager->newBasicBlock(jmptarget);
+                                                               assert(block);
+                                                               new_blocks.insert(std::make_pair(block->getStartAddress(), block));
                                                                function->addBasicBlock(block);
                                                                remaining_blocks.push(block);
                                                        } else {
-                                                               LOG4CXX_DEBUG(logger, "Reusing Block starting at " << std::hex << current_block->getStartAddress());
-                                                               function->addBasicBlock(blocks.find(jmptarget)->second);
+                                                               LOG4CXX_DEBUG(logger, "Reusing Block starting at " << std::hex
+                                                                             << current_block->getStartAddress());
+                                                               function->addBasicBlock(new_blocks.find(jmptarget)->second);
                                                        }
                                                        if (MIA->isConditionalBranch(inst)) {
                                                                jmptarget = base_address + current_address + inst_size;
                                                                current_block->setNextBlock(1, jmptarget);
-                                                               if (blocks.find(jmptarget) == blocks.end()) {
-                                                                       BasicBlock * block = new BasicBlock(jmptarget, this, manager);
-                                                                       blocks.insert(std::make_pair(block->getStartAddress(), block));
+                                                               if (new_blocks.find(jmptarget) == new_blocks.end()) {
+                                                                       BasicBlock * block = manager->newBasicBlock(jmptarget);
+                                                                       assert(block);
+                                                                       new_blocks.insert(std::make_pair(block->getStartAddress(), block));
                                                                        function->addBasicBlock(block);
                                                                        remaining_blocks.push(block);
                                                                } else {
-                                                                       LOG4CXX_DEBUG(logger, "Reusing Block starting at " << std::hex << current_block->getStartAddress());
-                                                                       function->addBasicBlock(blocks.find(jmptarget)->second);
+                                                                       LOG4CXX_DEBUG(logger, "Reusing Block starting at " << std::hex
+                                                                                     << current_block->getStartAddress());
+                                                                       function->addBasicBlock(new_blocks.find(jmptarget)->second);
                                                                }
                                                        }
                                                }
@@ -310,15 +311,16 @@ void LLVMDisassembler<ELFT>::disassemble() {
                        continue;
 
                if (!x->second.getAddress(result)) {
-                       Function * fun = new Function(x->first, result, manager);
+                       Function * fun = manager->newFunction(result);
+                       fun->setName(x->first);
                        remaining_functions.push_back(fun);
-                       functions.insert(std::make_pair(result, fun));
                        LOG4CXX_DEBUG(logger, "Disasembling " << x->first);
                }
        }
 
        for (Function* function : remaining_functions) {
                disassembleFunction(function);
+               manager->finishFunction(function);
        }
 
        if (binary->isELF()) {
@@ -333,12 +335,10 @@ void LLVMDisassembler<ELFT>::disassemble() {
                disassembleFunctionAt(_entryAddress, s.str());
        }
 
-       if (functions.empty()) {
-               uint64_t text_entry;
-               text_section.getAddress(text_entry);
-               LOG4CXX_INFO(logger, "No Symbols found, starting at the beginning of the text segment");
-               disassembleFunctionAt(text_entry);
-       }
+       uint64_t text_entry;
+       text_section.getAddress(text_entry);
+       LOG4CXX_INFO(logger, "No Symbols found, starting at the beginning of the text segment");
+       disassembleFunctionAt(text_entry);
 }
 
 template <typename ELFT>
@@ -365,20 +365,22 @@ void LLVMDisassembler<ELFT>::splitBlocks(Function* function) {
                        if(llvm::MCDisassembler::Success ==
                           DisAsm->getInstruction(inst, inst_size, ref, current_address, nulls(), nulls())) {
                                // See if some other block starts here
-                               auto other = blocks.find(current_address + inst_size + base_address);
+                               BasicBlock* other = manager->getBasicBlock(current_address
+                                                                          + inst_size
+                                                                          + base_address);
 
                                // Special case, other block starts here but we are at the end anyway
-                               if (other != blocks.end()) {
+                               if (other != NULL) {
                                        uint64_t endaddress = current_address + inst_size + base_address;
                                        if (endaddress != current_block->getEndAddress()) {
                                                LOG4CXX_DEBUG(logger, "Shortening block starting at "
                                                              << std::hex
                                                              << current_block->getStartAddress()
                                                              << " now ending at "
-                                                             << other->first);
-                                               function->addBasicBlock(other->second);
+                                                             << other->getStartAddress());
+                                               function->addBasicBlock(other);
                                                current_block->setEndAddress(endaddress);
-                                               current_block->setNextBlock(0, other->first);
+                                               current_block->setNextBlock(0, other->getStartAddress());
                                                current_block->setNextBlock(1, 0);
                                        }
                                }
@@ -440,13 +442,13 @@ void LLVMDisassembler<ELFT>::readSections() {
 
 }
 
-template <typename ELFT>
-void LLVMDisassembler<ELFT>::forEachFunction(std::function<void (uint64_t, Function*)> callback) {
-       std::for_each(functions.begin(), functions.end(),
-                     [&](std::pair<uint64_t, Function*> x) {
-                             callback(x.first, x.second);
-                     });
-}
+// template <typename ELFT>
+// void LLVMDisassembler<ELFT>::forEachFunction(std::function<void (uint64_t, Function*)> callback) {
+//     // std::for_each(functions.begin(), functions.end(),
+//     //               [&](std::pair<uint64_t, Function*> x) {
+//     //                    callback(x.first, x.second);
+//     //               });
+// }
 
 template <typename ELFT>
 void LLVMDisassembler<ELFT>::printEachInstruction(uint64_t start, uint64_t end,