1 files changed, 318 insertions, 0 deletions
diff --git a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
new file mode 100644
index 0000000..ebd9bc3
--- /dev/null
+++ b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
@@ -0,0 +1,318 @@
+//===-- 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);
+}