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

This reverts commit e50aea58d59c8cfae807a7fee21c4227472c0678 while I
investigate bot failures.

20 files changed, 666 insertions, 855 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/CMakeLists.txt b/llvm/lib/ExecutionEngine/JITLink/CMakeLists.txt
index 02157cb..4dfea9e 100644
--- a/llvm/lib/ExecutionEngine/JITLink/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/JITLink/CMakeLists.txt
@@ -3,7 +3,6 @@ add_llvm_component_library(LLVMJITLink
   JITLink.cpp
   JITLinkGeneric.cpp
   JITLinkMemoryManager.cpp
-  MemoryFlags.cpp
 
   # Formats:
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/ELFLinkGraphBuilder.h b/llvm/lib/ExecutionEngine/JITLink/ELFLinkGraphBuilder.h
index fdc9877..05f84f3 100644
--- a/llvm/lib/ExecutionEngine/JITLink/ELFLinkGraphBuilder.h
+++ b/llvm/lib/ExecutionEngine/JITLink/ELFLinkGraphBuilder.h
@@ -36,9 +36,11 @@ protected:
   }
 
   Section &getCommonSection() {
-    if (!CommonSection)
-      CommonSection =
-          &G->createSection(CommonSectionName, MemProt::Read | MemProt::Write);
+    if (!CommonSection) {
+      auto Prot = static_cast<sys::Memory::ProtectionFlags>(
+          sys::Memory::MF_READ | sys::Memory::MF_WRITE);
+      CommonSection = &G->createSection(CommonSectionName, Prot);
+    }
     return *CommonSection;
   }
 
@@ -293,11 +295,13 @@ template <typename ELFT> Error ELFLinkGraphBuilder<ELFT>::graphifySections() {
     });
 
     // Get the section's memory protection flags.
-    MemProt Prot;
+    sys::Memory::ProtectionFlags Prot;
     if (Sec.sh_flags & ELF::SHF_EXECINSTR)
-      Prot = MemProt::Read | MemProt::Exec;
+      Prot = static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                       sys::Memory::MF_EXEC);
     else
-      Prot = MemProt::Read | MemProt::Write;
+      Prot = static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                       sys::Memory::MF_WRITE);
 
     // For now we just use this to skip the "undefined" section, probably need
     // to revist.
diff --git a/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp b/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
index b057788..1f8dbad 100644
--- a/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
@@ -80,14 +80,16 @@ public:
 private:
   Section &getGOTSection() const {
     if (!GOTSection)
-      GOTSection = &G.createSection("$__GOT", MemProt::Read);
+      GOTSection = &G.createSection("$__GOT", sys::Memory::MF_READ);
     return *GOTSection;
   }
 
   Section &getStubsSection() const {
-    if (!StubsSection)
-      StubsSection =
-          &G.createSection("$__STUBS", MemProt::Read | MemProt::Exec);
+    if (!StubsSection) {
+      auto StubsProt = static_cast<sys::Memory::ProtectionFlags>(
+          sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+      StubsSection = &G.createSection("$__STUBS", StubsProt);
+    }
     return *StubsSection;
   }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp b/llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp
index 79c6049..0a50f40 100644
--- a/llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp
@@ -65,7 +65,8 @@ public:
 
   Section &getTLSInfoSection() const {
     if (!TLSInfoSection)
-      TLSInfoSection = &G.createSection(ELFTLSInfoSectionName, MemProt::Read);
+      TLSInfoSection =
+          &G.createSection(ELFTLSInfoSectionName, sys::Memory::MF_READ);
     return *TLSInfoSection;
   }
 
@@ -171,14 +172,16 @@ public:
 private:
   Section &getGOTSection() const {
     if (!GOTSection)
-      GOTSection = &G.createSection(ELFGOTSectionName, MemProt::Read);
+      GOTSection = &G.createSection(ELFGOTSectionName, sys::Memory::MF_READ);
     return *GOTSection;
   }
 
   Section &getStubsSection() const {
-    if (!StubsSection)
-      StubsSection =
-          &G.createSection("$__STUBS", MemProt::Read | MemProt::Exec);
+    if (!StubsSection) {
+      auto StubsProt = static_cast<sys::Memory::ProtectionFlags>(
+          sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+      StubsSection = &G.createSection("$__STUBS", StubsProt);
+    }
     return *StubsSection;
   }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index 706688a..8bb61a2 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -48,21 +48,12 @@ void JITLinkerBase::linkPhase1(std::unique_ptr<JITLinkerBase> Self) {
   if (auto Err = runPasses(Passes.PostPrunePasses))
     return Ctx->notifyFailed(std::move(Err));
 
-  Ctx->getMemoryManager().allocate(
-      Ctx->getJITLinkDylib(), *G,
-      [S = std::move(Self)](AllocResult AR) mutable {
-        auto *TmpSelf = S.get();
-        TmpSelf->linkPhase2(std::move(S), std::move(AR));
-      });
-}
+  // Sort blocks into segments.
+  auto Layout = layOutBlocks();
 
-void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
-                               AllocResult AR) {
-
-  if (AR)
-    Alloc = std::move(*AR);
-  else
-    return Ctx->notifyFailed(AR.takeError());
+  // Allocate memory for segments.
+  if (auto Err = allocateSegments(Layout))
+    return Ctx->notifyFailed(std::move(Err));
 
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName()
@@ -82,16 +73,16 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
 
   auto ExternalSymbols = getExternalSymbolNames();
 
-  // If there are no external symbols then proceed immediately with phase 3.
+  // If there are no external symbols then proceed immediately with phase 2.
   if (ExternalSymbols.empty()) {
     LLVM_DEBUG({
       dbgs() << "No external symbols for " << G->getName()
-             << ". Proceeding immediately with link phase 3.\n";
+             << ". Proceeding immediately with link phase 2.\n";
     });
     // FIXME: Once callee expressions are defined to be sequenced before
     //        argument expressions (c++17) we can simplify this. See below.
     auto &TmpSelf = *Self;
-    TmpSelf.linkPhase3(std::move(Self), AsyncLookupResult());
+    TmpSelf.linkPhase2(std::move(Self), AsyncLookupResult(), std::move(Layout));
     return;
   }
 
@@ -109,31 +100,37 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
   //
   // Ctx->lookup(std::move(UnresolvedExternals),
   //             [Self=std::move(Self)](Expected<AsyncLookupResult> Result) {
-  //               Self->linkPhase3(std::move(Self), std::move(Result));
+  //               Self->linkPhase2(std::move(Self), std::move(Result));
   //             });
-  Ctx->lookup(std::move(ExternalSymbols),
-              createLookupContinuation(
-                  [S = std::move(Self)](
-                      Expected<AsyncLookupResult> LookupResult) mutable {
-                    auto &TmpSelf = *S;
-                    TmpSelf.linkPhase3(std::move(S), std::move(LookupResult));
-                  }));
+  auto *TmpCtx = Ctx.get();
+  TmpCtx->lookup(std::move(ExternalSymbols),
+                 createLookupContinuation(
+                     [S = std::move(Self), L = std::move(Layout)](
+                         Expected<AsyncLookupResult> LookupResult) mutable {
+                       auto &TmpSelf = *S;
+                       TmpSelf.linkPhase2(std::move(S), std::move(LookupResult),
+                                          std::move(L));
+                     }));
 }
 
-void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
-                               Expected<AsyncLookupResult> LR) {
+void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
+                               Expected<AsyncLookupResult> LR,
+                               SegmentLayoutMap Layout) {
 
   LLVM_DEBUG({
-    dbgs() << "Starting link phase 3 for graph " << G->getName() << "\n";
+    dbgs() << "Starting link phase 2 for graph " << G->getName() << "\n";
   });
 
   // If the lookup failed, bail out.
   if (!LR)
-    return abandonAllocAndBailOut(std::move(Self), LR.takeError());
+    return deallocateAndBailOut(LR.takeError());
 
   // Assign addresses to external addressables.
   applyLookupResult(*LR);
 
+  // Copy block content to working memory.
+  copyBlockContentToWorkingMemory(Layout, *Alloc);
+
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName()
            << "\" before pre-fixup passes:\n";
@@ -141,7 +138,7 @@ void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
   });
 
   if (auto Err = runPasses(Passes.PreFixupPasses))
