3 files changed, 35 insertions, 31 deletions

diff --git a/bolt/include/bolt/Rewrite/ExecutableFileMemoryManager.h b/bolt/include/bolt/Rewrite/ExecutableFileMemoryManager.h
index 9623beb..3549f3d 100644
--- a/bolt/include/bolt/Rewrite/ExecutableFileMemoryManager.h
+++ b/bolt/include/bolt/Rewrite/ExecutableFileMemoryManager.h
@@ -10,7 +10,7 @@
 #define BOLT_REWRITE_EXECUTABLE_FILE_MEMORY_MANAGER_H
 
 #include "llvm/ADT/StringRef.h"
-#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include <cstdint>
 #include <string>
 
@@ -20,14 +20,21 @@ namespace bolt {
 class BinaryContext;
 
 /// Class responsible for allocating and managing code and data sections.
-class ExecutableFileMemoryManager : public SectionMemoryManager {
+class ExecutableFileMemoryManager : public RuntimeDyld::MemoryManager {
 private:
-  uint8_t *allocateSection(intptr_t Size, unsigned Alignment,
+  uint8_t *allocateSection(uintptr_t Size, unsigned Alignment,
                            unsigned SectionID, StringRef SectionName,
                            bool IsCode, bool IsReadOnly);
   BinaryContext &BC;
   bool AllowStubs;
 
+  struct AllocInfo {
+    uint8_t *Address;
+    size_t Size;
+    size_t Alignment;
+  };
+  SmallVector<AllocInfo, 8> AllocatedSections;
+
 public:
   // Our linker's main purpose is to handle a single object file, created
   // by RewriteInstance after reading the input binary and reordering it.
@@ -69,7 +76,16 @@ public:
 
   bool allowStubAllocation() const override { return AllowStubs; }
 
-  bool finalizeMemory(std::string *ErrMsg = nullptr) override;
+  /// Count processed objects and skip memory finalization.
+  bool finalizeMemory(std::string *ErrMsg) override {
+    ++ObjectsLoaded;
+    return false;
+  }
+
+  /// Ignore EH frames.
+  void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
+                        size_t Size) override {}
+  void deregisterEHFrames() override {}
 };
 
 } // namespace bolt
diff --git a/bolt/lib/Core/BinarySection.cpp b/bolt/lib/Core/BinarySection.cpp
index dc0e80c..c5cce79 100644
--- a/bolt/lib/Core/BinarySection.cpp
+++ b/bolt/lib/Core/BinarySection.cpp
@@ -167,7 +167,7 @@ BinarySection::~BinarySection() {
     return;
   }
 
-  if (!isAllocatable() &&
+  if (!isAllocatable() && !hasValidSectionID() &&
       (!hasSectionRef() ||
        OutputContents.data() != getContents(Section).data())) {
     delete[] getOutputData();
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