[Orc] Take offset inside slab into account in SharedMemoryMapper

SharedMemoryMapper assumed each reservation will have its corresponding
allocations starting from the beginning. However with the introduction
of the slab allocator, there can be a possible offset from the start
from where the initialization is being performed.

This commit also simplifies the logic for finding the parent reservation
and makes the assert messages consistent.

1 files changed, 8 insertions, 8 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
index 773e6d2..115289d 100644
--- a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
@@ -292,7 +292,7 @@ void SharedMemoryMapper::reserve(size_t NumBytes,
 
 char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
   auto R = Reservations.upper_bound(Addr);
-  assert(R != Reservations.begin() && "Attempt to prepare unknown range");
+  assert(R != Reservations.begin() && "Attempt to prepare unreserved range");
   R--;
 
   ExecutorAddrDiff Offset = Addr - R->first;
@@ -302,11 +302,11 @@ char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
 
 void SharedMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
                                     OnInitializedFunction OnInitialized) {
-  auto Reservation = std::lower_bound(
-      Reservations.rbegin(), Reservations.rend(), AI.MappingBase,
-      [](const auto &A, const auto &B) { return A.first > B; });
-  assert(Reservation != Reservations.rend() &&
-         "Attempt to initialize unreserved range");
+  auto Reservation = Reservations.upper_bound(AI.MappingBase);
+  assert(Reservation != Reservations.begin() && "Attempt to initialize unreserved range");
+  Reservation--;
+
+  auto AllocationOffset = AI.MappingBase - Reservation->first;
 
   tpctypes::SharedMemoryFinalizeRequest FR;
 
@@ -315,8 +315,8 @@ void SharedMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
   FR.Segments.reserve(AI.Segments.size());
 
   for (auto Segment : AI.Segments) {
-    char *Base =
-        static_cast<char *>(Reservation->second.LocalAddr) + Segment.Offset;
+    char *Base = static_cast<char *>(Reservation->second.LocalAddr) +
+                 AllocationOffset + Segment.Offset;
     std::memset(Base + Segment.ContentSize, 0, Segment.ZeroFillSize);
 
     tpctypes::SharedMemorySegFinalizeRequest SegReq;