-    return abandonAllocAndBailOut(std::move(Self), std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
 
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName() << "\" before copy-and-fixup:\n";
@@ -150,7 +147,7 @@ void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
 
   // Fix up block content.
   if (auto Err = fixUpBlocks(*G))
-    return abandonAllocAndBailOut(std::move(Self), std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
 
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName() << "\" after copy-and-fixup:\n";
@@ -158,25 +155,27 @@ void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
   });
 
   if (auto Err = runPasses(Passes.PostFixupPasses))
-    return abandonAllocAndBailOut(std::move(Self), std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
 
-  Alloc->finalize([S = std::move(Self)](FinalizeResult FR) mutable {
-    auto *TmpSelf = S.get();
-    TmpSelf->linkPhase4(std::move(S), std::move(FR));
-  });
+  // FIXME: Use move capture once we have c++14.
+  auto *UnownedSelf = Self.release();
+  auto Phase3Continuation = [UnownedSelf](Error Err) {
+    std::unique_ptr<JITLinkerBase> Self(UnownedSelf);
+    UnownedSelf->linkPhase3(std::move(Self), std::move(Err));
+  };
+
+  Alloc->finalizeAsync(std::move(Phase3Continuation));
 }
 
-void JITLinkerBase::linkPhase4(std::unique_ptr<JITLinkerBase> Self,
-                               FinalizeResult FR) {
+void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self, Error Err) {
 
   LLVM_DEBUG({
-    dbgs() << "Starting link phase 4 for graph " << G->getName() << "\n";
+    dbgs() << "Starting link phase 3 for graph " << G->getName() << "\n";
   });
 
-  if (!FR)
-    return Ctx->notifyFailed(FR.takeError());
-
-  Ctx->notifyFinalized(std::move(*FR));
+  if (Err)
+    return deallocateAndBailOut(std::move(Err));
+  Ctx->notifyFinalized(std::move(Alloc));
 
   LLVM_DEBUG({ dbgs() << "Link of graph " << G->getName() << " complete\n"; });
 }
@@ -188,6 +187,131 @@ Error JITLinkerBase::runPasses(LinkGraphPassList &Passes) {
   return Error::success();
 }
 
+JITLinkerBase::SegmentLayoutMap JITLinkerBase::layOutBlocks() {
+
+  SegmentLayoutMap Layout;
+
+  /// Partition blocks based on permissions and content vs. zero-fill.
+  for (auto *B : G->blocks()) {
+    auto &SegLists = Layout[B->getSection().getProtectionFlags()];
+    if (!B->isZeroFill())
+      SegLists.ContentBlocks.push_back(B);
+    else
+      SegLists.ZeroFillBlocks.push_back(B);
+  }
+
+  /// Sort blocks within each list.
+  for (auto &KV : Layout) {
+
+    auto CompareBlocks = [](const Block *LHS, const Block *RHS) {
+      // Sort by section, address and size
+      if (LHS->getSection().getOrdinal() != RHS->getSection().getOrdinal())
+        return LHS->getSection().getOrdinal() < RHS->getSection().getOrdinal();
+      if (LHS->getAddress() != RHS->getAddress())
+        return LHS->getAddress() < RHS->getAddress();
+      return LHS->getSize() < RHS->getSize();
+    };
+
+    auto &SegLists = KV.second;
+    llvm::sort(SegLists.ContentBlocks, CompareBlocks);
+    llvm::sort(SegLists.ZeroFillBlocks, CompareBlocks);
+  }
+
+  LLVM_DEBUG({
+    dbgs() << "Computed segment ordering:\n";
+    for (auto &KV : Layout) {
+      dbgs() << "  Segment "
+             << static_cast<sys::Memory::ProtectionFlags>(KV.first) << ":\n";
+      auto &SL = KV.second;
+      for (auto &SIEntry :
+           {std::make_pair(&SL.ContentBlocks, "content block"),
+            std::make_pair(&SL.ZeroFillBlocks, "zero-fill block")}) {
+        dbgs() << "    " << SIEntry.second << ":\n";
+        for (auto *B : *SIEntry.first)
+          dbgs() << "      " << *B << "\n";
+      }
+    }
+  });
+
+  return Layout;
+}
+
+Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
+
+  // Compute segment sizes and allocate memory.
+  LLVM_DEBUG(dbgs() << "JIT linker requesting: { ");
+  JITLinkMemoryManager::SegmentsRequestMap Segments;
+  for (auto &KV : Layout) {
+    auto &Prot = KV.first;
+    auto &SegLists = KV.second;
+
+    uint64_t SegAlign = 1;
+
+    // Calculate segment content size.
+    size_t SegContentSize = 0;
+    for (auto *B : SegLists.ContentBlocks) {
+      SegAlign = std::max(SegAlign, B->getAlignment());
+      SegContentSize = alignToBlock(SegContentSize, *B);
+      SegContentSize += B->getSize();
+    }
+
+    uint64_t SegZeroFillStart = SegContentSize;
+    uint64_t SegZeroFillEnd = SegZeroFillStart;
+
+    for (auto *B : SegLists.ZeroFillBlocks) {
+      SegAlign = std::max(SegAlign, B->getAlignment());
+      SegZeroFillEnd = alignToBlock(SegZeroFillEnd, *B);
+      SegZeroFillEnd += B->getSize();
+    }
+
+    Segments[Prot] = {SegAlign, SegContentSize,
+                      SegZeroFillEnd - SegZeroFillStart};
+
+    LLVM_DEBUG({
+      dbgs() << (&KV == &*Layout.begin() ? "" : "; ")
+             << static_cast<sys::Memory::ProtectionFlags>(Prot)
+             << ": alignment = " << SegAlign
+             << ", content size = " << SegContentSize
+             << ", zero-fill size = " << (SegZeroFillEnd - SegZeroFillStart);
+    });
+  }
+  LLVM_DEBUG(dbgs() << " }\n");
+
+  if (auto AllocOrErr =
+          Ctx->getMemoryManager().allocate(Ctx->getJITLinkDylib(), Segments))
+    Alloc = std::move(*AllocOrErr);
+  else
+    return AllocOrErr.takeError();
+
+  LLVM_DEBUG({
+    dbgs() << "JIT linker got memory (working -> target):\n";
+    for (auto &KV : Layout) {
+      auto Prot = static_cast<sys::Memory::ProtectionFlags>(KV.first);
+      dbgs() << "  " << Prot << ": "
+             << (const void *)Alloc->getWorkingMemory(Prot).data() << " -> "
+             << formatv("{0:x16}", Alloc->getTargetMemory(Prot)) << "\n";
+    }
+  });
+
+  // Update block target addresses.
+  for (auto &KV : Layout) {
+    auto &Prot = KV.first;
+    auto &SL = KV.second;
+
+    JITTargetAddress NextBlockAddr =
+        Alloc->getTargetMemory(static_cast<sys::Memory::ProtectionFlags>(Prot));
+
+    for (auto *SIList : {&SL.ContentBlocks, &SL.ZeroFillBlocks})
+      for (auto *B : *SIList) {
+        NextBlockAddr = alignToBlock(NextBlockAddr, *B);
+        B->setAddress(NextBlockAddr);
+        NextBlockAddr += B->getSize();
+      }
+  }
+
+  return Error::success();
+}
+
 JITLinkContext::LookupMap JITLinkerBase::getExternalSymbolNames() const {
   // Identify unresolved external symbols.
   JITLinkContext::LookupMap UnresolvedExternals;
@@ -227,13 +351,90 @@ void JITLinkerBase::applyLookupResult(AsyncLookupResult Result) {
   });
 }
 
