[orc-rt] Align SimpleNativeMemoryMap Segment with LLVM type. (llvm#162823)

This commit aims to align SimpleNativeMemoryMap::FinalizeRequest::Segment with
llvm::orc::tpctypes::SegFinalizeRequest. This will simplify construction of a
new LLVM JITLinkMemoryManager that's capable of using SimpleNativeMemoryMap as
a backend.

Author: lhames

orc-rt/include/orc-rt/SimpleNativeMemoryMap.h

@@ -54,18 +54,10 @@ class SimpleNativeMemoryMap : public ResourceManager {
 
   struct FinalizeRequest {
     struct Segment {
-      enum class ContentType : uint8_t { Uninitialized, ZeroFill, Regular };
-
-      Segment() = default;
-      Segment(void *Address, size_t Size, AllocGroup G, ContentType C)
-          : Address(Address), Size(Size), G(G), C(C) {}
-
-      void *Address = nullptr;
+      AllocGroup AG;
+      char *Address = nullptr;
       size_t Size = 0;
-      AllocGroup G;
-      ContentType C = ContentType::Uninitialized;
-      char *data() { return reinterpret_cast<char *>(Address); }
-      size_t size() const { return Size; }
+      span<const char> Content;
     };
 
     std::vector<Segment> Segments;

orc-rt/lib/executor/SimpleNativeMemoryMap.cpp

@@ -31,36 +31,22 @@ struct SPSSimpleNativeMemoryMapSegment;
 template <>
 class SPSSerializationTraits<SPSSimpleNativeMemoryMapSegment,
                              SimpleNativeMemoryMap::FinalizeRequest::Segment> {
-  using SPSType = SPSTuple<SPSExecutorAddr, uint64_t, SPSAllocGroup, uint8_t>;
+  using SPSType =
+      SPSTuple<SPSAllocGroup, SPSExecutorAddr, uint64_t, SPSSequence<char>>;
 
 public:
   static bool deserialize(SPSInputBuffer &IB,
                           SimpleNativeMemoryMap::FinalizeRequest::Segment &S) {
-    using ContentType =
-        SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType;
-
+    AllocGroup AG;
     ExecutorAddr Address;
     uint64_t Size;
-    AllocGroup G;
-    uint8_t C;
-    if (!SPSType::AsArgList::deserialize(IB, Address, Size, G, C))
+    span<const char> Content;
+    if (!SPSType::AsArgList::deserialize(IB, AG, Address, Size, Content))
       return false;
-    if (Size >= std::numeric_limits<size_t>::max())
+    if (Size > std::numeric_limits<size_t>::max())
       return false;
-    S.Address = Address.toPtr<void *>();
-    S.Size = Size;
-    S.G = G;
-    S.C = static_cast<ContentType>(C);
-    switch (S.C) {
-    case ContentType::Uninitialized:
-      return true;
-    case ContentType::ZeroFill:
-      memset(reinterpret_cast<char *>(S.Address), 0, S.Size);
-      return true;
-    case ContentType::Regular:
-      // Read content directly into target address.
-      return IB.read(reinterpret_cast<char *>(S.Address), S.Size);
-    }
+    S = {AG, Address.toPtr<char *>(), static_cast<size_t>(Size), Content};
+    return true;
   }
 };
 
@@ -138,10 +124,31 @@ void SimpleNativeMemoryMap::finalize(OnFinalizeCompleteFn &&OnComplete,
   // TODO: Record finalize segments for release.
   // std::vector<std::pair<void*, size_t>> FinalizeSegments;
 
+  // Check segment validity before proceeding.
   for (auto &S : FR.Segments) {
-    if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.G.getMemProt()))
+
+    if (S.Content.size() > S.Size) {
+      return OnComplete(make_error<StringError>(
+          (std::ostringstream()
+           << "For segment [" << (void *)S.Address << ".."
+           << (void *)(S.Address + S.Size) << "), "
+           << " content size (" << std::hex << S.Content.size()
+           << ") exceeds segment size (" << S.Size << ")")
+              .str()));
+    }
+
+    // Copy any requested content.
+    if (!S.Content.empty())
+      memcpy(S.Address, S.Content.data(), S.Content.size());
+
+    // Zero-fill the rest of the section.
+    if (size_t ZeroFillSize = S.Size - S.Content.size())
+      memset(S.Address + S.Content.size(), 0, ZeroFillSize);
+
+    if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.AG.getMemProt()))
       return OnComplete(std::move(Err));
-    switch (S.G.getMemLifetime()) {
+
+    switch (S.AG.getMemLifetime()) {
     case MemLifetime::Standard:
       if (!Base || S.Address < Base)
         Base = S.Address;

orc-rt/unittests/SimpleNativeMemoryMapTest.cpp

@@ -31,49 +31,32 @@ struct SPSSimpleNativeMemoryMapSegment;
 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)
+  TestSNMMSegment(AllocGroup AG, char *Address, size_t Size,
+                  std::vector<char> C = {})
       : SimpleNativeMemoryMap::FinalizeRequest::Segment(
-            Address, Size, G,
-            Content == ZeroFill ? ContentType::ZeroFill
-                                : ContentType::Uninitialized) {}
+            {AG, Address, Size, {}}),
+        OwnedContent(std::move(C)) {
+    this->Content = {OwnedContent.data(), OwnedContent.size()};
+  }
 
-  std::string Content;
+  std::vector<char> OwnedContent;
 };
 
 template <>
 class SPSSerializationTraits<SPSSimpleNativeMemoryMapSegment, TestSNMMSegment> {
-  using SPSType = SPSTuple<SPSExecutorAddr, uint64_t, SPSAllocGroup, uint8_t>;
+  using SPSType =
+      SPSTuple<SPSAllocGroup, SPSExecutorAddr, uint64_t, SPSSequence<char>>;
 
 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);
