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

This reverts commit e50aea58d59c8cfae807a7fee21c4227472c0678 while I
investigate bot failures.

9 files changed, 236 insertions, 277 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());
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 88d4cf9..79c602d 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -56,7 +56,7 @@ public:
         "<DSOHandleMU>", TT, PointerSize, Endianness,
         jitlink::getGenericEdgeKindName);
     auto &DSOHandleSection =
-        G->createSection(".data.__dso_handle", jitlink::MemProt::Read);
+        G->createSection(".data.__dso_handle", sys::Memory::MF_READ);
     auto &DSOHandleBlock = G->createContentBlock(
         DSOHandleSection, getDSOHandleContent(PointerSize), 0, 8, 0);
     auto &DSOHandleSymbol = G->addDefinedSymbol(
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp b/llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp
index f3fe055..e1639a8 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp
@@ -43,9 +43,10 @@ createJITLoaderGDBRegistrar(ExecutionSession &ES) {
       ES, ExecutorAddr((*Result)[0][0]));
 }
 
-Error EPCDebugObjectRegistrar::registerDebugObject(
-    ExecutorAddrRange TargetMem) {
-  return ES.callSPSWrapper<void(SPSExecutorAddrRange)>(RegisterFn, TargetMem);
+Error EPCDebugObjectRegistrar::registerDebugObject(sys::MemoryBlock TargetMem) {
+  return ES.callSPSWrapper<void(SPSExecutorAddr, uint64_t)>(
+      RegisterFn, ExecutorAddr::fromPtr(TargetMem.base()),
+      static_cast<uint64_t>(TargetMem.allocatedSize()));
 }
 
 } // namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp
index 9456837..234c085 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp
@@ -7,168 +7,132 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h"
-
-#include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
 
 #include <limits>
 
