Revert "[JITLink][ORC] Major JITLinkMemoryManager refactor."

This reverts commit e50aea58d59c8cfae807a7fee21c4227472c0678 while I
investigate bot failures.

1 files changed, 63 insertions, 72 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
index fcfe389..36efc744 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -1,15 +1,10 @@
-//===------- DebugObjectManagerPlugin.cpp - JITLink debug objects ---------===//
+//===---- DebugObjectManagerPlugin.h - JITLink debug objects ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-//
-// FIXME: Update Plugin to poke the debug object into a new JITLink section,
-//        rather than creating a new allocation.
-//
-//===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h"
 
@@ -113,77 +108,70 @@ void ELFDebugObjectSection<ELFT>::dump(raw_ostream &OS, StringRef Name) {
   }
 }
 
+static constexpr sys::Memory::ProtectionFlags ReadOnly =
+    static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ);
+
 enum class Requirement {
   // Request final target memory load-addresses for all sections.
   ReportFinalSectionLoadAddresses,
 };
 
-/// The plugin creates a debug object from when JITLink starts processing the
-/// corresponding LinkGraph. It provides access to the pass configuration of
-/// the LinkGraph and calls the finalization function, once the resulting link
-/// artifact was emitted.
+/// The plugin creates a debug object from JITLinkContext when JITLink starts
+/// processing the corresponding LinkGraph. It provides access to the pass
+/// configuration of the LinkGraph and calls the finalization function, once
+/// the resulting link artifact was emitted.
 ///
 class DebugObject {
 public:
-  DebugObject(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
-              ExecutionSession &ES)
-      : MemMgr(MemMgr), JD(JD), ES(ES) {}
+  DebugObject(JITLinkContext &Ctx, ExecutionSession &ES) : Ctx(Ctx), ES(ES) {}
 
   void set(Requirement Req) { Reqs.insert(Req); }
   bool has(Requirement Req) const { return Reqs.count(Req) > 0; }
 
-  using FinalizeContinuation = std::function<void(Expected<ExecutorAddrRange>)>;
-
+  using FinalizeContinuation = std::function<void(Expected<sys::MemoryBlock>)>;
   void finalizeAsync(FinalizeContinuation OnFinalize);
 
   virtual ~DebugObject() {
-    if (Alloc) {
-      std::vector<FinalizedAlloc> Allocs;
-      Allocs.push_back(std::move(Alloc));
-      if (Error Err = MemMgr.deallocate(std::move(Allocs)))
+    if (Alloc)
+      if (Error Err = Alloc->deallocate())
         ES.reportError(std::move(Err));
-    }
   }
 
   virtual void reportSectionTargetMemoryRange(StringRef Name,
                                               SectionRange TargetMem) {}
 
 protected:
-  using InFlightAlloc = JITLinkMemoryManager::InFlightAlloc;
-  using FinalizedAlloc = JITLinkMemoryManager::FinalizedAlloc;
+  using Allocation = JITLinkMemoryManager::Allocation;
 
-  virtual Expected<SimpleSegmentAlloc> finalizeWorkingMemory() = 0;
-
-  JITLinkMemoryManager &MemMgr;
-  const JITLinkDylib *JD = nullptr;
+  virtual Expected<std::unique_ptr<Allocation>>
+  finalizeWorkingMemory(JITLinkContext &Ctx) = 0;
 
 private:
+  JITLinkContext &Ctx;
   ExecutionSession &ES;
   std::set<Requirement> Reqs;
-  FinalizedAlloc Alloc;
+  std::unique_ptr<Allocation> Alloc{nullptr};
 };
 
 // Finalize working memory and take ownership of the resulting allocation. Start
 // copying memory over to the target and pass on the result once we're done.
 // Ownership of the allocation remains with us for the rest of our lifetime.
 void DebugObject::finalizeAsync(FinalizeContinuation OnFinalize) {
-  assert(!Alloc && "Cannot finalize more than once");
-
-  if (auto SimpleSegAlloc = finalizeWorkingMemory()) {
-    auto ROSeg = SimpleSegAlloc->getSegInfo(MemProt::Read);
-    ExecutorAddrRange DebugObjRange(ExecutorAddr(ROSeg.Addr),
-                                    ExecutorAddrDiff(ROSeg.WorkingMem.size()));
-    SimpleSegAlloc->finalize(
-        [this, DebugObjRange,
-         OnFinalize = std::move(OnFinalize)](Expected<FinalizedAlloc> FA) {
-          if (FA) {
-            Alloc = std::move(*FA);
-            OnFinalize(DebugObjRange);
-          } else
-            OnFinalize(FA.takeError());
-        });
-  } else
-    OnFinalize(SimpleSegAlloc.takeError());
+  assert(Alloc == nullptr && "Cannot finalize more than once");
+
+  auto AllocOrErr = finalizeWorkingMemory(Ctx);
+  if (!AllocOrErr)
+    OnFinalize(AllocOrErr.takeError());
+  Alloc = std::move(*AllocOrErr);
+
+  Alloc->finalizeAsync([this, OnFinalize](Error Err) {
+    if (Err)
+      OnFinalize(std::move(Err));
+    else
+      OnFinalize(sys::MemoryBlock(
+          jitTargetAddressToPointer<void *>(Alloc->getTargetMemory(ReadOnly)),
+          Alloc->getWorkingMemory(ReadOnly).size()));
+  });
 }
 
 /// The current implementation of ELFDebugObject replicates the approach used in