-void JITLinkerBase::abandonAllocAndBailOut(std::unique_ptr<JITLinkerBase> Self,
-                                           Error Err) {
+void JITLinkerBase::copyBlockContentToWorkingMemory(
+    const SegmentLayoutMap &Layout, JITLinkMemoryManager::Allocation &Alloc) {
+
+  LLVM_DEBUG(dbgs() << "Copying block content:\n");
+  for (auto &KV : Layout) {
+    auto &Prot = KV.first;
+    auto &SegLayout = KV.second;
+
+    auto SegMem =
+        Alloc.getWorkingMemory(static_cast<sys::Memory::ProtectionFlags>(Prot));
+
+    LLVM_DEBUG({
+      dbgs() << "  Processing segment "
+             << static_cast<sys::Memory::ProtectionFlags>(Prot) << " [ "
+             << (const void *)SegMem.data() << " .. "
+             << (const void *)((char *)SegMem.data() + SegMem.size())
+             << " ]\n    Processing content sections:\n";
+    });
+
+    if (SegLayout.ContentBlocks.empty()) {
+      LLVM_DEBUG(dbgs() << "    No content blocks.\n");
+      continue;
+    }
+
+    size_t BlockOffset = 0;
+    size_t LastBlockEnd = 0;
+
+    for (auto *B : SegLayout.ContentBlocks) {
+      LLVM_DEBUG(dbgs() << "    " << *B << ":\n");
+
+      // Pad to alignment/alignment-offset.
+      BlockOffset = alignToBlock(BlockOffset, *B);
+
+      LLVM_DEBUG({
+        dbgs() << "      Bumped block offset to "
+               << formatv("{0:x}", BlockOffset) << " to meet block alignment "
+               << B->getAlignment() << " and alignment offset "
+               << B->getAlignmentOffset() << "\n";
+      });
+
+      // Zero pad up to alignment.
+      LLVM_DEBUG({
+        if (LastBlockEnd != BlockOffset)
+          dbgs() << "      Zero padding from " << formatv("{0:x}", LastBlockEnd)
+                 << " to " << formatv("{0:x}", BlockOffset) << "\n";
+      });
+
+      for (; LastBlockEnd != BlockOffset; ++LastBlockEnd)
+        *(SegMem.data() + LastBlockEnd) = 0;
+
+      // Copy initial block content.
+      LLVM_DEBUG({
+        dbgs() << "      Copying block " << *B << " content, "
+               << B->getContent().size() << " bytes, from "
+               << (const void *)B->getContent().data() << " to offset "
+               << formatv("{0:x}", BlockOffset) << "\n";
+      });
+      memcpy(SegMem.data() + BlockOffset, B->getContent().data(),
+             B->getContent().size());
+
+      // Point the block's content to the fixed up buffer.
+      B->setMutableContent(
+          {SegMem.data() + BlockOffset, B->getContent().size()});
+
+      // Update block end pointer.
+      LastBlockEnd = BlockOffset + B->getContent().size();
+      BlockOffset = LastBlockEnd;
+    }
+
+    // Zero pad the rest of the segment.
+    LLVM_DEBUG({
+      dbgs() << "    Zero padding end of segment from offset "
+             << formatv("{0:x}", LastBlockEnd) << " to "
+             << formatv("{0:x}", SegMem.size()) << "\n";
+    });
+    for (; LastBlockEnd != SegMem.size(); ++LastBlockEnd)
+      *(SegMem.data() + LastBlockEnd) = 0;
+  }
+}
+
+void JITLinkerBase::deallocateAndBailOut(Error Err) {
   assert(Err && "Should not be bailing out on success value");
-  assert(Alloc && "can not call abandonAllocAndBailOut before allocation");
-  Alloc->abandon([S = std::move(Self), E1 = std::move(Err)](Error E2) mutable {
-    S->Ctx->notifyFailed(joinErrors(std::move(E1), std::move(E2)));
-  });
+  assert(Alloc && "can not call deallocateAndBailOut before allocation");
+  Ctx->notifyFailed(joinErrors(std::move(Err), Alloc->deallocate()));
 }
 
 void prune(LinkGraph &G) {
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
index e4fdda0..6b815fe 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
@@ -42,9 +42,14 @@ public:
   virtual ~JITLinkerBase();
 
 protected:
-  using InFlightAlloc = JITLinkMemoryManager::InFlightAlloc;
-  using AllocResult = Expected<std::unique_ptr<InFlightAlloc>>;
-  using FinalizeResult = Expected<JITLinkMemoryManager::FinalizedAlloc>;
+  struct SegmentLayout {
+    using BlocksList = std::vector<Block *>;
+
+    BlocksList ContentBlocks;
+    BlocksList ZeroFillBlocks;
+  };
+
+  using SegmentLayoutMap = DenseMap<unsigned, SegmentLayout>;
 
   // Returns the PassConfiguration for this instance. This can be used by
   // JITLinkerBase implementations to add late passes that reference their
@@ -56,27 +61,39 @@ protected:
   //   1.1: Run pre-prune passes
   //   1.2: Prune graph
   //   1.3: Run post-prune passes
-  //   1.4: Allocate memory.
+  //   1.4: Sort blocks into segments
+  //   1.5: Allocate segment memory, update node vmaddrs to target vmaddrs
+  //   1.6: Run post-allocation passes
+  //   1.7: Notify context of final assigned symbol addresses
+  //   1.8: Identify external symbols and make an async call to resolve
   void linkPhase1(std::unique_ptr<JITLinkerBase> Self);
 
   // Phase 2:
-  //   2.2: Run post-allocation passes
-  //   2.3: Notify context of final assigned symbol addresses
-  //   2.4: Identify external symbols and make an async call to resolve
-  void linkPhase2(std::unique_ptr<JITLinkerBase> Self, AllocResult AR);
+  //   2.1: Apply resolution results
+  //   2.2: Run pre-fixup passes
+  //   2.3: Fix up block contents
+  //   2.4: Run post-fixup passes
+  //   2.5: Make an async call to transfer and finalize memory.
+  void linkPhase2(std::unique_ptr<JITLinkerBase> Self,
+                  Expected<AsyncLookupResult> LookupResult,
+                  SegmentLayoutMap Layout);
 
   // Phase 3:
-  //   3.1: Apply resolution results
-  //   3.2: Run pre-fixup passes
-  //   3.3: Fix up block contents
-  //   3.4: Run post-fixup passes
-  //   3.5: Make an async call to transfer and finalize memory.
-  void linkPhase3(std::unique_ptr<JITLinkerBase> Self,
-                  Expected<AsyncLookupResult> LookupResult);
-
-  // Phase 4:
-  //   4.1: Call OnFinalized callback, handing off allocation.
-  void linkPhase4(std::unique_ptr<JITLinkerBase> Self, FinalizeResult FR);
+  //   3.1: Call OnFinalized callback, handing off allocation.
+  void linkPhase3(std::unique_ptr<JITLinkerBase> Self, Error Err);
+
+  // Align a JITTargetAddress to conform with block alignment requirements.
+  static JITTargetAddress alignToBlock(JITTargetAddress Addr, Block &B) {
+    uint64_t Delta = (B.getAlignmentOffset() - Addr) % B.getAlignment();
+    return Addr + Delta;
+  }
+
+  // Align a pointer to conform with block alignment requirements.
+  static char *alignToBlock(char *P, Block &B) {
+    uint64_t PAddr = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(P));
+    uint64_t Delta = (B.getAlignmentOffset() - PAddr) % B.getAlignment();
+    return P + Delta;
+  }
 
 private:
   // Run all passes in the given pass list, bailing out immediately if any pass
@@ -87,14 +104,18 @@ private:
   // Implemented in JITLinker.
   virtual Error fixUpBlocks(LinkGraph &G) const = 0;
 
+  SegmentLayoutMap layOutBlocks();
+  Error allocateSegments(const SegmentLayoutMap &Layout);
   JITLinkContext::LookupMap getExternalSymbolNames() const;
   void applyLookupResult(AsyncLookupResult LR);
-  void abandonAllocAndBailOut(std::unique_ptr<JITLinkerBase> Self, Error Err);
+  void copyBlockContentToWorkingMemory(const SegmentLayoutMap &Layout,
+                                       JITLinkMemoryManager::Allocation &Alloc);
+  void deallocateAndBailOut(Error Err);
 
   std::unique_ptr<JITLinkContext> Ctx;
   std::unique_ptr<LinkGraph> G;
   PassConfiguration Passes;
-  std::unique_ptr<InFlightAlloc> Alloc;
+  std::unique_ptr<JITLinkMemoryManager::Allocation> Alloc;
 };
 
 template <typename LinkerImpl> class JITLinker : public JITLinkerBase {
@@ -131,8 +152,6 @@ private:
 
       // Copy Block data and apply fixups.
       LLVM_DEBUG(dbgs() << "    Applying fixups.\n");
-      assert((!B->isZeroFill() || B->edges_size() == 0) &&
-             "Edges in zero-fill block?");
       for (auto &E : B->edges()) {
 
         // Skip non-relocation edges.
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index 9069ad8..d572a0b 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -7,479 +7,128 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
-#include "llvm/ExecutionEngine/JITLink/JITLink.h"
-#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Process.h"
 
-#define DEBUG_TYPE "jitlink"
-
 namespace llvm {
 namespace jitlink {
 
 JITLinkMemoryManager::~JITLinkMemoryManager() = default;
-JITLinkMemoryManager::InFlightAlloc::~InFlightAlloc() = default;
-
-static Error runAllocAction(JITLinkMemoryManager::AllocActionCall &C) {
-  using DeallocFnTy = char *(*)(const void *, size_t);
-  auto *Fn = jitTargetAddressToPointer<DeallocFnTy>(C.FnAddr);
-
-  if (char *ErrMsg = Fn(jitTargetAddressToPointer<const void *>(C.CtxAddr),
-                        static_cast<size_t>(C.CtxSize))) {
-    auto E = make_error<StringError>(ErrMsg, inconvertibleErrorCode());
-    free(ErrMsg);
-    return E;
-  }
-
-  return Error::success();
-}
-
-// Align a JITTargetAddress to conform with block alignment requirements.
-static JITTargetAddress alignToBlock(JITTargetAddress Addr, Block &B) {
-  uint64_t Delta = (B.getAlignmentOffset() - Addr) % B.getAlignment();
-  return Addr + Delta;
-}
-
-BasicLayout::BasicLayout(LinkGraph &G) : G(G) {
-
-  for (auto &Sec : G.sections()) {
-    // Skip empty sections.
-    if (empty(Sec.blocks()))
-      continue;
-
-    auto &Seg = Segments[{Sec.getMemProt(), Sec.getMemDeallocPolicy()}];
-    for (auto *B : Sec.blocks())
-      if (LLVM_LIKELY(!B->isZeroFill()))
-        Seg.ContentBlocks.push_back(B);
-      else
-        Seg.ZeroFillBlocks.push_back(B);
-  }
-
-  // Build Segments map.
-  auto CompareBlocks = [](const Block *LHS, const Block *RHS) {
-    // Sort by section, address and size
-    if (LHS->getSection().getOrdinal() != RHS->getSection().getOrdinal())
-      return LHS->getSection().getOrdinal() < RHS->getSection().getOrdinal();
-    if (LHS->getAddress() != RHS->getAddress())
-      return LHS->getAddress() < RHS->getAddress();
-    return LHS->getSize() < RHS->getSize();
-  };
-
-  LLVM_DEBUG(dbgs() << "Generated BasicLayout for " << G.getName() << ":\n");
-  for (auto &KV : Segments) {
-    auto &Seg = KV.second;
-
-    llvm::sort(Seg.ContentBlocks, CompareBlocks);
-    llvm::sort(Seg.ZeroFillBlocks, CompareBlocks);
-
-    for (auto *B : Seg.ContentBlocks) {
-      Seg.ContentSize = alignToBlock(Seg.ContentSize, *B);
-      Seg.ContentSize += B->getSize();
-      Seg.Alignment = std::max(Seg.Alignment, Align(B->getAlignment()));
-    }
-
-    uint64_t SegEndOffset = Seg.ContentSize;
-    for (auto *B : Seg.ZeroFillBlocks) {
-      SegEndOffset = alignToBlock(SegEndOffset, *B);
-      SegEndOffset += B->getSize();
-      Seg.Alignment = std::max(Seg.Alignment, Align(B->getAlignment()));
-    }
-    Seg.ZeroFillSize = SegEndOffset - Seg.ContentSize;
-
-    LLVM_DEBUG({
-      dbgs() << "  Seg " << KV.first
-             << ": content-size=" << formatv("{0:x}", Seg.ContentSize)
-             << ", zero-fill-size=" << formatv("{0:x}", Seg.ZeroFillSize)
-             << ", align=" << formatv("{0:x}", Seg.Alignment.value()) << "\n";
-    });
-  }
-}
-
-Expected<BasicLayout::ContiguousPageBasedLayoutSizes>
-BasicLayout::getContiguousPageBasedLayoutSizes(uint64_t PageSize) {
-  ContiguousPageBasedLayoutSizes SegsSizes;
-
-  for (auto &KV : segments()) {
-    auto &AG = KV.first;
-    auto &Seg = KV.second;
-
-    if (Seg.Alignment > PageSize)
-      return make_error<StringError>("Segment alignment greater than page size",
-                                     inconvertibleErrorCode());
-
-    uint64_t SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize);
-    if (AG.getMemDeallocPolicy() == MemDeallocPolicy::Standard)
-      SegsSizes.StandardSegs += SegSize;
-    else
-      SegsSizes.FinalizeSegs += SegSize;
-  }
-
-  return SegsSizes;
-}
-
-Error BasicLayout::apply() {
-  for (auto &KV : Segments) {
-    auto &Seg = KV.second;
-
-    assert(!(Seg.ContentBlocks.empty() && Seg.ZeroFillBlocks.empty()) &&
-           "Empty section recorded?");
-
-    for (auto *B : Seg.ContentBlocks) {
-      // Align addr and working-mem-offset.
-      Seg.Addr = alignToBlock(Seg.Addr, *B);
-      Seg.NextWorkingMemOffset = alignToBlock(Seg.NextWorkingMemOffset, *B);
-
-      // Update block addr.
-      B->setAddress(Seg.Addr);
-      Seg.Addr += B->getSize();
-
-      // Copy content to working memory, then update content to point at working
-      // memory.
-      memcpy(Seg.WorkingMem + Seg.NextWorkingMemOffset, B->getContent().data(),
-             B->getSize());
-      B->setMutableContent(
-          {Seg.WorkingMem + Seg.NextWorkingMemOffset, B->getSize()});
-      Seg.NextWorkingMemOffset += B->getSize();
-    }
-
-    for (auto *B : Seg.ZeroFillBlocks) {
-      // Align addr.
-      Seg.Addr = alignToBlock(Seg.Addr, *B);
-      // Update block addr.
-      B->setAddress(Seg.Addr);
-      Seg.Addr += B->getSize();
-    }
-
-    Seg.ContentBlocks.clear();
-    Seg.ZeroFillBlocks.clear();
-  }
-
-  return Error::success();
-}
-
-JITLinkMemoryManager::AllocActions &BasicLayout::graphAllocActions() {
-  return G.allocActions();
-}
-
-void SimpleSegmentAlloc::Create(JITLinkMemoryManager &MemMgr,
-                                const JITLinkDylib *JD, SegmentMap Segments,
-                                OnCreatedFunction OnCreated) {
-
-  static_assert(AllocGroup::NumGroups == 16,
-                "AllocGroup has changed. Section names below must be updated");
-  StringRef AGSectionNames[] = {
-      "__---.standard", "__R--.standard", "__-W-.standard", "__RW-.standard",
-      "__--X.standard", "__R-X.standard", "__-WX.standard", "__RWX.standard",
-      "__---.finalize", "__R--.finalize", "__-W-.finalize", "__RW-.finalize",
-      "__--X.finalize", "__R-X.finalize", "__-WX.finalize", "__RWX.finalize"};
-
-  auto G =
-      std::make_unique<LinkGraph>("", Triple(), 0, support::native, nullptr);
-  AllocGroupSmallMap<Block *> ContentBlocks;
-
-  JITTargetAddress NextAddr = 0x100000;
-  for (auto &KV : Segments) {
-    auto &AG = KV.first;
-    auto &Seg = KV.second;
-
-    auto AGSectionName =
-        AGSectionNames[static_cast<unsigned>(AG.getMemProt()) |
-                       static_cast<bool>(AG.getMemDeallocPolicy()) << 3];
-
-    auto &Sec = G->createSection(AGSectionName, AG.getMemProt());
-    Sec.setMemDeallocPolicy(AG.getMemDeallocPolicy());
-
-    if (Seg.ContentSize != 0) {
-      NextAddr = alignTo(NextAddr, Seg.ContentAlign);
-      auto &B =
-          G->createMutableContentBlock(Sec, G->allocateBuffer(Seg.ContentSize),
-                                       NextAddr, Seg.ContentAlign.value(), 0);
-      ContentBlocks[AG] = &B;
-      NextAddr += Seg.ContentSize;
+JITLinkMemoryManager::Allocation::~Allocation() = default;
+
+Expected<std::unique_ptr<JITLinkMemoryManager::Allocation>>
+InProcessMemoryManager::allocate(const JITLinkDylib *JD,
+                                 const SegmentsRequestMap &Request) {
+
+  using AllocationMap = DenseMap<unsigned, sys::MemoryBlock>;
+
+  // Local class for allocation.
+  class IPMMAlloc : public Allocation {
+  public:
+    IPMMAlloc(AllocationMap SegBlocks) : SegBlocks(std::move(SegBlocks)) {}
+    MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
+      assert(SegBlocks.count(Seg) && "No allocation for segment");
+      return {static_cast<char *>(SegBlocks[Seg].base()),
+              SegBlocks[Seg].allocatedSize()};
     }
-  }
-
-  // GRef declared separately since order-of-argument-eval isn't specified.
-  auto &GRef = *G;
-  MemMgr.allocate(JD, GRef,
-                  [G = std::move(G), ContentBlocks = std::move(ContentBlocks),
-                   OnCreated = std::move(OnCreated)](
-                      JITLinkMemoryManager::AllocResult Alloc) mutable {
-                    if (!Alloc)
-                      OnCreated(Alloc.takeError());
-                    else
-                      OnCreated(SimpleSegmentAlloc(std::move(G),
-                                                   std::move(ContentBlocks),
-                                                   std::move(*Alloc)));
-                  });
-}
-
-Expected<SimpleSegmentAlloc>
-SimpleSegmentAlloc::Create(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
-                           SegmentMap Segments) {
-  std::promise<MSVCPExpected<SimpleSegmentAlloc>> AllocP;
-  auto AllocF = AllocP.get_future();
-  Create(MemMgr, JD, std::move(Segments),
-         [&](Expected<SimpleSegmentAlloc> Result) {
-           AllocP.set_value(std::move(Result));
-         });
-  return AllocF.get();
-}
-
-SimpleSegmentAlloc::SimpleSegmentAlloc(SimpleSegmentAlloc &&) = default;
-SimpleSegmentAlloc &
-SimpleSegmentAlloc::operator=(SimpleSegmentAlloc &&) = default;
-SimpleSegmentAlloc::~SimpleSegmentAlloc() {}
-
-SimpleSegmentAlloc::SegmentInfo SimpleSegmentAlloc::getSegInfo(AllocGroup AG) {
-  auto I = ContentBlocks.find(AG);
-  if (I != ContentBlocks.end()) {
-    auto &B = *I->second;
-    return {B.getAddress(), B.getAlreadyMutableContent()};
-  }
-  return {};
-}
-
-SimpleSegmentAlloc::SimpleSegmentAlloc(
-    std::unique_ptr<LinkGraph> G, AllocGroupSmallMap<Block *> ContentBlocks,
-    std::unique_ptr<JITLinkMemoryManager::InFlightAlloc> Alloc)
-    : G(std::move(G)), ContentBlocks(std::move(ContentBlocks)),
-      Alloc(std::move(Alloc)) {}
-
-class InProcessMemoryManager::IPInFlightAlloc
-    : public JITLinkMemoryManager::InFlightAlloc {
-public:
-  IPInFlightAlloc(InProcessMemoryManager &MemMgr, LinkGraph &G, BasicLayout BL,
-                  sys::MemoryBlock StandardSegments,
-                  sys::MemoryBlock FinalizationSegments)
-      : MemMgr(MemMgr), G(G), BL(std::move(BL)),
-        StandardSegments(std::move(StandardSegments)),
-        FinalizationSegments(std::move(FinalizationSegments)) {}
-
-  void finalize(OnFinalizedFunction OnFinalized) override {
-
-    // Apply memory protections to all segments.
-    if (auto Err = applyProtections()) {
-      OnFinalized(std::move(Err));
-      return;
-    }
-
-    // Run finalization actions.
-    // FIXME: Roll back previous successful actions on failure.
-    std::vector<AllocActionCall> DeallocActions;
-    DeallocActions.reserve(G.allocActions().size());
-    for (auto &ActPair : G.allocActions()) {
-      if (ActPair.Finalize.FnAddr)
-        if (auto Err = runAllocAction(ActPair.Finalize)) {
-          OnFinalized(std::move(Err));
-          return;
-        }
-      if (ActPair.Dealloc.FnAddr)
-        DeallocActions.push_back(ActPair.Dealloc);
+    JITTargetAddress getTargetMemory(ProtectionFlags Seg) override {
+      assert(SegBlocks.count(Seg) && "No allocation for segment");
+      return pointerToJITTargetAddress(SegBlocks[Seg].base());
     }
-    G.allocActions().clear();
-
-    // Release the finalize segments slab.
-    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
-      OnFinalized(errorCodeToError(EC));
-      return;
+    void finalizeAsync(FinalizeContinuation OnFinalize) override {
+      OnFinalize(applyProtections());
     }
-
-    // Continue with finalized allocation.
-    OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
-                                            std::move(DeallocActions)));
-  }
-
-  void abandon(OnAbandonedFunction OnAbandoned) override {
-    Error Err = Error::success();
-    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments))
-      Err = joinErrors(std::move(Err), errorCodeToError(EC));
-    if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
-      Err = joinErrors(std::move(Err), errorCodeToError(EC));
-    OnAbandoned(std::move(Err));
-  }
-
-private:
-  Error applyProtections() {
-    for (auto &KV : BL.segments()) {
-      const auto &AG = KV.first;
-      auto &Seg = KV.second;
-
-      auto Prot = toSysMemoryProtectionFlags(AG.getMemProt());
-
-      uint64_t SegSize =
-          alignTo(Seg.ContentSize + Seg.ZeroFillSize, MemMgr.PageSize);
-      sys::MemoryBlock MB(Seg.WorkingMem, SegSize);
-      if (auto EC = sys::Memory::protectMappedMemory(MB, Prot))
+    Error deallocate() override {
+      if (SegBlocks.empty())
+        return Error::success();
+      void *SlabStart = SegBlocks.begin()->second.base();
+      char *SlabEnd = (char *)SlabStart;
+      for (auto &KV : SegBlocks) {
+        SlabStart = std::min(SlabStart, KV.second.base());
+        SlabEnd = std::max(SlabEnd, (char *)(KV.second.base()) +
+                                        KV.second.allocatedSize());
+      }
+      size_t SlabSize = SlabEnd - (char *)SlabStart;
+      assert((SlabSize % sys::Process::getPageSizeEstimate()) == 0 &&
+             "Slab size is not a multiple of page size");
+      sys::MemoryBlock Slab(SlabStart, SlabSize);
+      if (auto EC = sys::Memory::releaseMappedMemory(Slab))
         return errorCodeToError(EC);
-      if (Prot & sys::Memory::MF_EXEC)
-        sys::Memory::InvalidateInstructionCache(MB.base(), MB.allocatedSize());
+      return Error::success();
     }
-    return Error::success();
-  }
-
-  InProcessMemoryManager &MemMgr;
-  LinkGraph &G;
-  BasicLayout BL;
-  sys::MemoryBlock StandardSegments;
-  sys::MemoryBlock FinalizationSegments;
-};
-
-Expected<std::unique_ptr<InProcessMemoryManager>>
-InProcessMemoryManager::Create() {
-  if (auto PageSize = sys::Process::getPageSize())
-    return std::make_unique<InProcessMemoryManager>(*PageSize);
-  else
-    return PageSize.takeError();
-}
-
-void InProcessMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
-                                      OnAllocatedFunction OnAllocated) {
-
-  // FIXME: Just check this once on startup.
-  if (!isPowerOf2_64((uint64_t)PageSize)) {
-    OnAllocated(make_error<StringError>("Page size is not a power of 2",
-                                        inconvertibleErrorCode()));
-    return;
-  }
 
-  BasicLayout BL(G);
-
-  /// Scan the request and calculate the group and total sizes.
-  /// Check that segment size is no larger than a page.
-  auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(PageSize);
-  if (!SegsSizes) {
-    OnAllocated(SegsSizes.takeError());
-    return;
-  }
-
-  // Allocate one slab for the whole thing (to make sure everything is
-  // in-range), then partition into standard and finalization blocks.
-  //
-  // FIXME: Make two separate allocations in the future to reduce
-  // fragmentation: finalization segments will usually be a single page, and
-  // standard segments are likely to be more than one page. Where multiple
-  // allocations are in-flight at once (likely) the current approach will leave
-  // a lot of single-page holes.
-  sys::MemoryBlock Slab;
-  sys::MemoryBlock StandardSegsMem;
-  sys::MemoryBlock FinalizeSegsMem;
-  {
-    const sys::Memory::ProtectionFlags ReadWrite =
-        static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
-                                                  sys::Memory::MF_WRITE);
-
-    std::error_code EC;
-    Slab = sys::Memory::allocateMappedMemory(SegsSizes->total(), nullptr,
-                                             ReadWrite, EC);
-
-    if (EC) {
-      OnAllocated(errorCodeToError(EC));
-      return;
+  private:
+    Error applyProtections() {
+      for (auto &KV : SegBlocks) {
+        auto &Prot = KV.first;
+        auto &Block = KV.second;
+        if (auto EC = sys::Memory::protectMappedMemory(Block, Prot))
+          return errorCodeToError(EC);
+        if (Prot & sys::Memory::MF_EXEC)
+          sys::Memory::InvalidateInstructionCache(Block.base(),
+                                                  Block.allocatedSize());
+      }
+      return Error::success();
     }
 
-    // Zero-fill the whole slab up-front.
-    memset(Slab.base(), 0, Slab.allocatedSize());
-
-    StandardSegsMem = {Slab.base(), SegsSizes->StandardSegs};
-    FinalizeSegsMem = {(void *)((char *)Slab.base() + SegsSizes->StandardSegs),
-                       SegsSizes->FinalizeSegs};
-  }
-
-  auto NextStandardSegAddr = pointerToJITTargetAddress(StandardSegsMem.base());
-  auto NextFinalizeSegAddr = pointerToJITTargetAddress(FinalizeSegsMem.base());
-
-  LLVM_DEBUG({
-    dbgs() << "InProcessMemoryManager allocated:\n";
-    if (SegsSizes->StandardSegs)
-      dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextStandardSegAddr,
-                        NextStandardSegAddr + StandardSegsMem.allocatedSize())
-             << " to stardard segs\n";
-    else
-      dbgs() << "  no standard segs\n";
-    if (SegsSizes->FinalizeSegs)
-      dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextFinalizeSegAddr,
-                        NextFinalizeSegAddr + FinalizeSegsMem.allocatedSize())
-             << " to finalize segs\n";
-    else
-      dbgs() << "  no finalize segs\n";
-  });
+    AllocationMap SegBlocks;
+  };
 