-using namespace llvm::jitlink;
-
 namespace llvm {
 namespace orc {
 
-class EPCGenericJITLinkMemoryManager::InFlightAlloc
-    : public jitlink::JITLinkMemoryManager::InFlightAlloc {
+class EPCGenericJITLinkMemoryManager::Alloc
+    : public jitlink::JITLinkMemoryManager::Allocation {
 public:
   struct SegInfo {
     char *WorkingMem = nullptr;
-    ExecutorAddr Addr;
+    ExecutorAddr TargetAddr;
     uint64_t ContentSize = 0;
     uint64_t ZeroFillSize = 0;
   };
-  using SegInfoMap = AllocGroupSmallMap<SegInfo>;
+  using SegInfoMap = DenseMap<unsigned, SegInfo>;
+
+  Alloc(EPCGenericJITLinkMemoryManager &Parent, ExecutorAddr TargetAddr,
+        std::unique_ptr<char[]> WorkingBuffer, SegInfoMap Segs)
+      : Parent(Parent), TargetAddr(TargetAddr),
+        WorkingBuffer(std::move(WorkingBuffer)), Segs(std::move(Segs)) {}
+
+  MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
+    auto I = Segs.find(Seg);
+    assert(I != Segs.end() && "No allocation for seg");
+    assert(I->second.ContentSize <= std::numeric_limits<size_t>::max());
+    return {I->second.WorkingMem, static_cast<size_t>(I->second.ContentSize)};
+  }
 
-  InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
-                ExecutorAddr AllocAddr, SegInfoMap Segs)
-      : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
+  JITTargetAddress getTargetMemory(ProtectionFlags Seg) override {
+    auto I = Segs.find(Seg);
+    assert(I != Segs.end() && "No allocation for seg");
+    return I->second.TargetAddr.getValue();
+  }
 
-  void finalize(OnFinalizedFunction OnFinalize) override {
+  void finalizeAsync(FinalizeContinuation OnFinalize) override {
+    char *WorkingMem = WorkingBuffer.get();
     tpctypes::FinalizeRequest FR;
     for (auto &KV : Segs) {
       assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
       FR.Segments.push_back(tpctypes::SegFinalizeRequest{
           tpctypes::toWireProtectionFlags(
-              toSysMemoryProtectionFlags(KV.first.getMemProt())),
-          KV.second.Addr,
+              static_cast<sys::Memory::ProtectionFlags>(KV.first)),
+          KV.second.TargetAddr,
           alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
                   Parent.EPC.getPageSize()),
-          {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
+          {WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
+      WorkingMem += KV.second.ContentSize;
     }
-
-    // Transfer allocation actions.
-    // FIXME: Merge JITLink and ORC SupportFunctionCall and Action list types,
-    //        turn this into a std::swap.
-    FR.Actions.reserve(G.allocActions().size());
-    for (auto &ActPair : G.allocActions())
-      FR.Actions.push_back(
-          {{ExecutorAddr(ActPair.Finalize.FnAddr),
-            ExecutorAddr(ActPair.Finalize.CtxAddr), ActPair.Finalize.CtxSize},
-           {ExecutorAddr(ActPair.Dealloc.FnAddr),
-            ExecutorAddr(ActPair.Dealloc.CtxAddr), ActPair.Dealloc.CtxSize}});
-    G.allocActions().clear();
-
     Parent.EPC.callSPSWrapperAsync<
         rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
         Parent.SAs.Finalize,
-        [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
-            Error SerializationErr, Error FinalizeErr) mutable {
-          // FIXME: Release abandoned alloc.
+        [OnFinalize = std::move(OnFinalize)](Error SerializationErr,
+                                             Error FinalizeErr) {
           if (SerializationErr) {
             cantFail(std::move(FinalizeErr));
             OnFinalize(std::move(SerializationErr));
-          } else if (FinalizeErr)
+          } else
             OnFinalize(std::move(FinalizeErr));
-          else
-            OnFinalize(FinalizedAlloc(AllocAddr.getValue()));
         },
         Parent.SAs.Allocator, std::move(FR));
   }
 
-  void abandon(OnAbandonedFunction OnAbandoned) override {
-    // FIXME: Return memory to pool instead.
-    Parent.EPC.callSPSWrapperAsync<
-        rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
-        Parent.SAs.Deallocate,
-        [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
-                                               Error DeallocateErr) mutable {
-          if (SerializationErr) {
-            cantFail(std::move(DeallocateErr));
-            OnAbandoned(std::move(SerializationErr));
-          } else
-            OnAbandoned(std::move(DeallocateErr));
-        },
-        Parent.SAs.Allocator, ArrayRef<ExecutorAddr>(AllocAddr));
+  Error deallocate() override {
+    Error Err = Error::success();
+    if (auto E2 = Parent.EPC.callSPSWrapper<
+                  rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
+            Parent.SAs.Deallocate, Err, Parent.SAs.Allocator,
+            ArrayRef<ExecutorAddr>(TargetAddr)))
+      return E2;
+    return Err;
   }
 
 private:
   EPCGenericJITLinkMemoryManager &Parent;
-  LinkGraph &G;
-  ExecutorAddr AllocAddr;
+  ExecutorAddr TargetAddr;
+  std::unique_ptr<char[]> WorkingBuffer;
   SegInfoMap Segs;
 };
 
-void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
-                                              LinkGraph &G,
-                                              OnAllocatedFunction OnAllocated) {
-  BasicLayout BL(G);
-
-  auto Pages = BL.getContiguousPageBasedLayoutSizes(EPC.getPageSize());
-  if (!Pages)
-    return OnAllocated(Pages.takeError());
-
-  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
-      SAs.Reserve,
-      [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
-          Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
-        if (SerializationErr) {
-          cantFail(AllocAddr.takeError());
-          return OnAllocated(std::move(SerializationErr));
-        }
-        if (!AllocAddr)
-          return OnAllocated(AllocAddr.takeError());
-
-        completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
-      },
-      SAs.Allocator, Pages->total());
-}
-
-void EPCGenericJITLinkMemoryManager::deallocate(
-    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
-  EPC.callSPSWrapperAsync<
-      rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
-      SAs.Deallocate,
-      [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
-                                                 Error DeallocErr) mutable {
-        if (SerErr) {
-          cantFail(std::move(DeallocErr));
-          OnDeallocated(std::move(SerErr));
-        } else
-          OnDeallocated(std::move(DeallocErr));
-      },
-      SAs.Allocator, Allocs);
-  for (auto &A : Allocs)
-    A.release();
-}
-
-void EPCGenericJITLinkMemoryManager::completeAllocation(
-    ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
-
-  InFlightAlloc::SegInfoMap SegInfos;
+Expected<std::unique_ptr<jitlink::JITLinkMemoryManager::Allocation>>
+EPCGenericJITLinkMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
+                                         const SegmentsRequestMap &Request) {
+  Alloc::SegInfoMap Segs;
+  uint64_t AllocSize = 0;
+  size_t WorkingSize = 0;
+  for (auto &KV : Request) {
+    if (!isPowerOf2_64(KV.second.getAlignment()))
+      return make_error<StringError>("Alignment is not a power of two",
+                                     inconvertibleErrorCode());
+    if (KV.second.getAlignment() > EPC.getPageSize())
+      return make_error<StringError>("Alignment exceeds page size",
+                                     inconvertibleErrorCode());
+
+    auto &Seg = Segs[KV.first];
+    Seg.ContentSize = KV.second.getContentSize();
+    Seg.ZeroFillSize = KV.second.getZeroFillSize();
+    AllocSize += alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize());
+    WorkingSize += Seg.ContentSize;
+  }
 
