1 files changed, 45 insertions, 18 deletions

diff --git a/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index 267307cfde05..9e0d207e8bdb 100644
--- a/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -38,9 +38,21 @@ InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) {
       OnFinalize(applyProtections());
     }
     Error deallocate() override {
-      for (auto &KV : SegBlocks)
-        if (auto EC = sys::Memory::releaseMappedMemory(KV.second))
-          return errorCodeToError(EC);
+      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);
       return Error::success();
     }
 
@@ -61,37 +73,52 @@ InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) {
     AllocationMap SegBlocks;
   };
 
+  if (!isPowerOf2_64((uint64_t)sys::Process::getPageSizeEstimate()))
+    return make_error<StringError>("Page size is not a power of 2",
+                                   inconvertibleErrorCode());
+
   AllocationMap Blocks;
   const sys::Memory::ProtectionFlags ReadWrite =
       static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
                                                 sys::Memory::MF_WRITE);
 
+  // Compute the total number of pages to allocate.
+  size_t TotalSize = 0;
   for (auto &KV : Request) {
-    auto &Seg = KV.second;
+    const auto &Seg = KV.second;
 
-    if (Seg.getContentAlignment() > sys::Process::getPageSizeEstimate())
+    if (Seg.getAlignment() > sys::Process::getPageSizeEstimate())
       return make_error<StringError>("Cannot request higher than page "
                                      "alignment",
                                      inconvertibleErrorCode());
 
-    if (sys::Process::getPageSizeEstimate() % Seg.getContentAlignment() != 0)
-      return make_error<StringError>("Page size is not a multiple of "
-                                     "alignment",
-                                     inconvertibleErrorCode());
+    TotalSize = alignTo(TotalSize, sys::Process::getPageSizeEstimate());
+    TotalSize += Seg.getContentSize();
+    TotalSize += Seg.getZeroFillSize();
+  }
+
+  // Allocate one slab to cover all the segments.
+  std::error_code EC;
+  auto SlabRemaining =
+      sys::Memory::allocateMappedMemory(TotalSize, nullptr, ReadWrite, EC);
+
+  if (EC)
+    return errorCodeToError(EC);
+
+  // Allocate segment memory from the slab.
+  for (auto &KV : Request) {
 
-    uint64_t ZeroFillStart =
-        alignTo(Seg.getContentSize(), Seg.getZeroFillAlignment());
-    uint64_t SegmentSize = ZeroFillStart + Seg.getZeroFillSize();
+    const auto &Seg = KV.second;
 
-    std::error_code EC;
-    auto SegMem =
-        sys::Memory::allocateMappedMemory(SegmentSize, nullptr, ReadWrite, EC);
+    uint64_t SegmentSize = alignTo(Seg.getContentSize() + Seg.getZeroFillSize(),
+                                   sys::Process::getPageSizeEstimate());
 
-    if (EC)
-      return errorCodeToError(EC);
+    sys::MemoryBlock SegMem(SlabRemaining.base(), SegmentSize);
+    SlabRemaining = sys::MemoryBlock((char *)SlabRemaining.base() + SegmentSize,
+                                     SegmentSize);
 
     // Zero out the zero-fill memory.
-    memset(static_cast<char *>(SegMem.base()) + ZeroFillStart, 0,
+    memset(static_cast<char *>(SegMem.base()) + Seg.getContentSize(), 0,
            Seg.getZeroFillSize());
 
     // Record the block for this segment.