-  // Build ProtMap, assign addresses.
-  for (auto &KV : BL.segments()) {
-    auto &AG = KV.first;
-    auto &Seg = KV.second;
+  if (!isPowerOf2_64((uint64_t)sys::Process::getPageSizeEstimate()))
+    return make_error<StringError>("Page size is not a power of 2",
+                                   inconvertibleErrorCode());
 
-    auto &SegAddr = (AG.getMemDeallocPolicy() == MemDeallocPolicy::Standard)
-                        ? NextStandardSegAddr
-                        : NextFinalizeSegAddr;
+  AllocationMap Blocks;
+  const sys::Memory::ProtectionFlags ReadWrite =
+      static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                sys::Memory::MF_WRITE);
 
-    Seg.WorkingMem = jitTargetAddressToPointer<char *>(SegAddr);
-    Seg.Addr = SegAddr;
+  // Compute the total number of pages to allocate.
+  size_t TotalSize = 0;
+  for (auto &KV : Request) {
+    const auto &Seg = KV.second;
 
-    SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize);
-  }
+    if (Seg.getAlignment() > sys::Process::getPageSizeEstimate())
+      return make_error<StringError>("Cannot request higher than page "
+                                     "alignment",
+                                     inconvertibleErrorCode());
 
-  if (auto Err = BL.apply()) {
-    OnAllocated(std::move(Err));
-    return;
+    TotalSize = alignTo(TotalSize, sys::Process::getPageSizeEstimate());
+    TotalSize += Seg.getContentSize();
+    TotalSize += Seg.getZeroFillSize();
   }
 