-  ExecutorAddr NextSegAddr = AllocAddr;
-  for (auto &KV : BL.segments()) {
-    const auto &AG = KV.first;
+  std::unique_ptr<char[]> WorkingBuffer;
+  if (WorkingSize > 0)
+    WorkingBuffer = std::make_unique<char[]>(WorkingSize);
+  Expected<ExecutorAddr> TargetAllocAddr((ExecutorAddr()));
+  if (auto Err = EPC.callSPSWrapper<
+                 rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
+          SAs.Reserve, TargetAllocAddr, SAs.Allocator, AllocSize))
+    return std::move(Err);
+  if (!TargetAllocAddr)
+    return TargetAllocAddr.takeError();
+
+  char *WorkingMem = WorkingBuffer.get();
+  JITTargetAddress SegAddr = TargetAllocAddr->getValue();
+  for (auto &KV : Segs) {
     auto &Seg = KV.second;
-
-    Seg.Addr = NextSegAddr.getValue();
-    KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
-    NextSegAddr += ExecutorAddrDiff(
-        alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize()));
-
-    auto &SegInfo = SegInfos[AG];
-    SegInfo.ContentSize = Seg.ContentSize;
-    SegInfo.ZeroFillSize = Seg.ZeroFillSize;
-    SegInfo.Addr = ExecutorAddr(Seg.Addr);
-    SegInfo.WorkingMem = Seg.WorkingMem;
+    Seg.TargetAddr.setValue(SegAddr);
+    SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize());
+    Seg.WorkingMem = WorkingMem;
+    WorkingMem += Seg.ContentSize;
   }
 
-  if (auto Err = BL.apply())
-    return OnAllocated(std::move(Err));
-
-  OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
-                                              std::move(SegInfos)));
+  return std::make_unique<Alloc>(*this, *TargetAllocAddr,
+                                 std::move(WorkingBuffer), std::move(Segs));
 }
 
 } // end namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCIndirectionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/EPCIndirectionUtils.cpp
index 948d1d9..b9c70b0 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCIndirectionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCIndirectionUtils.cpp
@@ -43,12 +43,12 @@ public:
 protected:
   Error grow() override;
 