+    return SPSType::AsArgList::size(S.AG, ExecutorAddr::fromPtr(S.Address),
+                                    static_cast<uint64_t>(S.Size), S.Content);
   }
 
   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;
+    return SPSType::AsArgList::serialize(
+        OB, S.AG, ExecutorAddr::fromPtr(S.Address),
+        static_cast<uint64_t>(S.Size), S.Content);
   }
 };
 
@@ -207,30 +190,61 @@ TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
 
   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});
+  char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
+      reinterpret_cast<char *>(Addr) + 64 * 1024;
+  uint64_t SentinelValue1 = 0; // Read from pre-filled content
+  uint64_t SentinelValue2 = 0; // Written in finalize, read back during dealloc.
+  uint64_t SentinelValue3 = 42; // Read from zero-filled region.
+
+  // Build initial content vector.
+  std::vector<char> Content;
+  Content.resize(sizeof(uint64_t) * 2);
+  memcpy(Content.data(), &SentinelValue3, sizeof(uint64_t));
+  memcpy(Content.data() + sizeof(uint64_t), &SentinelValue1, sizeof(uint64_t));
+
+  FR.Segments.push_back({MemProt::Read | MemProt::Write, FinalizeBase,
+                         64 * 1024, std::move(Content)});
+
+  // Read initial content into Sentinel 1.
+  FR.AAPs.push_back({
+      *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
+          read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue1),
+          ExecutorAddr::fromPtr(FinalizeBase)),
+      {} // No dealloc action.
+  });
+
+  // Write value in finalize action, then read back into Sentinel 2.
   FR.AAPs.push_back(
       {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
-           write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
+           write_value_sps_allocaction,
+           ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t),
            uint64_t(42)),
        *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
-           read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue),
-           ExecutorAddr::fromPtr(FinalizeBase))});
+           read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue2),
+           ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t))});
+
+  // Read first 64 bits of the zero-fill region.
+  FR.AAPs.push_back({
+      *MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
+          read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue3),
+          ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t) * 2),
+      {} // No dealloc action.
+  });
+
   snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR));
   void *FinalizeKeyAddr = cantFail(cantFail(FinalizeKey.get()));
 
-  EXPECT_EQ(SentinelValue, 0U);
+  EXPECT_EQ(SentinelValue1, 42U);
+  EXPECT_EQ(SentinelValue2, 0U);
+  EXPECT_EQ(SentinelValue3, 0U);
 
   std::future<Expected<Error>> DeallocResult;
   snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr);
   cantFail(cantFail(DeallocResult.get()));
 
-  EXPECT_EQ(SentinelValue, 42);
+  EXPECT_EQ(SentinelValue1, 42U);
+  EXPECT_EQ(SentinelValue2, 42U);
+  EXPECT_EQ(SentinelValue3, 0U);
 
   std::future<Expected<Error>> ReleaseResult;
   snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
@@ -248,13 +262,13 @@ TEST(SimpleNativeMemoryMap, ReserveFinalizeShutdown) {
 
   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);
+  char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
+      reinterpret_cast<char *>(Addr) + 64 * 1024;
   uint64_t SentinelValue = 0;
 
-  FR.Segments.push_back({FinalizeBase, 64 * 1024,
-                         MemProt::Read | MemProt::Write,
-                         TestSNMMSegment::ZeroFill});
+  FR.Segments.push_back(
+      {MemProt::Read | MemProt::Write, FinalizeBase, 64 * 1024});
+
   FR.AAPs.push_back(
       {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
            write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
@@ -285,13 +299,13 @@ TEST(SimpleNativeMemoryMap, ReserveFinalizeDetachShutdown) {
 
   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);
+  char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
+      reinterpret_cast<char *>(Addr) + 64 * 1024;
   uint64_t SentinelValue = 0;
 
-  FR.Segments.push_back({FinalizeBase, 64 * 1024,
-                         MemProt::Read | MemProt::Write,
-                         TestSNMMSegment::ZeroFill});
+  FR.Segments.push_back(
+      {MemProt::Read | MemProt::Write, FinalizeBase, 64 * 1024});
+
   FR.AAPs.push_back(
       {*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
            write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),