//===-- SPSNativeMemoryMapTest.cpp ----------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Test SPS serialization for MemoryFlags APIs.
//
//===----------------------------------------------------------------------===//

#include "orc-rt/SimpleNativeMemoryMap.h"
#include "orc-rt/SPSAllocAction.h"
#include "orc-rt/SPSMemoryFlags.h"

#include "AllocActionTestUtils.h"
#include "DirectCaller.h"
#include "gtest/gtest.h"

#include <future>

using namespace orc_rt;

namespace orc_rt {

struct SPSSimpleNativeMemoryMapSegment;

/// A SimpleNativeMemoryMap::FinalizeRequest::Segment plus segment content (if
/// segment content type is regular).
struct TestSNMMSegment
    : public SimpleNativeMemoryMap::FinalizeRequest::Segment {

  enum TestSNMMSegmentContent { Uninitialized, ZeroFill };

  TestSNMMSegment(void *Address, AllocGroup G, std::string Content)
      : SimpleNativeMemoryMap::FinalizeRequest::Segment(
            Address, Content.size(), G, ContentType::Regular),
        Content(std::move(Content)) {}

  TestSNMMSegment(void *Address, size_t Size, AllocGroup G,
                  TestSNMMSegmentContent Content)
      : SimpleNativeMemoryMap::FinalizeRequest::Segment(
            Address, Size, G,
            Content == ZeroFill ? ContentType::ZeroFill
                                : ContentType::Uninitialized) {}

  std::string Content;
};

template <>
class SPSSerializationTraits<SPSSimpleNativeMemoryMapSegment, TestSNMMSegment> {
  using SPSType = SPSTuple<SPSExecutorAddr, uint64_t, SPSAllocGroup, uint8_t>;

public:
  static size_t size(const TestSNMMSegment &S) {
    using ContentType =
        SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType;
    assert((S.C != ContentType::Regular || S.Size == S.Content.size()));
    return SPSType::AsArgList::size(ExecutorAddr::fromPtr(S.Address),
                                    static_cast<uint64_t>(S.Size), S.G,
                                    static_cast<uint8_t>(S.C)) +
           (S.C == ContentType::Regular ? S.Size : 0);
  }

  static bool serialize(SPSOutputBuffer &OB, const TestSNMMSegment &S) {
    using ContentType =
        SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType;
    assert((S.C != ContentType::Regular || S.Size == S.Content.size()));
    if (!SPSType::AsArgList::serialize(OB, ExecutorAddr::fromPtr(S.Address),
                                       static_cast<uint64_t>(S.Size), S.G,
                                       static_cast<uint8_t>(S.C)))
      return false;
    if (S.C == ContentType::Regular)
      return OB.write(S.Content.data(), S.Content.size());
    return true;
  }
};

struct SPSSimpleNativeMemoryMapFinalizeRequest;

struct TestSNMMFinalizeRequest {
  std::vector<TestSNMMSegment> Segments;
  std::vector<AllocActionPair> AAPs;
};

template <>
class SPSSerializationTraits<SPSSimpleNativeMemoryMapFinalizeRequest,
                             TestSNMMFinalizeRequest> {
  using SPSType = SPSTuple<SPSSequence<SPSSimpleNativeMemoryMapSegment>,
                           SPSSequence<SPSAllocActionPair>>;

public:
  static size_t size(const TestSNMMFinalizeRequest &FR) {
    return SPSType::AsArgList::size(FR.Segments, FR.AAPs);
  }
  static bool serialize(SPSOutputBuffer &OB,
                        const TestSNMMFinalizeRequest &FR) {
    return SPSType::AsArgList::serialize(OB, FR.Segments, FR.AAPs);
  }
};

} // namespace orc_rt

template <typename T> move_only_function<void(T)> waitFor(std::future<T> &F) {
  std::promise<T> P;
  F = P.get_future();
  return [P = std::move(P)](T Val) mutable { P.set_value(std::move(Val)); };
}

