3 files changed, 22 insertions, 19 deletions

diff --git a/include/llvm/ExecutionEngine/RuntimeDyld.h b/include/llvm/ExecutionEngine/RuntimeDyld.h
index 87ba119457..d0d7dd60db 100644
--- a/include/llvm/ExecutionEngine/RuntimeDyld.h
+++ b/include/llvm/ExecutionEngine/RuntimeDyld.h
@@ -63,9 +63,6 @@ public:
   // and resolve relocatons based on where they put it).
   void *getSymbolAddress(StringRef Name);
   void reassignSymbolAddress(StringRef Name, uint64_t Addr);
-  // FIXME: Should be parameterized to get the memory block associated with
-  // a particular loaded object.
-  sys::MemoryBlock getMemoryBlock();
   StringRef getErrorString();
 };
 
diff --git a/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index 5084b9664d..dbb83faaa0 100644
--- a/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -51,7 +51,7 @@ class RuntimeDyldImpl {
 
   // FIXME: Should have multiple data blocks, one for each loaded chunk of
   //        compiled code.
-  sys::MemoryBlock Data;
+//  sys::MemoryBlock Data;
 
   bool HasError;
   std::string ErrorStr;
@@ -91,8 +91,6 @@ public:
     return Functions.lookup(Name).base();
   }
 
-  sys::MemoryBlock getMemoryBlock() { return Data; }
-
   // Is the linker in an error state?
   bool hasError() { return HasError; }
 
@@ -528,10 +526,6 @@ void *RuntimeDyld::getSymbolAddress(StringRef Name) {
   return Dyld->getSymbolAddress(Name);
 }
 
-sys::MemoryBlock RuntimeDyld::getMemoryBlock() {
-  return Dyld->getMemoryBlock();
-}
-
 StringRef RuntimeDyld::getErrorString() {
   return Dyld->getErrorString();
 }
diff --git a/tools/llvm-rtdyld/llvm-rtdyld.cpp b/tools/llvm-rtdyld/llvm-rtdyld.cpp
index c2f1c3c48d..faec7c3cc5 100644
--- a/tools/llvm-rtdyld/llvm-rtdyld.cpp
+++ b/tools/llvm-rtdyld/llvm-rtdyld.cpp
@@ -44,9 +44,11 @@ Action(cl::desc("Action to perform:"),
 // support library allocation routines directly.
 class TrivialMemoryManager : public RTDyldMemoryManager {
 public:
+  SmallVector<sys::MemoryBlock, 16> FunctionMemory;
+
   uint8_t *startFunctionBody(const char *Name, uintptr_t &Size);
   void endFunctionBody(const char *Name, uint8_t *FunctionStart,
-                       uint8_t *FunctionEnd) {}
+                       uint8_t *FunctionEnd);
 };
 
 uint8_t *TrivialMemoryManager::startFunctionBody(const char *Name,
@@ -54,6 +56,13 @@ uint8_t *TrivialMemoryManager::startFunctionBody(const char *Name,
   return (uint8_t*)sys::Memory::AllocateRWX(Size, 0, 0).base();
 }
 
+void TrivialMemoryManager::endFunctionBody(const char *Name,
+                                           uint8_t *FunctionStart,
+                                           uint8_t *FunctionEnd) {
+  uintptr_t Size = FunctionEnd - FunctionStart + 1;
+  FunctionMemory.push_back(sys::MemoryBlock(FunctionStart, Size));
+}
+
 static const char *ProgramName;
 
 static void Message(const char *Type, const Twine &Msg) {
@@ -74,7 +83,8 @@ static int executeInput() {
     return Error("unable to read input: '" + ec.message() + "'");
 
   // Instantiate a dynamic linker.
-  RuntimeDyld Dyld(new TrivialMemoryManager);
+  TrivialMemoryManager *MemMgr = new TrivialMemoryManager;
+  RuntimeDyld Dyld(MemMgr);
 
   // Load the object file into it.
   if (Dyld.loadObject(InputBuffer.take())) {
@@ -86,14 +96,16 @@ static int executeInput() {
   if (MainAddress == 0)
     return Error("no definition for '_main'");
 
-  // Invalidate the instruction cache.
-  sys::MemoryBlock Data = Dyld.getMemoryBlock();
-  sys::Memory::InvalidateInstructionCache(Data.base(), Data.size());
+  // Invalidate the instruction cache for each loaded function.
+  for (unsigned i = 0, e = MemMgr->FunctionMemory.size(); i != e; ++i) {
+    sys::MemoryBlock &Data = MemMgr->FunctionMemory[i];
+    // Make sure the memory is executable.
+    std::string ErrorStr;
+    sys::Memory::InvalidateInstructionCache(Data.base(), Data.size());
+    if (!sys::Memory::setExecutable(Data, &ErrorStr))
+      return Error("unable to mark function executable: '" + ErrorStr + "'");
+  }
 
-  // Make sure the memory is executable.
-  std::string ErrorStr;
-  if (!sys::Memory::setExecutable(Data, &ErrorStr))
-    return Error("unable to mark function executable: '" + ErrorStr + "'");
 
   // Dispatch to _main().
   errs() << "loaded '_main' at: " << (void*)MainAddress << "\n";