-  OnAllocated(std::make_unique<IPInFlightAlloc>(*this, G, std::move(BL),
-                                                std::move(StandardSegsMem),
-                                                std::move(FinalizeSegsMem)));
-}
+  // Allocate one slab to cover all the segments.
+  std::error_code EC;
+  auto SlabRemaining =
+      sys::Memory::allocateMappedMemory(TotalSize, nullptr, ReadWrite, EC);
 
-void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
-                                        OnDeallocatedFunction OnDeallocated) {
-  std::vector<sys::MemoryBlock> StandardSegmentsList;
-  std::vector<std::vector<AllocActionCall>> DeallocActionsList;
+  if (EC)
+    return errorCodeToError(EC);
 
-  {
-    std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
-    for (auto &Alloc : Allocs) {
-      auto *FA =
-          jitTargetAddressToPointer<FinalizedAllocInfo *>(Alloc.release());
-      StandardSegmentsList.push_back(std::move(FA->StandardSegments));
-      if (!FA->DeallocActions.empty())
-        DeallocActionsList.push_back(std::move(FA->DeallocActions));
-      FA->~FinalizedAllocInfo();
-      FinalizedAllocInfos.Deallocate(FA);
-    }
-  }
+  // Allocate segment memory from the slab.
+  for (auto &KV : Request) {
 
-  Error DeallocErr = Error::success();
+    const auto &Seg = KV.second;
 
-  while (!DeallocActionsList.empty()) {
-    auto &DeallocActions = DeallocActionsList.back();
-    auto &StandardSegments = StandardSegmentsList.back();
+    uint64_t SegmentSize = alignTo(Seg.getContentSize() + Seg.getZeroFillSize(),
+                                   sys::Process::getPageSizeEstimate());
+    assert(SlabRemaining.allocatedSize() >= SegmentSize &&
+           "Mapping exceeds allocation");
 
-    /// Run any deallocate calls.
-    while (!DeallocActions.empty()) {
-      if (auto Err = runAllocAction(DeallocActions.back()))
-        DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
-      DeallocActions.pop_back();
-    }
+    sys::MemoryBlock SegMem(SlabRemaining.base(), SegmentSize);
+    SlabRemaining = sys::MemoryBlock((char *)SlabRemaining.base() + SegmentSize,
+                                     SlabRemaining.allocatedSize() - SegmentSize);
 
-    /// Release the standard segments slab.
-    if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
-      DeallocErr = joinErrors(std::move(DeallocErr), errorCodeToError(EC));
+    // Zero out the zero-fill memory.
+    memset(static_cast<char *>(SegMem.base()) + Seg.getContentSize(), 0,
+           Seg.getZeroFillSize());
 
-    DeallocActionsList.pop_back();
-    StandardSegmentsList.pop_back();
+    // Record the block for this segment.
+    Blocks[KV.first] = std::move(SegMem);
   }
 
-  OnDeallocated(std::move(DeallocErr));
-}
-
-JITLinkMemoryManager::FinalizedAlloc
-InProcessMemoryManager::createFinalizedAlloc(
-    sys::MemoryBlock StandardSegments,
-    std::vector<AllocActionCall> DeallocActions) {
-  std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
-  auto *FA = FinalizedAllocInfos.Allocate<FinalizedAllocInfo>();
-  new (FA) FinalizedAllocInfo(
-      {std::move(StandardSegments), std::move(DeallocActions)});
-  return FinalizedAlloc(pointerToJITTargetAddress(FA));
+  return std::unique_ptr<InProcessMemoryManager::Allocation>(
+      new IPMMAlloc(std::move(Blocks)));
 }
 
 } // end namespace jitlink
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.cpp b/llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.cpp
index a54f2a3..22f17331 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.cpp
@@ -107,9 +107,11 @@ MachOLinkGraphBuilder::getEndianness(const object::MachOObjectFile &Obj) {
 }
 
 Section &MachOLinkGraphBuilder::getCommonSection() {
-  if (!CommonSection)
-    CommonSection =
-        &G->createSection(CommonSectionName, MemProt::Read | MemProt::Write);
+  if (!CommonSection) {
+    auto Prot = static_cast<sys::Memory::ProtectionFlags>(
+        sys::Memory::MF_READ | sys::Memory::MF_WRITE);
+    CommonSection = &G->createSection(CommonSectionName, Prot);
+  }
   return *CommonSection;
 }
 
