2 files changed, 73 insertions, 6 deletions

diff --git a/include/llvm/MC/MCObjectDisassembler.h b/include/llvm/MC/MCObjectDisassembler.h
index 749a54e7f4..de2aae7c6a 100644
--- a/include/llvm/MC/MCObjectDisassembler.h
+++ b/include/llvm/MC/MCObjectDisassembler.h
@@ -15,6 +15,10 @@
 #ifndef LLVM_MC_MCOBJECTDISASSEMBLER_H
 #define LLVM_MC_MCOBJECTDISASSEMBLER_H
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/DataTypes.h"
+
 namespace llvm {
 
 namespace object {
@@ -33,14 +37,11 @@ class MCModule;
 /// It can also be used to create a control flow graph consisting of MCFunctions
 /// and MCBasicBlocks.
 class MCObjectDisassembler {
-  const object::ObjectFile &Obj;
-  const MCDisassembler &Dis;
-  const MCInstrAnalysis &MIA;
-
 public:
   MCObjectDisassembler(const object::ObjectFile &Obj,
                        const MCDisassembler &Dis,
                        const MCInstrAnalysis &MIA);
+  virtual ~MCObjectDisassembler() {}
 
   /// \brief Build an MCModule, creating atoms and optionally functions.
   /// \param withCFG Also build a CFG by adding MCFunctions to the Module.
@@ -50,6 +51,25 @@ public:
   /// block atoms, which then each back an MCBasicBlock.
   MCModule *buildModule(bool withCFG = false);
 
+  MCModule *buildEmptyModule();
+
+  /// \brief Get the effective address of the entrypoint, or 0 if there is none.
+  virtual uint64_t getEntrypoint();
+
+  /// \name Get the addresses of static constructors/destructors in the object.
+  /// The caller is expected to know how to interpret the addresses;
+  /// for example, Mach-O init functions expect 5 arguments, not for ELF.
+  /// The addresses are original object file load addresses, not effective.
+  /// @{
+  virtual ArrayRef<uint64_t> getStaticInitFunctions();
+  virtual ArrayRef<uint64_t> getStaticExitFunctions();
+  /// @}
+
+protected:
+  const object::ObjectFile &Obj;
+  const MCDisassembler &Dis;
+  const MCInstrAnalysis &MIA;
+
 private:
   /// \brief Fill \p Module by creating an atom for each section.
   /// This could be made much smarter, using information like symbols, but also
diff --git a/lib/MC/MCObjectDisassembler.cpp b/lib/MC/MCObjectDisassembler.cpp
index 1ea6eed307..ef8fef1a5a 100644
--- a/lib/MC/MCObjectDisassembler.cpp
+++ b/lib/MC/MCObjectDisassembler.cpp
@@ -33,8 +33,40 @@ MCObjectDisassembler::MCObjectDisassembler(const ObjectFile &Obj,
                                            const MCInstrAnalysis &MIA)
   : Obj(Obj), Dis(Dis), MIA(MIA) {}
 
-MCModule *MCObjectDisassembler::buildModule(bool withCFG) {
+uint64_t MCObjectDisassembler::getEntrypoint() {
+  error_code ec;
+  for (symbol_iterator SI = Obj.begin_symbols(), SE = Obj.end_symbols();
+       SI != SE; SI.increment(ec)) {
+    if (ec)
+      break;
+    StringRef Name;
+    SI->getName(Name);
+    if (Name == "main" || Name == "_main") {
+      uint64_t Entrypoint;
+      SI->getAddress(Entrypoint);
+      return Entrypoint;
+    }
+  }
+  return 0;
+}
+
+ArrayRef<uint64_t> MCObjectDisassembler::getStaticInitFunctions() {
+  return ArrayRef<uint64_t>();
+}
+
+ArrayRef<uint64_t> MCObjectDisassembler::getStaticExitFunctions() {
+  return ArrayRef<uint64_t>();
+}
+
+MCModule *MCObjectDisassembler::buildEmptyModule() {
   MCModule *Module = new MCModule;
+  Module->Entrypoint = getEntrypoint();
+  return Module;
+}
+
+MCModule *MCObjectDisassembler::buildModule(bool withCFG) {
+  MCModule *Module = buildEmptyModule();
+
   buildSectionAtoms(Module);
   if (withCFG)
     buildCFG(Module);
@@ -60,7 +92,7 @@ void MCObjectDisassembler::buildSectionAtoms(MCModule *Module) {
       continue;
 
     StringRef Contents; SI->getContents(Contents);
-    StringRefMemoryObject memoryObject(Contents);
+    StringRefMemoryObject memoryObject(Contents, StartAddr);
 
     // We don't care about things like non-file-backed sections yet.
     if (Contents.size() != SecSize || !SecSize)
@@ -116,6 +148,21 @@ void MCObjectDisassembler::buildCFG(MCModule *Module) {
   AddressSetTy Splits;
   AddressSetTy Calls;
 
+  error_code ec;
+  for (symbol_iterator SI = Obj.begin_symbols(), SE = Obj.end_symbols();
+       SI != SE; SI.increment(ec)) {
+    if (ec)
+      break;
+    SymbolRef::Type SymType;
+    SI->getType(SymType);
+    if (SymType == SymbolRef::ST_Function) {
+      uint64_t SymAddr;
+      SI->getAddress(SymAddr);
+      Calls.insert(SymAddr);
+      Splits.insert(SymAddr);
+    }
+  }
+
   assert(Module->func_begin() == Module->func_end()
          && "Module already has a CFG!");