3 files changed, 23 insertions, 1 deletions
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index 18827978d9a..99524a3b002 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -54,6 +54,14 @@ void RuntimeDyldImpl::resolveRelocations() {
   }
 }
 
+void RuntimeDyldImpl::mapSectionAddress(void *LocalAddress,
+                                        uint64_t TargetAddress) {
+  assert(SectionLocalMemToID.count(LocalAddress) &&
+         "Attempting to remap address of unknown section!");
+  unsigned SectionID = SectionLocalMemToID[LocalAddress];
+  reassignSectionAddress(SectionID, TargetAddress);
+}
+
 //===----------------------------------------------------------------------===//
 // RuntimeDyld class implementation
 RuntimeDyld::RuntimeDyld(RTDyldMemoryManager *mm) {
@@ -116,6 +124,11 @@ void RuntimeDyld::reassignSectionAddress(unsigned SectionID,
   Dyld->reassignSectionAddress(SectionID, Addr);
 }
 
+void RuntimeDyld::mapSectionAddress(void *LocalAddress,
+                                    uint64_t TargetAddress) {
+  Dyld->mapSectionAddress(LocalAddress, TargetAddress);
+}
+
 StringRef RuntimeDyld::getErrorString() {
   return Dyld->getErrorString();
 }
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
index bbfef768c4c..59c863679ef 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
@@ -16,6 +16,7 @@
 
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/Object/MachOObject.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/Twine.h"
@@ -45,10 +46,14 @@ protected:
   // Indexed by SectionID.
   SmallVector<sys::MemoryBlock, 32> Sections;
   // For each section, the address it will be considered to live at for
-  // relocations. The same as the pointer the above memory block for hosted
+  // relocations. The same as the pointer to the above memory block for hosted
   // JITs. Indexed by SectionID.
   SmallVector<uint64_t, 32> SectionLoadAddress;
 
+  // Keep a map of starting local address to the SectionID which references it.
+  // Lookup function for when we assign virtual addresses.
+  DenseMap<void *, unsigned> SectionLocalMemToID;
+
   // Master symbol table. As modules are loaded and external symbols are
   // resolved, their addresses are stored here as a SectionID/Offset pair.
   typedef std::pair<unsigned, uint64_t> SymbolLoc;
@@ -90,6 +95,8 @@ public:
 
   virtual void reassignSectionAddress(unsigned SectionID, uint64_t Addr) = 0;
 
+  void mapSectionAddress(void *LocalAddress, uint64_t TargetAddress);
+
   // Is the linker in an error state?
   bool hasError() { return HasError; }
 
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldMachO.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldMachO.cpp
index a3b4919f68b..1cc021ade51 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldMachO.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldMachO.cpp
@@ -172,6 +172,7 @@ loadSegment32(const MachOObject *Obj,
 
     // Remember what got allocated for this SectionID.
     Sections.push_back(sys::MemoryBlock(Buffer, Sect->Size));
+    SectionLocalMemToID[Buffer] = SectionID;
 
     // By default, the load address of a section is its memory buffer.
     SectionLoadAddress.push_back((uint64_t)Buffer);
@@ -291,6 +292,7 @@ loadSegment64(const MachOObject *Obj,
 
     // Remember what got allocated for this SectionID.
     Sections.push_back(sys::MemoryBlock(Buffer, Sect->Size));
+    SectionLocalMemToID[Buffer] = SectionID;
 
     // By default, the load address of a section is its memory buffer.
     SectionLoadAddress.push_back((uint64_t)Buffer);