//===---- EPCGenericJITLinkMemoryManager.cpp -- Mem management via EPC ----===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h"

#include "llvm/ExecutionEngine/JITLink/JITLink.h"

#include <limits>

using namespace llvm::jitlink;

namespace llvm {
namespace orc {

class EPCGenericJITLinkMemoryManager::InFlightAlloc
    : public jitlink::JITLinkMemoryManager::InFlightAlloc {
public:

  // FIXME: The C++98 initializer is an attempt to work around compile failures
  // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
  // We should be able to switch this back to member initialization once that
  // issue is fixed.
  struct SegInfo {
    SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}

    char *WorkingMem;
    ExecutorAddr Addr;
    uint64_t ContentSize;
    uint64_t ZeroFillSize;
  };

  using SegInfoMap = AllocGroupSmallMap<SegInfo>;

  InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
                ExecutorAddr AllocAddr, SegInfoMap Segs)
      : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}

  void finalize(OnFinalizedFunction OnFinalize) override {
    tpctypes::FinalizeRequest FR;
    for (auto &KV : Segs) {
      assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
      FR.Segments.push_back(tpctypes::SegFinalizeRequest{
          KV.first,
          KV.second.Addr,
          alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
                  Parent.ES.getExecutorProcessControl().getPageSize()),
          {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
    }

    // Transfer allocation actions.
    std::swap(FR.Actions, G.allocActions());

    Parent.B.Initialize(
        [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
            Expected<ExecutorAddr> InitializeKey) mutable {
          // FIXME: Release abandoned alloc.
          if (!InitializeKey)
            return OnFinalize(InitializeKey.takeError());
          OnFinalize(FinalizedAlloc(AllocAddr));
        },
        Parent.ES, Parent.B.Instance, FR);
  }

  void abandon(OnAbandonedFunction OnAbandoned) override {
    // FIXME: Return memory to pool instead.
    Parent.B.Release(std::move(OnAbandoned), Parent.ES, Parent.B.Instance,
                     ArrayRef<ExecutorAddr>(AllocAddr));
  }

private:
  EPCGenericJITLinkMemoryManager &Parent;
  LinkGraph &G;
  ExecutorAddr AllocAddr;
  SegInfoMap Segs;
};

void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
                                              LinkGraph &G,
                                              OnAllocatedFunction OnAllocated) {
  BasicLayout BL(G);

  auto Pages = BL.getContiguousPageBasedLayoutSizes(
      ES.getExecutorProcessControl().getPageSize());
  if (!Pages)
    return OnAllocated(Pages.takeError());

  B.Reserve(
      [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
          Expected<ExecutorAddr> AllocAddr) mutable {
        if (!AllocAddr)
          return OnAllocated(AllocAddr.takeError());
        completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
      },
      ES, B.Instance, Pages->total());
}

void EPCGenericJITLinkMemoryManager::deallocate(
    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
  std::vector<ExecutorAddr> Bases;
  Bases.reserve(Allocs.size());
  for (auto &A : Allocs)
    Bases.push_back(ExecutorAddr(A.getAddress()));
  B.Release(std::move(OnDeallocated), ES, B.Instance, Bases);
  for (auto &A : Allocs)
    A.release();
}

void EPCGenericJITLinkMemoryManager::completeAllocation(
    ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {

  InFlightAlloc::SegInfoMap SegInfos;

  ExecutorAddr NextSegAddr = AllocAddr;
  for (auto &KV : BL.segments()) {
    const auto &AG = KV.first;
    auto &Seg = KV.second;

    Seg.Addr = NextSegAddr;
    KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
    NextSegAddr +=
        ExecutorAddrDiff(alignTo(Seg.ContentSize + Seg.ZeroFillSize,
                                 ES.getExecutorProcessControl().getPageSize()));

    auto &SegInfo = SegInfos[AG];
    SegInfo.ContentSize = Seg.ContentSize;
    SegInfo.ZeroFillSize = Seg.ZeroFillSize;
    SegInfo.Addr = Seg.Addr;
    SegInfo.WorkingMem = Seg.WorkingMem;
  }

  if (auto Err = BL.apply())
    return OnAllocated(std::move(Err));

  OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
                                              std::move(SegInfos)));
}

} // end namespace orc
} // end namespace llvm