1 files changed, 131 insertions, 20 deletions

diff --git a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
index bc0dd2f..2718389 100644
--- a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
+++ b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Rewrite/ExecutableFileMemoryManager.h"
+#include "bolt/Rewrite/JITLinkLinker.h"
 #include "bolt/Rewrite/RewriteInstance.h"
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/MemAlloc.h"
 
 #undef  DEBUG_TYPE
@@ -21,30 +23,105 @@ namespace llvm {
 
 namespace bolt {
 
-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});
+namespace {
 
-  // A Size of 1 might mean an empty section for which RuntimeDyld decided to
-  // allocate 1 byte. In this case, the allocation will never be initialized
-  // causing non-deterministic output section contents.
-  if (Size == 1)
-    *Ret = 0;
+SmallVector<jitlink::Section *> orderedSections(jitlink::LinkGraph &G) {
+  SmallVector<jitlink::Section *> Sections(
+      llvm::map_range(G.sections(), [](auto &S) { return &S; }));
+  llvm::sort(Sections, [](const auto *LHS, const auto *RHS) {
+    return LHS->getOrdinal() < RHS->getOrdinal();
+  });
+  return Sections;
+}
+
+size_t sectionAlignment(const jitlink::Section &Section) {
+  assert(!Section.empty() && "Cannot get alignment for empty section");
+  return JITLinkLinker::orderedBlocks(Section).front()->getAlignment();
+}
+
+StringRef sectionName(const jitlink::Section &Section,
+                      const BinaryContext &BC) {
+  auto Name = Section.getName();
+
+  if (BC.isMachO()) {
+    // JITLink "normalizes" section names as "SegmentName,SectionName" on
+    // Mach-O. BOLT internally refers to sections just by the section name so
+    // strip-off the segment name.
+    auto SegmentEnd = Name.find(',');
+    assert(SegmentEnd != StringRef::npos && "Mach-O segment not found");
+    Name = Name.substr(SegmentEnd + 1);
+  }
+
+  return Name;
+}
+
+struct SectionAllocInfo {
+  void *Address;
+  size_t Size;
+  size_t Alignment;
+};
+
+struct AllocInfo {
+  SmallVector<SectionAllocInfo, 8> AllocatedSections;
+
+  ~AllocInfo() {
+    for (auto &Section : AllocatedSections)
+      deallocate_buffer(Section.Address, Section.Size, Section.Alignment);
+  }
+
+  SectionAllocInfo allocateSection(const jitlink::Section &Section) {
+    auto Size = JITLinkLinker::sectionSize(Section);
+    auto Alignment = sectionAlignment(Section);
+    auto *Buf = allocate_buffer(Size, Alignment);
+    SectionAllocInfo Alloc{Buf, Size, Alignment};
+    AllocatedSections.push_back(Alloc);
+    return Alloc;
+  }
+};
+
+struct BOLTInFlightAlloc : ExecutableFileMemoryManager::InFlightAlloc {
+  // Even though this is passed using a raw pointer in FinalizedAlloc, we keep
+  // it in a unique_ptr as long as possible to enjoy automatic cleanup when
+  // something goes wrong.
+  std::unique_ptr<AllocInfo> Alloc;
+
+public:
+  BOLTInFlightAlloc(std::unique_ptr<AllocInfo> Alloc)
+      : Alloc(std::move(Alloc)) {}
+
+  virtual void abandon(OnAbandonedFunction OnAbandoned) override {
+    llvm_unreachable("unexpected abandoned allocation");
+  }
+
+  virtual void finalize(OnFinalizedFunction OnFinalized) override {
+    OnFinalized(ExecutableFileMemoryManager::FinalizedAlloc(
+        orc::ExecutorAddr::fromPtr(Alloc.release())));
+  }
+};
+
+} // anonymous namespace
+
+void ExecutableFileMemoryManager::updateSection(
+    const jitlink::Section &JLSection, uint8_t *Contents, size_t Size,
+    size_t Alignment) {
+  auto SectionID = JLSection.getName();
+  auto SectionName = sectionName(JLSection, BC);
+  auto Prot = JLSection.getMemProt();
+  auto IsCode = (Prot & orc::MemProt::Exec) != orc::MemProt::None;
+  auto IsReadOnly = (Prot & orc::MemProt::Write) == orc::MemProt::None;
 
   // Register a debug section as a note section.
   if (!ObjectsLoaded && RewriteInstance::isDebugSection(SectionName)) {
     BinarySection &Section =
-        BC.registerOrUpdateNoteSection(SectionName, Ret, Size, Alignment);
+        BC.registerOrUpdateNoteSection(SectionName, Contents, Size, Alignment);
     Section.setSectionID(SectionID);
     assert(!Section.isAllocatable() && "note sections cannot be allocatable");
-    return Ret;
+    return;
   }
 
   if (!IsCode && (SectionName == ".strtab" || SectionName == ".symtab" ||
                   SectionName == "" || SectionName.startswith(".rela.")))
-    return Ret;
+    return;
 
   SmallVector<char, 256> Buf;
   if (ObjectsLoaded > 0) {
@@ -70,7 +147,7 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
            "Original section must exist and be allocatable.");
 
     Section = &OrgSection.get();
-    Section->updateContents(Ret, Size);
+    Section->updateContents(Contents, Size);
   } else {
     // If the input contains a section with the section name, rename it in the
     // output file to avoid the section name conflict and emit the new section
@@ -87,7 +164,7 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
     // sections in the input file.
     BinarySection &NewSection = BC.registerOrUpdateSection(
         UsePrefix ? NewSecPrefix + SectionName : SectionName, ELF::SHT_PROGBITS,
-        BinarySection::getFlags(IsReadOnly, IsCode, true), Ret, Size,
+        BinarySection::getFlags(IsReadOnly, IsCode, true), Contents, Size,
         Alignment);
     if (UsePrefix)
       NewSection.setOutputName(SectionName);
@@ -100,17 +177,51 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
            << " section : " << Section->getOutputName() << " ("
            << Section->getName() << ")"
            << " with size " << Size << ", alignment " << Alignment << " at "
-           << Ret << ", ID = " << SectionID << "\n";
+           << Contents << ", ID = " << SectionID << "\n";
   });
 
   Section->setSectionID(SectionID);