-  using FinalizedAlloc = jitlink::JITLinkMemoryManager::FinalizedAlloc;
+  using Allocation = jitlink::JITLinkMemoryManager::Allocation;
 
   EPCIndirectionUtils &EPCIU;
   unsigned TrampolineSize = 0;
   unsigned TrampolinesPerPage = 0;
-  std::vector<FinalizedAlloc> TrampolineBlocks;
+  std::vector<std::unique_ptr<Allocation>> TrampolineBlocks;
 };
 
 class EPCIndirectStubsManager : public IndirectStubsManager,
@@ -89,19 +89,12 @@ EPCTrampolinePool::EPCTrampolinePool(EPCIndirectionUtils &EPCIU)
 
 Error EPCTrampolinePool::deallocatePool() {
   Error Err = Error::success();
-  std::promise<MSVCPError> DeallocResultP;
-  auto DeallocResultF = DeallocResultP.get_future();
-
-  EPCIU.getExecutorProcessControl().getMemMgr().deallocate(
-      std::move(TrampolineBlocks),
-      [&](Error Err) { DeallocResultP.set_value(std::move(Err)); });
-
-  return DeallocResultF.get();
+  for (auto &Alloc : TrampolineBlocks)
+    Err = joinErrors(std::move(Err), Alloc->deallocate());
+  return Err;
 }
 
 Error EPCTrampolinePool::grow() {
-  using namespace jitlink;
-
   assert(AvailableTrampolines.empty() &&
          "Grow called with trampolines still available");
 
@@ -109,26 +102,34 @@ Error EPCTrampolinePool::grow() {
   assert(ResolverAddress && "Resolver address can not be null");
 
   auto &EPC = EPCIU.getExecutorProcessControl();
+  constexpr auto TrampolinePagePermissions =
+      static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                sys::Memory::MF_EXEC);
   auto PageSize = EPC.getPageSize();
-  auto Alloc = SimpleSegmentAlloc::Create(
-      EPC.getMemMgr(), nullptr,
-      {{MemProt::Read | MemProt::Exec, {PageSize, Align(PageSize)}}});
+  jitlink::JITLinkMemoryManager::SegmentsRequestMap Request;
+  Request[TrampolinePagePermissions] = {PageSize, static_cast<size_t>(PageSize),
+                                        0};
+  auto Alloc = EPC.getMemMgr().allocate(nullptr, Request);
+
   if (!Alloc)
     return Alloc.takeError();
 
   unsigned NumTrampolines = TrampolinesPerPage;
 
-  auto SegInfo = Alloc->getSegInfo(MemProt::Read | MemProt::Exec);
-  EPCIU.getABISupport().writeTrampolines(
-      SegInfo.WorkingMem.data(), SegInfo.Addr, ResolverAddress, NumTrampolines);
+  auto WorkingMemory = (*Alloc)->getWorkingMemory(TrampolinePagePermissions);
+  auto TargetAddress = (*Alloc)->getTargetMemory(TrampolinePagePermissions);
+
+  EPCIU.getABISupport().writeTrampolines(WorkingMemory.data(), TargetAddress,
+                                         ResolverAddress, NumTrampolines);
+
+  auto TargetAddr = (*Alloc)->getTargetMemory(TrampolinePagePermissions);
   for (unsigned I = 0; I < NumTrampolines; ++I)
-    AvailableTrampolines.push_back(SegInfo.Addr + (I * TrampolineSize));
+    AvailableTrampolines.push_back(TargetAddr + (I * TrampolineSize));
 
-  auto FA = Alloc->finalize();
-  if (!FA)
-    return FA.takeError();
+  if (auto Err = (*Alloc)->finalize())
+    return Err;
 
-  TrampolineBlocks.push_back(std::move(*FA));
+  TrampolineBlocks.push_back(std::move(*Alloc));
 
   return Error::success();
 }
