1 files changed, 14 insertions, 26 deletions

diff --git a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
index b186e80..68695d1 100644
--- a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
+++ b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
@@ -8,6 +8,7 @@
 
 #include "bolt/Rewrite/ExecutableFileMemoryManager.h"
 #include "bolt/Rewrite/RewriteInstance.h"
+#include "llvm/Support/MemAlloc.h"
 
 #undef  DEBUG_TYPE
 #define DEBUG_TYPE "efmm"
@@ -20,34 +21,24 @@ namespace llvm {
 
 namespace bolt {
 
-uint8_t *ExecutableFileMemoryManager::allocateSection(intptr_t Size,
-                                                      unsigned Alignment,
-                                                      unsigned SectionID,
-                                                      StringRef SectionName,
-                                                      bool IsCode,
-                                                      bool IsReadOnly) {
+uint8_t *ExecutableFileMemoryManager::allocateSection(
+    uintptr_t Size, unsigned Alignment, unsigned SectionID,
+    StringRef SectionName, bool IsCode, bool IsReadOnly) {
+  uint8_t *Ret = static_cast<uint8_t *>(llvm::allocate_buffer(Size, Alignment));
+  AllocatedSections.push_back(AllocInfo{Ret, Size, Alignment});
+
   // Register a debug section as a note section.
   if (!ObjectsLoaded && RewriteInstance::isDebugSection(SectionName)) {
-    uint8_t *DataCopy = new uint8_t[Size];
     BinarySection &Section =
-        BC.registerOrUpdateNoteSection(SectionName, DataCopy, Size, Alignment);
+        BC.registerOrUpdateNoteSection(SectionName, Ret, Size, Alignment);
     Section.setSectionID(SectionID);
     assert(!Section.isAllocatable() && "note sections cannot be allocatable");
-    return DataCopy;
+    return Ret;
   }
 
   if (!IsCode && (SectionName == ".strtab" || SectionName == ".symtab" ||
                   SectionName == "" || SectionName.startswith(".rela.")))
-    return SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID,
-                                                     SectionName, IsReadOnly);
-
-  uint8_t *Ret;
-  if (IsCode)
-    Ret = SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID,
-                                                    SectionName);
-  else
-    Ret = SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID,
-                                                    SectionName, IsReadOnly);
+    return Ret;
 
   SmallVector<char, 256> Buf;
   if (ObjectsLoaded > 0) {
@@ -75,19 +66,16 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(intptr_t Size,
       dbgs() << "BOLT: allocating "
              << (IsCode ? "code" : (IsReadOnly ? "read-only data" : "data"))
              << " section : " << SectionName << " with size " << Size
-             << ", alignment " << Alignment << " at 0x" << Ret
+             << ", alignment " << Alignment << " at " << Ret
              << ", ID = " << SectionID << "\n");
   return Ret;
 }
 
-bool ExecutableFileMemoryManager::finalizeMemory(std::string *ErrMsg) {
-  LLVM_DEBUG(dbgs() << "BOLT: finalizeMemory()\n");
-  ++ObjectsLoaded;
-  return SectionMemoryManager::finalizeMemory(ErrMsg);
+ExecutableFileMemoryManager::~ExecutableFileMemoryManager() {
+  for (const AllocInfo &AI : AllocatedSections)
+    llvm::deallocate_buffer(AI.Address, AI.Size, AI.Alignment);
 }
 
-ExecutableFileMemoryManager::~ExecutableFileMemoryManager() {}
-
 } // namespace bolt
 
 } // namespace llvm