@@ -174,11 +176,13 @@ Error MachOLinkGraphBuilder::createNormalizedSections() {
     // Get prot flags.
     // FIXME: Make sure this test is correct (it's probably missing cases
     // as-is).
-    MemProt Prot;
+    sys::Memory::ProtectionFlags Prot;
     if (NSec.Flags & MachO::S_ATTR_PURE_INSTRUCTIONS)
-      Prot = MemProt::Read | MemProt::Exec;
+      Prot = static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                       sys::Memory::MF_EXEC);
     else
-      Prot = MemProt::Read | MemProt::Write;
+      Prot = static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                       sys::Memory::MF_WRITE);
 
     if (!isDebugSection(NSec)) {
       auto FullyQualifiedName =
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp b/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp
index bf81fc4..8a2c7e5 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp
@@ -457,14 +457,16 @@ public:
 private:
   Section &getGOTSection() {
     if (!GOTSection)
-      GOTSection = &G.createSection("$__GOT", MemProt::Read);
+      GOTSection = &G.createSection("$__GOT", sys::Memory::MF_READ);
     return *GOTSection;
   }
 
   Section &getStubsSection() {
-    if (!StubsSection)
-      StubsSection =
-          &G.createSection("$__STUBS", MemProt::Read | MemProt::Exec);
+    if (!StubsSection) {
+      auto StubsProt = static_cast<sys::Memory::ProtectionFlags>(
+          sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+      StubsSection = &G.createSection("$__STUBS", StubsProt);
+    }
     return *StubsSection;
   }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
index 40ded07..b860816 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
@@ -479,14 +479,16 @@ public:
 private:
   Section &getGOTSection() {
     if (!GOTSection)
-      GOTSection = &G.createSection("$__GOT", MemProt::Read);
+      GOTSection = &G.createSection("$__GOT", sys::Memory::MF_READ);
     return *GOTSection;
   }
 
   Section &getStubsSection() {
-    if (!StubsSection)
-      StubsSection =
-          &G.createSection("$__STUBS", MemProt::Read | MemProt::Exec);
+    if (!StubsSection) {
+      auto StubsProt = static_cast<sys::Memory::ProtectionFlags>(
+          sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+      StubsSection = &G.createSection("$__STUBS", StubsProt);
+    }
     return *StubsSection;
   }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/MemoryFlags.cpp b/llvm/lib/ExecutionEngine/JITLink/MemoryFlags.cpp
deleted file mode 100644
index b73a310b..0000000
--- a/llvm/lib/ExecutionEngine/JITLink/MemoryFlags.cpp
+++ /dev/null
@@ -1,33 +0,0 @@
-//===------------- MemoryFlags.cpp - Memory allocation flags --------------===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#include "llvm/ExecutionEngine/JITLink/MemoryFlags.h"
-
-#define DEBUG_TYPE "jitlink"
-
-namespace llvm {
-namespace jitlink {
-
-raw_ostream &operator<<(raw_ostream &OS, MemProt MP) {
-  return OS << (((MP & MemProt::Read) != MemProt::None) ? 'R' : '-')
-            << (((MP & MemProt::Write) != MemProt::None) ? 'W' : '-')
-            << (((MP & MemProt::Exec) != MemProt::None) ? 'X' : '-');
-}
-
-raw_ostream &operator<<(raw_ostream &OS, MemDeallocPolicy MDP) {
-  return OS << (MDP == MemDeallocPolicy::Standard ? "standard" : "finalize");
-}
-
-raw_ostream &operator<<(raw_ostream &OS, AllocGroup AG) {
-  return OS << '(' << AG.getMemProt() << ", " << AG.getMemDeallocPolicy()
-            << ')';
-}
-
-} // end namespace jitlink
-} // end namespace llvm
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();
 }