[ORC] Move 'CompleteAllocation' lambda body into its own method.

[lhames]

Summary

Upcoming changes will add more complexity to CompleteAllocation. Moving it into
its own method gives it room to grow without impacting the readability of the
allocate method.

[sgraenitz]

I am not the author of the underlying code, but it looks to me like a pure API extension with no functional change. Is that right?

[lhames]

Yep -- just moves the lambda body out into a method. NFC.

[sgraenitz]

LGTM

llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h

public:
   void deallocate(std::vector<FinalizedAlloc> Allocs,
                   OnDeallocatedFunction OnDeallocated) override;
   // synchronous overload
   using JITLinkMemoryManager::deallocate;
 
 private:
   class InFlightAlloc;
 
+  void completeAllocation(jitlink::LinkGraph &G, jitlink::BasicLayout BL,
+                          OnAllocatedFunction OnAllocated,
+                          Expected<ExecutorAddrRange> Result);
+
   std::mutex Mutex;
 
   // We reserve multiples of this from the executor address space
   size_t ReservationUnits;
 
   // Ranges that have been reserved in executor but not yet allocated
   using AvailableMemoryMap = IntervalMap<ExecutorAddr, bool>;
   AvailableMemoryMap::Allocator AMAllocator;

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
   auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
   if (!SegsSizes) {
     OnAllocated(SegsSizes.takeError());
     return;
   }
 
   auto TotalSize = SegsSizes->total();
 
-  auto CompleteAllocation = [this, &G, BL = std::move(BL),
-                             OnAllocated = std::move(OnAllocated)](
-                                Expected<ExecutorAddrRange> Result) mutable {
-    if (!Result) {
-      Mutex.unlock();
-      return OnAllocated(Result.takeError());
-    }
-
-    auto NextSegAddr = Result->Start;
-
-    std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
-
-    for (auto &KV : BL.segments()) {
-      auto &AG = KV.first;
-      auto &Seg = KV.second;
-
-      auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
-
-      Seg.Addr = NextSegAddr;
-      Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
-
-      NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
-
-      MemoryMapper::AllocInfo::SegInfo SI;
-      SI.Offset = Seg.Addr - Result->Start;
-      SI.ContentSize = Seg.ContentSize;
-      SI.ZeroFillSize = Seg.ZeroFillSize;
-      SI.AG = AG;
-      SI.WorkingMem = Seg.WorkingMem;
-
-      SegInfos.push_back(SI);
-    }
-
-    UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
-
-    if (NextSegAddr < Result->End) {
-      // Save the remaining memory for reuse in next allocation(s)
-      AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
-    }
-    Mutex.unlock();
-
-    if (auto Err = BL.apply()) {
-      OnAllocated(std::move(Err));
-      return;
-    }
-
-    OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
-                                                std::move(SegInfos)));
+  auto CompleteAlloc = [this, &G, BL = std::move(BL),
+    completeAllocation(G, std::move(BL), std::move(OnAllocated),
+                       std::move(Result));
   };
 
   Mutex.lock();
 
   // find an already reserved range that is large enough
   ExecutorAddrRange SelectedRange{};
 
   for (AvailableMemoryMap::iterator It = AvailableMemory.begin();
        It != AvailableMemory.end(); It++) {
     if (It.stop() - It.start() + 1 >= TotalSize) {
       SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);
       It.erase();
       break;
     }
   }
 
   if (SelectedRange.empty()) { // no already reserved range was found
     auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
-    Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
+    Mapper->reserve(TotalAllocation, std::move(CompleteAlloc));
   } else {
-    CompleteAllocation(SelectedRange);
+    CompleteAlloc(SelectedRange);
   }
 }
 
 void MapperJITLinkMemoryManager::deallocate(
     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
   std::vector<ExecutorAddr> Bases;
   Bases.reserve(Allocs.size());
   for (auto &FA : Allocs) {
if (Err) {
         FA.release();
       }
     }
 
     OnDeallocated(Error::success());
   });
 }
 
+void MapperJITLinkMemoryManager::completeAllocation(
+    jitlink::LinkGraph &G, jitlink::BasicLayout BL,
+    OnAllocatedFunction OnAllocated, Expected<ExecutorAddrRange> Result) {
+  if (!Result) {
+    Mutex.unlock();
+    return OnAllocated(Result.takeError());
+  }
+
+  auto NextSegAddr = Result->Start;
+
+  std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
+
+  for (auto &KV : BL.segments()) {
+    auto &AG = KV.first;
+    auto &Seg = KV.second;
+
+    auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
+
+    Seg.Addr = NextSegAddr;
+    Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
+
+    NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
+
+    MemoryMapper::AllocInfo::SegInfo SI;
+    SI.Offset = Seg.Addr - Result->Start;
+    SI.ContentSize = Seg.ContentSize;
+    SI.ZeroFillSize = Seg.ZeroFillSize;
+    SI.AG = AG;
+    SI.WorkingMem = Seg.WorkingMem;
+
+    SegInfos.push_back(SI);
+  }
+
+  UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
+
+  if (NextSegAddr < Result->End) {
+    // Save the remaining memory for reuse in next allocation(s)
+    AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
+  }
+  Mutex.unlock();
+
+  if (auto Err = BL.apply()) {
+    OnAllocated(std::move(Err));
+    return;
+  }
+
+  OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
+                                              std::move(SegInfos)));
+}
+
 } // end namespace orc
 } // end namespace llvm