TEST(SimpleNativeMemoryMapTest, CreateAndDestroy) {
  // Test that we can create and destroy a SimpleNativeMemoryMap instance as
  // expected.
  auto SNMM = std::make_unique<SimpleNativeMemoryMap>();
}

template <typename OnCompleteFn>
static void snmm_reserve(OnCompleteFn &&OnComplete,
                         SimpleNativeMemoryMap *Instance, size_t Size) {
  using SPSSig = SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSSize);
  SPSWrapperFunction<SPSSig>::call(
      DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper),
      std::forward<OnCompleteFn>(OnComplete), Instance, Size);
}

template <typename OnCompleteFn>
static void snmm_finalize(OnCompleteFn &&OnComplete,
                          SimpleNativeMemoryMap *Instance,
                          TestSNMMFinalizeRequest FR) {
  using SPSSig = SPSExpected<SPSExecutorAddr>(
      SPSExecutorAddr, SPSSimpleNativeMemoryMapFinalizeRequest);
  SPSWrapperFunction<SPSSig>::call(
      DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper),
      std::forward<OnCompleteFn>(OnComplete), Instance, std::move(FR));
}

template <typename OnCompleteFn>
static void snmm_deallocate(OnCompleteFn &&OnComplete,
                            SimpleNativeMemoryMap *Instance, void *Base) {
  using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
  SPSWrapperFunction<SPSSig>::call(
      DirectCaller(nullptr,
                   orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper),
      std::forward<OnCompleteFn>(OnComplete), Instance, Base);
}

template <typename OnCompleteFn>
static void snmm_release(OnCompleteFn &&OnComplete,
                         SimpleNativeMemoryMap *Instance, void *Addr) {
  using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
  SPSWrapperFunction<SPSSig>::call(
      DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_release_sps_wrapper),
      std::forward<OnCompleteFn>(OnComplete), Instance, Addr);
}

TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) {
  // Test that we can reserve and release a slab of address space as expected,
  // without finalizing any memory within it.
  auto SNMM = std::make_unique<SimpleNativeMemoryMap>();
  std::future<Expected<Expected<void *>>> ReserveAddr;
  snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024);
  auto Addr = cantFail(cantFail(ReserveAddr.get()));

  std::future<Expected<Error>> ReleaseResult;
  snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
  cantFail(cantFail(ReleaseResult.get()));
}

// Write the given value to the address pointed to by P.
static orc_rt_WrapperFunctionBuffer
write_value_sps_allocaction(const char *ArgData, size_t ArgSize) {
  return SPSAllocActionFunction<SPSExecutorAddr, uint64_t>::handle(
             ArgData, ArgSize,
             [](ExecutorAddr P, uint64_t Val) {
               *P.toPtr<uint64_t *>() = Val;
               return WrapperFunctionBuffer();
             })
      .release();
}

// Read the uint64_t value at Src and write it to Dst.
// Increments int via pointer.
static orc_rt_WrapperFunctionBuffer
read_value_sps_allocaction(const char *ArgData, size_t ArgSize) {
  return SPSAllocActionFunction<SPSExecutorAddr, SPSExecutorAddr>::handle(
             ArgData, ArgSize,
             [](ExecutorAddr Dst, ExecutorAddr Src) {
               *Dst.toPtr<uint64_t *>() = *Src.toPtr<uint64_t *>();
               return WrapperFunctionBuffer();
             })
      .release();
}

TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
  // Test that we can:
  // 1. reserve some address space.
  // 2. finalize a range within it as read/write, and that finalize actions
  //    are applied as expected.
  // 3. deallocate the finalized range, with deallocation actions applied as
  //    expected.
  // 4. release the address range.

  auto SNMM = std::make_unique<SimpleNativeMemoryMap>();
  std::future<Expected<Expected<void *>>> ReserveAddr;
  snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024);
  void *Addr = cantFail(cantFail(ReserveAddr.get()));

  std::future<Expected<Expected<void *>>> FinalizeKey;
  TestSNMMFinalizeRequest FR;
  void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
      reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
  uint64_t SentinelValue = 0;

  FR.Segments.push_back({FinalizeBase, 64 * 1024,
                         MemProt::Read | MemProt::Write,
                         TestSNMMSegment::ZeroFill});
  FR.AAPs.push_back(
      {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
           write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
           uint64_t(42)),
       *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
           read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue),
           ExecutorAddr::fromPtr(FinalizeBase))});
  snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR));
  void *FinalizeKeyAddr = cantFail(cantFail(FinalizeKey.get()));

  EXPECT_EQ(SentinelValue, 0U);

  std::future<Expected<Error>> DeallocResult;
  snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr);
  cantFail(cantFail(DeallocResult.get()));

  EXPECT_EQ(SentinelValue, 42);

  std::future<Expected<Error>> ReleaseResult;
  snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
  cantFail(cantFail(ReleaseResult.get()));
}

TEST(SimpleNativeMemoryMap, ReserveFinalizeShutdown) {
  // Test that memory is deallocated in the case where we reserve and finalize
  // some memory, then just shut down the memory manager.

  auto SNMM = std::make_unique<SimpleNativeMemoryMap>();
  std::future<Expected<Expected<void *>>> ReserveAddr;
  snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024);
  void *Addr = cantFail(cantFail(ReserveAddr.get()));

  std::future<Expected<Expected<void *>>> FinalizeKey;
  TestSNMMFinalizeRequest FR;
  void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
      reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
  uint64_t SentinelValue = 0;

  FR.Segments.push_back({FinalizeBase, 64 * 1024,
                         MemProt::Read | MemProt::Write,
                         TestSNMMSegment::ZeroFill});
  FR.AAPs.push_back(
      {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
           write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
           uint64_t(42)),
       *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
           read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue),
           ExecutorAddr::fromPtr(FinalizeBase))});
  snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR));
  cantFail(cantFail(FinalizeKey.get()));

  EXPECT_EQ(SentinelValue, 0U);

  std::future<Error> ShutdownResult;
  SNMM->shutdown(waitFor(ShutdownResult));
  cantFail(ShutdownResult.get());

  EXPECT_EQ(SentinelValue, 42);
}

TEST(SimpleNativeMemoryMap, ReserveFinalizeDetachShutdown) {
  // Test that memory is deallocated in the case where we reserve and finalize
  // some memory, then just shut down the memory manager.

  auto SNMM = std::make_unique<SimpleNativeMemoryMap>();
  std::future<Expected<Expected<void *>>> ReserveAddr;
  snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024);
  void *Addr = cantFail(cantFail(ReserveAddr.get()));

  std::future<Expected<Expected<void *>>> FinalizeKey;
  TestSNMMFinalizeRequest FR;
  void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
      reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
  uint64_t SentinelValue = 0;

  FR.Segments.push_back({FinalizeBase, 64 * 1024,
                         MemProt::Read | MemProt::Write,
                         TestSNMMSegment::ZeroFill});
  FR.AAPs.push_back(
      {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
           write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
           uint64_t(42)),
       *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
           read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue),
           ExecutorAddr::fromPtr(FinalizeBase))});
  snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR));
  cantFail(cantFail(FinalizeKey.get()));

  EXPECT_EQ(SentinelValue, 0U);

  std::future<Error> DetachResult;
  SNMM->detach(waitFor(DetachResult));
  cantFail(DetachResult.get());

  EXPECT_EQ(SentinelValue, 0);

  std::future<Error> ShutdownResult;
  SNMM->shutdown(waitFor(ShutdownResult));
  cantFail(ShutdownResult.get());

  EXPECT_EQ(SentinelValue, 42);
}