@@ -266,17 +267,17 @@ EPCIndirectionUtils::Create(ExecutorProcessControl &EPC) {
 }
 
 Error EPCIndirectionUtils::cleanup() {
+  Error Err = Error::success();
 
-  auto &MemMgr = EPC.getMemMgr();
-  auto Err = MemMgr.deallocate(std::move(IndirectStubAllocs));
+  for (auto &A : IndirectStubAllocs)
+    Err = joinErrors(std::move(Err), A->deallocate());
 
   if (TP)
     Err = joinErrors(std::move(Err),
                      static_cast<EPCTrampolinePool &>(*TP).deallocatePool());
 
   if (ResolverBlock)
-    Err =
-        joinErrors(std::move(Err), MemMgr.deallocate(std::move(ResolverBlock)));
+    Err = joinErrors(std::move(Err), ResolverBlock->deallocate());
 
   return Err;
 }
@@ -284,29 +285,29 @@ Error EPCIndirectionUtils::cleanup() {
 Expected<JITTargetAddress>
 EPCIndirectionUtils::writeResolverBlock(JITTargetAddress ReentryFnAddr,
                                         JITTargetAddress ReentryCtxAddr) {
-  using namespace jitlink;
-
   assert(ABI && "ABI can not be null");
+  constexpr auto ResolverBlockPermissions =
+      static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                sys::Memory::MF_EXEC);
   auto ResolverSize = ABI->getResolverCodeSize();
 
-  auto Alloc =
-      SimpleSegmentAlloc::Create(EPC.getMemMgr(), nullptr,
-                                 {{MemProt::Read | MemProt::Exec,
-                                   {ResolverSize, Align(EPC.getPageSize())}}});
-
+  jitlink::JITLinkMemoryManager::SegmentsRequestMap Request;
+  Request[ResolverBlockPermissions] = {EPC.getPageSize(),
+                                       static_cast<size_t>(ResolverSize), 0};
+  auto Alloc = EPC.getMemMgr().allocate(nullptr, Request);
   if (!Alloc)
     return Alloc.takeError();
 
-  auto SegInfo = Alloc->getSegInfo(MemProt::Read | MemProt::Exec);
-  ABI->writeResolverCode(SegInfo.WorkingMem.data(), SegInfo.Addr, ReentryFnAddr,
+  auto WorkingMemory = (*Alloc)->getWorkingMemory(ResolverBlockPermissions);
+  ResolverBlockAddr = (*Alloc)->getTargetMemory(ResolverBlockPermissions);
+  ABI->writeResolverCode(WorkingMemory.data(), ResolverBlockAddr, ReentryFnAddr,
                          ReentryCtxAddr);
 
-  auto FA = Alloc->finalize();
-  if (!FA)
-    return FA.takeError();
+  if (auto Err = (*Alloc)->finalize())
+    return std::move(Err);
 
-  ResolverBlock = std::move(*FA);
-  return SegInfo.Addr;
+  ResolverBlock = std::move(*Alloc);
+  return ResolverBlockAddr;
 }
 
 std::unique_ptr<IndirectStubsManager>
@@ -340,7 +341,6 @@ EPCIndirectionUtils::EPCIndirectionUtils(ExecutorProcessControl &EPC,
 
 Expected<EPCIndirectionUtils::IndirectStubInfoVector>
 EPCIndirectionUtils::getIndirectStubs(unsigned NumStubs) {
-  using namespace jitlink;
 
   std::lock_guard<std::mutex> Lock(EPCUIMutex);
 
@@ -350,40 +350,42 @@ EPCIndirectionUtils::getIndirectStubs(unsigned NumStubs) {
     auto PageSize = EPC.getPageSize();
     auto StubBytes = alignTo(NumStubsToAllocate * ABI->getStubSize(), PageSize);
     NumStubsToAllocate = StubBytes / ABI->getStubSize();
-    auto PtrBytes =
+    auto PointerBytes =
         alignTo(NumStubsToAllocate * ABI->getPointerSize(), PageSize);
 
-    auto StubProt = MemProt::Read | MemProt::Exec;
-    auto PtrProt = MemProt::Read | MemProt::Write;
-
-    auto Alloc = SimpleSegmentAlloc::Create(
-        EPC.getMemMgr(), nullptr,
-        {{StubProt, {static_cast<size_t>(StubBytes), Align(PageSize)}},
-         {PtrProt, {PtrBytes, Align(PageSize)}}});
-
+    constexpr auto StubPagePermissions =
+        static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                  sys::Memory::MF_EXEC);
+    constexpr auto PointerPagePermissions =
+        static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                  sys::Memory::MF_WRITE);
+
+    jitlink::JITLinkMemoryManager::SegmentsRequestMap Request;
+    Request[StubPagePermissions] = {PageSize, static_cast<size_t>(StubBytes),
+                                    0};
+    Request[PointerPagePermissions] = {PageSize, 0, PointerBytes};
+    auto Alloc = EPC.getMemMgr().allocate(nullptr, Request);
     if (!Alloc)
       return Alloc.takeError();
 
-    auto StubSeg = Alloc->getSegInfo(StubProt);
-    auto PtrSeg = Alloc->getSegInfo(PtrProt);
-
-    ABI->writeIndirectStubsBlock(StubSeg.WorkingMem.data(), StubSeg.Addr,
-                                 PtrSeg.Addr, NumStubsToAllocate);
+    auto StubTargetAddr = (*Alloc)->getTargetMemory(StubPagePermissions);
+    auto PointerTargetAddr = (*Alloc)->getTargetMemory(PointerPagePermissions);
 
-    auto FA = Alloc->finalize();
-    if (!FA)
-      return FA.takeError();
+    ABI->writeIndirectStubsBlock(
+        (*Alloc)->getWorkingMemory(StubPagePermissions).data(), StubTargetAddr,
+        PointerTargetAddr, NumStubsToAllocate);
 
-    IndirectStubAllocs.push_back(std::move(*FA));
+    if (auto Err = (*Alloc)->finalize())
+      return std::move(Err);
 
-    auto StubExecutorAddr = StubSeg.Addr;
-    auto PtrExecutorAddr = PtrSeg.Addr;
     for (unsigned I = 0; I != NumStubsToAllocate; ++I) {
       AvailableIndirectStubs.push_back(
-          IndirectStubInfo(StubExecutorAddr, PtrExecutorAddr));
-      StubExecutorAddr += ABI->getStubSize();
-      PtrExecutorAddr += ABI->getPointerSize();
+          IndirectStubInfo(StubTargetAddr, PointerTargetAddr));
+      StubTargetAddr += ABI->getStubSize();
+      PointerTargetAddr += ABI->getPointerSize();
     }
+
+    IndirectStubAllocs.push_back(std::move(*Alloc));
   }
 
   assert(NumStubs <= AvailableIndirectStubs.size() &&
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
index 0145387..6fb8b52 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
@@ -31,8 +31,7 @@ SelfExecutorProcessControl::SelfExecutorProcessControl(
 
   OwnedMemMgr = std::move(MemMgr);
   if (!OwnedMemMgr)
-    OwnedMemMgr = std::make_unique<jitlink::InProcessMemoryManager>(
-        sys::Process::getPageSizeEstimate());
+    OwnedMemMgr = std::make_unique<jitlink::InProcessMemoryManager>();
 
   this->TargetTriple = std::move(TargetTriple);
   this->PageSize = PageSize;
diff --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
index e346262..1705010 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -53,7 +53,7 @@ public:
     auto G = std::make_unique<jitlink::LinkGraph>(
         "<MachOHeaderMU>", TT, PointerSize, Endianness,
         jitlink::getGenericEdgeKindName);
-    auto &HeaderSection = G->createSection("__header", jitlink::MemProt::Read);
+    auto &HeaderSection = G->createSection("__header", sys::Memory::MF_READ);
     auto &HeaderBlock = createHeaderBlock(*G, HeaderSection);
 
     // Init symbol is header-start symbol.
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 22a6425..40d4f19 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -306,7 +306,8 @@ public:
     return Error::success();
   }
 
-  void notifyFinalized(JITLinkMemoryManager::FinalizedAlloc A) override {
+  void notifyFinalized(
+      std::unique_ptr<JITLinkMemoryManager::Allocation> A) override {
     if (auto Err = Layer.notifyEmitted(*MR, std::move(A))) {
       Layer.getExecutionSession().reportError(std::move(Err));
       MR->failMaterialization();
@@ -679,7 +680,7 @@ void ObjectLinkingLayer::notifyLoaded(MaterializationResponsibility &MR) {
 }
 
 Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
-                                        FinalizedAlloc FA) {
+                                        AllocPtr Alloc) {
   Error Err = Error::success();
   for (auto &P : Plugins)
     Err = joinErrors(std::move(Err), P->notifyEmitted(MR));
@@ -688,20 +689,17 @@ Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
     return Err;
 
   return MR.withResourceKeyDo(
-      [&](ResourceKey K) { Allocs[K].push_back(std::move(FA)); });
+      [&](ResourceKey K) { Allocs[K].push_back(std::move(Alloc)); });
 }
 
 Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
 
-  {
-    Error Err = Error::success();
-    for (auto &P : Plugins)
-      Err = joinErrors(std::move(Err), P->notifyRemovingResources(K));
-    if (Err)
-      return Err;
-  }
+  Error Err = Error::success();
 
-  std::vector<FinalizedAlloc> AllocsToRemove;
+  for (auto &P : Plugins)
+    Err = joinErrors(std::move(Err), P->notifyRemovingResources(K));
+
+  std::vector<AllocPtr> AllocsToRemove;
   getExecutionSession().runSessionLocked([&] {
     auto I = Allocs.find(K);
     if (I != Allocs.end()) {
@@ -710,7 +708,12 @@ Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
     }
   });
 
-  return MemMgr.deallocate(std::move(AllocsToRemove));
+  while (!AllocsToRemove.empty()) {
+    Err = joinErrors(std::move(Err), AllocsToRemove.back()->deallocate());
+    AllocsToRemove.pop_back();
+  }
+
+  return Err;
 }
 
 void ObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
diff --git a/llvm/lib/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.cpp b/llvm/lib/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.cpp
index 787a6c6..0ee194c 100644
--- a/llvm/lib/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.cpp
@@ -64,16 +64,15 @@ LLVM_ATTRIBUTE_NOINLINE void __jit_debug_register_code() {
 }
 
 using namespace llvm;
-using namespace llvm::orc;
 
 // Serialize rendezvous with the debugger as well as access to shared data.
 ManagedStatic<std::mutex> JITDebugLock;
 
 // Register debug object, return error message or null for success.
-static void registerJITLoaderGDBImpl(ExecutorAddrRange DebugObjRange) {
+static void registerJITLoaderGDBImpl(JITTargetAddress Addr, uint64_t Size) {
   jit_code_entry *E = new jit_code_entry;
-  E->symfile_addr = DebugObjRange.Start.toPtr<const char *>();
-  E->symfile_size = DebugObjRange.size().getValue();
+  E->symfile_addr = jitTargetAddressToPointer<const char *>(Addr);
+  E->symfile_size = Size;
   E->prev_entry = nullptr;
 
   std::lock_guard<std::mutex> Lock(*JITDebugLock);
@@ -96,7 +95,7 @@ static void registerJITLoaderGDBImpl(ExecutorAddrRange DebugObjRange) {
 extern "C" orc::shared::detail::CWrapperFunctionResult
 llvm_orc_registerJITLoaderGDBWrapper(const char *Data, uint64_t Size) {
   using namespace orc::shared;
-  return WrapperFunction<void(SPSExecutorAddrRange)>::handle(
+  return WrapperFunction<void(SPSExecutorAddr, uint64_t)>::handle(
              Data, Size, registerJITLoaderGDBImpl)
       .release();
 }