+}
+
+void ExecutableFileMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
+                                           jitlink::LinkGraph &G,
+                                           OnAllocatedFunction OnAllocated) {
+  auto Alloc = std::make_unique<AllocInfo>();
+
+  for (auto *Section : orderedSections(G)) {
+    if (Section->empty())
+      continue;
+
+    auto SectionAlloc = Alloc->allocateSection(*Section);
+    updateSection(*Section, static_cast<uint8_t *>(SectionAlloc.Address),
+                  SectionAlloc.Size, SectionAlloc.Alignment);
+
+    size_t CurrentOffset = 0;
+    auto *Buf = static_cast<char *>(SectionAlloc.Address);
+    for (auto *Block : JITLinkLinker::orderedBlocks(*Section)) {
+      CurrentOffset = jitlink::alignToBlock(CurrentOffset, *Block);
+      auto BlockSize = Block->getSize();
+      auto *BlockBuf = Buf + CurrentOffset;
+
+      if (Block->isZeroFill())
+        std::memset(BlockBuf, 0, BlockSize);
+      else
+        std::memcpy(BlockBuf, Block->getContent().data(), BlockSize);
+
+      Block->setMutableContent({BlockBuf, Block->getSize()});
+      CurrentOffset += BlockSize;
+    }
+  }
 
-  return Ret;
+  OnAllocated(std::make_unique<BOLTInFlightAlloc>(std::move(Alloc)));
 }
 
-ExecutableFileMemoryManager::~ExecutableFileMemoryManager() {
-  for (const AllocInfo &AI : AllocatedSections)
-    llvm::deallocate_buffer(AI.Address, AI.Size, AI.Alignment);
+void ExecutableFileMemoryManager::deallocate(
+    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
+  for (auto &Alloc : Allocs)
+    delete Alloc.release().toPtr<AllocInfo *>();
+
+  OnDeallocated(Error::success());
 }
 
 } // namespace bolt