4 files changed, 44 insertions, 7 deletions

diff --git a/include/llvm/CodeGen/MachineCodeEmitter.h b/include/llvm/CodeGen/MachineCodeEmitter.h
index 86684d7..4ef230c 100644
--- a/include/llvm/CodeGen/MachineCodeEmitter.h
+++ b/include/llvm/CodeGen/MachineCodeEmitter.h
@@ -207,7 +207,7 @@ public:
   /// allocateSpace - Allocate a block of space in the current output buffer,
   /// returning null (and setting conditions to indicate buffer overflow) on
   /// failure.  Alignment is the alignment in bytes of the buffer desired.
-  void *allocateSpace(intptr_t Size, unsigned Alignment) {
+  virtual void *allocateSpace(intptr_t Size, unsigned Alignment) {
     emitAlignment(Alignment);
     void *Result = CurBufferPtr;
     
diff --git a/include/llvm/ExecutionEngine/JITMemoryManager.h b/include/llvm/ExecutionEngine/JITMemoryManager.h
index 0d79d14..61c3443 100644
--- a/include/llvm/ExecutionEngine/JITMemoryManager.h
+++ b/include/llvm/ExecutionEngine/JITMemoryManager.h
@@ -101,6 +101,9 @@ public:
   /// and remember where it is in case the client wants to deallocate it.
   virtual void endFunctionBody(const Function *F, unsigned char *FunctionStart,
                                unsigned char *FunctionEnd) = 0;
+
+  /// allocateSpace - Allocate a memory block of the given size.
+  virtual unsigned char *allocateSpace(intptr_t Size, unsigned Alignment) = 0;
   
   /// deallocateMemForFunction - Free JIT memory for the specified function.
   /// This is never called when the JIT is currently emitting a function.
diff --git a/lib/ExecutionEngine/JIT/JITEmitter.cpp b/lib/ExecutionEngine/JIT/JITEmitter.cpp
index a90a6a5..688d498 100644
--- a/lib/ExecutionEngine/JIT/JITEmitter.cpp
+++ b/lib/ExecutionEngine/JIT/JITEmitter.cpp
@@ -518,6 +518,10 @@ namespace {
                                    unsigned Alignment = 1);
     virtual void* finishFunctionStub(const GlobalValue *F);
 
+    /// allocateSpace - Reserves space in the current block if any, or
+    /// allocate a new one of the given size.
+    virtual void *allocateSpace(intptr_t Size, unsigned Alignment);
+
     virtual void addRelocation(const MachineRelocation &MR) {
       Relocations.push_back(MR);
     }
@@ -915,11 +919,6 @@ bool JITEmitter::finishFunction(MachineFunction &F) {
                                   Relocations.size(), MemMgr->getGOTBase());
   }
 
-  unsigned char *FnEnd   = CurBufferPtr;
-  
-  MemMgr->endFunctionBody(F.getFunction(), BufferBegin, FnEnd);
-  NumBytes += FnEnd-FnStart;
-
   // Update the GOT entry for F to point to the new code.
   if (MemMgr->isManagingGOT()) {
     unsigned idx = Resolver.getGOTIndexForAddr((void*)BufferBegin);
@@ -930,6 +929,12 @@ bool JITEmitter::finishFunction(MachineFunction &F) {
     }
   }
 
+  unsigned char *FnEnd = CurBufferPtr;
+
+  MemMgr->endFunctionBody(F.getFunction(), BufferBegin, FnEnd);
+  BufferBegin = CurBufferPtr = 0;
+  NumBytes += FnEnd-FnStart;
+
   // Invalidate the icache if necessary.
   sys::Memory::InvalidateInstructionCache(FnStart, FnEnd-FnStart);
   
@@ -993,6 +998,18 @@ bool JITEmitter::finishFunction(MachineFunction &F) {
   return false;
 }
 
+void* JITEmitter::allocateSpace(intptr_t Size, unsigned Alignment) {
+  if (BufferBegin)
+    return MachineCodeEmitter::allocateSpace(Size, Alignment);
+
+  // create a new memory block if there is no active one.
+  // care must be taken so that BufferBegin is invalidated when a
+  // block is trimmed
+  BufferBegin = CurBufferPtr = MemMgr->allocateSpace(Size, Alignment);
+  BufferEnd = BufferBegin+Size;
+  return CurBufferPtr;
+}
+
 void JITEmitter::emitConstantPool(MachineConstantPool *MCP) {
   const std::vector<MachineConstantPoolEntry> &Constants = MCP->getConstants();
   if (Constants.empty()) return;
diff --git a/lib/ExecutionEngine/JIT/JITMemoryManager.cpp b/lib/ExecutionEngine/JIT/JITMemoryManager.cpp
index 618f144..cc072a8 100644
--- a/lib/ExecutionEngine/JIT/JITMemoryManager.cpp
+++ b/lib/ExecutionEngine/JIT/JITMemoryManager.cpp
@@ -298,7 +298,24 @@ namespace {
       // Release the memory at the end of this block that isn't needed.
       FreeMemoryList =CurBlock->TrimAllocationToSize(FreeMemoryList, BlockSize);
     }
-    
+
+    /// allocateSpace - Allocate a memory block of the given size.
+    unsigned char *allocateSpace(intptr_t Size, unsigned Alignment) {
+      CurBlock = FreeMemoryList;
+      FreeMemoryList = FreeMemoryList->AllocateBlock();
+
+      unsigned char *result = (unsigned char *)CurBlock+1;
+
+      if (Alignment == 0) Alignment = 1;
+      result = (unsigned char*)(((intptr_t)result+Alignment-1) &
+               ~(intptr_t)(Alignment-1));
+
+      uintptr_t BlockSize = result + Size - (unsigned char *)CurBlock;
+      FreeMemoryList =CurBlock->TrimAllocationToSize(FreeMemoryList, BlockSize);
+
+      return result;
+    }
+
     /// startExceptionTable - Use startFunctionBody to allocate memory for the 
     /// function's exception table.
     unsigned char* startExceptionTable(const Function* F,