@@ -202,7 +190,8 @@ public:
   StringRef getBuffer() const { return Buffer->getMemBufferRef().getBuffer(); }
 
 protected:
-  Expected<SimpleSegmentAlloc> finalizeWorkingMemory() override;
+  Expected<std::unique_ptr<Allocation>>
+  finalizeWorkingMemory(JITLinkContext &Ctx) override;
 
   template <typename ELFT>
   Error recordSection(StringRef Name,
@@ -212,16 +201,15 @@ protected:
 private:
   template <typename ELFT>
   static Expected<std::unique_ptr<ELFDebugObject>>
-  CreateArchType(MemoryBufferRef Buffer, JITLinkMemoryManager &MemMgr,
-                 const JITLinkDylib *JD, ExecutionSession &ES);
+  CreateArchType(MemoryBufferRef Buffer, JITLinkContext &Ctx,
+                 ExecutionSession &ES);
 
   static std::unique_ptr<WritableMemoryBuffer>
   CopyBuffer(MemoryBufferRef Buffer, Error &Err);
 
   ELFDebugObject(std::unique_ptr<WritableMemoryBuffer> Buffer,
-                 JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
-                 ExecutionSession &ES)
-      : DebugObject(MemMgr, JD, ES), Buffer(std::move(Buffer)) {
+                 JITLinkContext &Ctx, ExecutionSession &ES)
+      : DebugObject(Ctx, ES), Buffer(std::move(Buffer)) {
     set(Requirement::ReportFinalSectionLoadAddresses);
   }
 
@@ -256,14 +244,13 @@ ELFDebugObject::CopyBuffer(MemoryBufferRef Buffer, Error &Err) {
 
 template <typename ELFT>
 Expected<std::unique_ptr<ELFDebugObject>>
-ELFDebugObject::CreateArchType(MemoryBufferRef Buffer,
-                               JITLinkMemoryManager &MemMgr,
-                               const JITLinkDylib *JD, ExecutionSession &ES) {
+ELFDebugObject::CreateArchType(MemoryBufferRef Buffer, JITLinkContext &Ctx,
+                               ExecutionSession &ES) {
   using SectionHeader = typename ELFT::Shdr;
 
   Error Err = Error::success();
   std::unique_ptr<ELFDebugObject> DebugObj(
-      new ELFDebugObject(CopyBuffer(Buffer, Err), MemMgr, JD, ES));
+      new ELFDebugObject(CopyBuffer(Buffer, Err), Ctx, ES));
   if (Err)
     return std::move(Err);
 
@@ -312,26 +299,23 @@ ELFDebugObject::Create(MemoryBufferRef Buffer, JITLinkContext &Ctx,
 
   if (Class == ELF::ELFCLASS32) {
     if (Endian == ELF::ELFDATA2LSB)
-      return CreateArchType<ELF32LE>(Buffer, Ctx.getMemoryManager(),
-                                     Ctx.getJITLinkDylib(), ES);
+      return CreateArchType<ELF32LE>(Buffer, Ctx, ES);
     if (Endian == ELF::ELFDATA2MSB)
-      return CreateArchType<ELF32BE>(Buffer, Ctx.getMemoryManager(),
-                                     Ctx.getJITLinkDylib(), ES);
+      return CreateArchType<ELF32BE>(Buffer, Ctx, ES);
     return nullptr;
   }
   if (Class == ELF::ELFCLASS64) {
     if (Endian == ELF::ELFDATA2LSB)
-      return CreateArchType<ELF64LE>(Buffer, Ctx.getMemoryManager(),
-                                     Ctx.getJITLinkDylib(), ES);
+      return CreateArchType<ELF64LE>(Buffer, Ctx, ES);
     if (Endian == ELF::ELFDATA2MSB)
-      return CreateArchType<ELF64BE>(Buffer, Ctx.getMemoryManager(),
-                                     Ctx.getJITLinkDylib(), ES);
+      return CreateArchType<ELF64BE>(Buffer, Ctx, ES);
     return nullptr;
   }
   return nullptr;
 }
 
-Expected<SimpleSegmentAlloc> ELFDebugObject::finalizeWorkingMemory() {
+Expected<std::unique_ptr<DebugObject::Allocation>>
+ELFDebugObject::finalizeWorkingMemory(JITLinkContext &Ctx) {
   LLVM_DEBUG({
     dbgs() << "Section load-addresses in debug object for \""
            << Buffer->getBufferIdentifier() << "\":\n";
@@ -340,21 +324,28 @@ Expected<SimpleSegmentAlloc> ELFDebugObject::finalizeWorkingMemory() {
   });
 
   // TODO: This works, but what actual alignment requirements do we have?
-  unsigned PageSize = sys::Process::getPageSizeEstimate();
+  unsigned Alignment = sys::Process::getPageSizeEstimate();
+  JITLinkMemoryManager &MemMgr = Ctx.getMemoryManager();
+  const JITLinkDylib *JD = Ctx.getJITLinkDylib();
   size_t Size = Buffer->getBufferSize();
 
   // Allocate working memory for debug object in read-only segment.
-  auto Alloc = SimpleSegmentAlloc::Create(
-      MemMgr, JD, {{MemProt::Read, {Size, Align(PageSize)}}});
-  if (!Alloc)
-    return Alloc;
+  JITLinkMemoryManager::SegmentsRequestMap SingleReadOnlySegment;
+  SingleReadOnlySegment[ReadOnly] =
+      JITLinkMemoryManager::SegmentRequest(Alignment, Size, 0);
+
+  auto AllocOrErr = MemMgr.allocate(JD, SingleReadOnlySegment);
+  if (!AllocOrErr)
+    return AllocOrErr.takeError();
 
   // Initialize working memory with a copy of our object buffer.
-  auto SegInfo = Alloc->getSegInfo(MemProt::Read);
-  memcpy(SegInfo.WorkingMem.data(), Buffer->getBufferStart(), Size);
+  // TODO: Use our buffer as working memory directly.
+  std::unique_ptr<Allocation> Alloc = std::move(*AllocOrErr);
+  MutableArrayRef<char> WorkingMem = Alloc->getWorkingMemory(ReadOnly);
+  memcpy(WorkingMem.data(), Buffer->getBufferStart(), Size);
   Buffer.reset();
 
-  return Alloc;
+  return std::move(Alloc);
 }
 
 void ELFDebugObject::reportSectionTargetMemoryRange(StringRef Name,
@@ -456,7 +447,7 @@ Error DebugObjectManagerPlugin::notifyEmitted(
   std::future<MSVCPError> FinalizeErr = FinalizePromise.get_future();
 
   It->second->finalizeAsync(
-      [this, &FinalizePromise, &MR](Expected<ExecutorAddrRange> TargetMem) {
+      [this, &FinalizePromise, &MR](Expected<sys::MemoryBlock> TargetMem) {
         // Any failure here will fail materialization.
         if (!TargetMem) {
           FinalizePromise.set_value(TargetMem.takeError());