Split BasicBlocks if there are backward jumps

diff --git a/src/disassembler/llvm/LLVMDisassembler.cxx b/src/disassembler/llvm/LLVMDisassembler.cxx
index a194d2144d459d407efa769596a2975e72ed0a44..25a4ecc25303087bee7c3b3e5265c5aa3bc91257 100644 (file)
--- a/src/disassembler/llvm/LLVMDisassembler.cxx
+++ b/src/disassembler/llvm/LLVMDisassembler.cxx
@@ -121,9 +121,6 @@ LLVMDisassembler::~LLVMDisassembler() {
                   });
 }
 
-/*
- * TODO: If we jump into some Basic Block we need to split it there into two
- */
 void LLVMDisassembler::disassemble() {
     std::stack<LLVMFunction*> remaining_functions;
     std::stack<LLVMBasicBlock*> remaining_blocks;
@@ -245,6 +242,39 @@ void LLVMDisassembler::disassemble() {
         }
         LOG4CXX_DEBUG(logger, "Finished function " << current_function->getName());
     }
+
+       // Split blocks where jumps are going inside the block
+       for (auto it = blocks.begin(); it != blocks.end(); ++it) {
+               LLVMBasicBlock * current_block = it->second;
+               uint64_t inst_size;
+               uint64_t base_address;
+               text_section.getAddress(base_address);
+               uint64_t current_address = current_block->getStartAddress() - base_address;
+               while(current_block->getEndAddress() - base_address != current_address) {
+                       MCInst inst;
+                       std::string buf;
+                       llvm::raw_string_ostream s(buf);
+
+                       if(llvm::MCDisassembler::Success ==
+                          DisAsm->getInstruction(inst, inst_size, ref, current_address, nulls(), nulls())) {
+                               auto other = blocks.find(current_address + inst_size + base_address);
+
+                               if (other != blocks.end()) {
+                                       LOG4CXX_DEBUG(logger, "Shortening block starting at "
+                                                                 << std::hex
+                                                                 << current_block->getStartAddress()
+                                                                 << " now ending at "
+                                                                 << other->first);
+                                       current_block->setEndAddress(current_address + inst_size + base_address);
+                                       current_block->setNextBlock(0, other->first);
+                                       current_block->setNextBlock(1, 0);
+                               }
+                       } else {
+                               inst_size = 1;
+                       }
+                       current_address += inst_size;
+               }
+       }
 }
 
 void LLVMDisassembler::readSymbols() {